You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
73 lines
2.9 KiB
73 lines
2.9 KiB
# Ultralytics YOLO 🚀, GPL-3.0 license |
|
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University |
|
# Example usage: yolo train data=VisDrone.yaml |
|
# parent |
|
# ├── ultralytics |
|
# └── datasets |
|
# └── VisDrone ← downloads here (2.3 GB) |
|
|
|
|
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] |
|
path: ../datasets/VisDrone # dataset root dir |
|
train: VisDrone2019-DET-train/images # train images (relative to 'path') 6471 images |
|
val: VisDrone2019-DET-val/images # val images (relative to 'path') 548 images |
|
test: VisDrone2019-DET-test-dev/images # test images (optional) 1610 images |
|
|
|
# Classes |
|
names: |
|
0: pedestrian |
|
1: people |
|
2: bicycle |
|
3: car |
|
4: van |
|
5: truck |
|
6: tricycle |
|
7: awning-tricycle |
|
8: bus |
|
9: motor |
|
|
|
|
|
# Download script/URL (optional) --------------------------------------------------------------------------------------- |
|
download: | |
|
import os |
|
from pathlib import Path |
|
|
|
from ultralytics.yolo.utils.downloads import download |
|
|
|
def visdrone2yolo(dir): |
|
from PIL import Image |
|
from tqdm import tqdm |
|
|
|
def convert_box(size, box): |
|
# Convert VisDrone box to YOLO xywh box |
|
dw = 1. / size[0] |
|
dh = 1. / size[1] |
|
return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh |
|
|
|
(dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory |
|
pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}') |
|
for f in pbar: |
|
img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size |
|
lines = [] |
|
with open(f, 'r') as file: # read annotation.txt |
|
for row in [x.split(',') for x in file.read().strip().splitlines()]: |
|
if row[4] == '0': # VisDrone 'ignored regions' class 0 |
|
continue |
|
cls = int(row[5]) - 1 |
|
box = convert_box(img_size, tuple(map(int, row[:4]))) |
|
lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n") |
|
with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl: |
|
fl.writelines(lines) # write label.txt |
|
|
|
|
|
# Download |
|
dir = Path(yaml['path']) # dataset root dir |
|
urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip', |
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip', |
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip', |
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip'] |
|
download(urls, dir=dir, curl=True, threads=4) |
|
|
|
# Convert |
|
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev': |
|
visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels
|
|
|