Spaces:
Running
on
T4
Running
on
T4
| # Ultralytics YOLO 🚀, AGPL-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(f'{os.sep}annotations{os.sep}', f'{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 | |