88 lines
3.5 KiB
YAML
88 lines
3.5 KiB
YAML
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University
|
|
# Documentation: https://docs.ultralytics.com/datasets/detect/visdrone/
|
|
# 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: VisDrone # dataset root dir
|
|
train: images/train # train images (relative to 'path') 6471 images
|
|
val: images/val # val images (relative to 'path') 548 images
|
|
test: images/test # test-dev 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
|
|
import shutil
|
|
|
|
from ultralytics.utils.downloads import download
|
|
from ultralytics.utils import TQDM
|
|
|
|
|
|
def visdrone2yolo(dir, split, source_name=None):
|
|
"""Convert VisDrone annotations to YOLO format with images/{split} and labels/{split} structure."""
|
|
from PIL import Image
|
|
|
|
source_dir = dir / (source_name or f"VisDrone2019-DET-{split}")
|
|
images_dir = dir / "images" / split
|
|
labels_dir = dir / "labels" / split
|
|
labels_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Move images to new structure
|
|
if (source_images_dir := source_dir / "images").exists():
|
|
images_dir.mkdir(parents=True, exist_ok=True)
|
|
for img in source_images_dir.glob("*.jpg"):
|
|
img.rename(images_dir / img.name)
|
|
|
|
for f in TQDM((source_dir / "annotations").glob("*.txt"), desc=f"Converting {split}"):
|
|
img_size = Image.open(images_dir / f.with_suffix(".jpg").name).size
|
|
dw, dh = 1.0 / img_size[0], 1.0 / img_size[1]
|
|
lines = []
|
|
|
|
with open(f, encoding="utf-8") as file:
|
|
for row in [x.split(",") for x in file.read().strip().splitlines()]:
|
|
if row[4] != "0": # Skip ignored regions
|
|
x, y, w, h = map(int, row[:4])
|
|
cls = int(row[5]) - 1
|
|
# Convert to YOLO format
|
|
x_center, y_center = (x + w / 2) * dw, (y + h / 2) * dh
|
|
w_norm, h_norm = w * dw, h * dh
|
|
lines.append(f"{cls} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}\n")
|
|
|
|
(labels_dir / f.name).write_text("".join(lines), encoding="utf-8")
|
|
|
|
|
|
# Download (ignores test-challenge split)
|
|
dir = Path(yaml["path"]) # dataset root dir
|
|
urls = [
|
|
"https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-train.zip",
|
|
"https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-val.zip",
|
|
"https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-dev.zip",
|
|
# "https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-challenge.zip",
|
|
]
|
|
download(urls, dir=dir, threads=4)
|
|
|
|
# Convert
|
|
splits = {"VisDrone2019-DET-train": "train", "VisDrone2019-DET-val": "val", "VisDrone2019-DET-test-dev": "test"}
|
|
for folder, split in splits.items():
|
|
visdrone2yolo(dir, split, folder) # convert VisDrone annotations to YOLO labels
|
|
shutil.rmtree(dir / folder) # cleanup original directory
|