40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import os
|
|
import pandas as pd
|
|
from PIL import Image
|
|
|
|
YOLO_LABELS_PATH = "../datasets/VisDrone/VisDrone2019-DET-val/labels"
|
|
VISANN_PATH = "../datasets/VisDrone/VisDrone2019-DET-val/annotations/"
|
|
VISIMG_PATH = "../datasets//VisDrone/VisDrone2019-DET-val/images/"
|
|
|
|
def convert(bbox, img_size):
|
|
#将标注visDrone数据集标注转为yolov5
|
|
#bbox top_left_x top_left_y width height
|
|
dw = 1/(img_size[0])
|
|
dh = 1/(img_size[1])
|
|
x = bbox[0] + bbox[2]/2
|
|
y = bbox[1] + bbox[3]/2
|
|
x = x * dw
|
|
y = y * dh
|
|
w = bbox[2] * dw
|
|
h = bbox[3] * dh
|
|
return (x,y,w,h)
|
|
|
|
def ChangeToYolo5():
|
|
if not os.path.exists(YOLO_LABELS_PATH):
|
|
os.makedirs(YOLO_LABELS_PATH)
|
|
print(len(os.listdir(VISANN_PATH)))
|
|
for file in os.listdir(VISANN_PATH):
|
|
image_path = VISIMG_PATH + '/' + file.replace('txt', 'jpg')
|
|
ann_file = VISANN_PATH + '/' + file
|
|
out_file = open(YOLO_LABELS_PATH + '/' + file, 'w')
|
|
bbox = pd.read_csv(ann_file,header=None).values
|
|
img = Image.open(image_path)
|
|
img_size = img.size
|
|
for row in bbox:
|
|
if(row[4]==1 and 0<row[5]<11):
|
|
label = convert(row[:4], img_size)
|
|
out_file.write(str(row[5]-1) + " " + " ".join(str(f'{x:.6f}') for x in label) + '\n')
|
|
out_file.close()
|
|
|
|
if __name__ == '__main__':
|
|
ChangeToYolo5() |