Spaces:
Sleeping
Sleeping
# Ultralytics YOLO π, AGPL-3.0 license | |
# Argoverse-HD dataset (ring-front-center camera) https://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI | |
# Documentation: https://docs.ultralytics.com/datasets/detect/argoverse/ | |
# Example usage: yolo train data=Argoverse.yaml | |
# parent | |
# βββ ultralytics | |
# βββ datasets | |
# βββ Argoverse β downloads here (31.5 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/Argoverse # dataset root dir | |
train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images | |
val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images | |
test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview | |
# Classes | |
names: | |
0: person | |
1: bicycle | |
2: car | |
3: motorcycle | |
4: bus | |
5: truck | |
6: traffic_light | |
7: stop_sign | |
# Download script/URL (optional) --------------------------------------------------------------------------------------- | |
download: | | |
import json | |
from tqdm import tqdm | |
from ultralytics.utils.downloads import download | |
from pathlib import Path | |
def argoverse2yolo(set): | |
labels = {} | |
a = json.load(open(set, "rb")) | |
for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."): | |
img_id = annot['image_id'] | |
img_name = a['images'][img_id]['name'] | |
img_label_name = f'{img_name[:-3]}txt' | |
cls = annot['category_id'] # instance class id | |
x_center, y_center, width, height = annot['bbox'] | |
x_center = (x_center + width / 2) / 1920.0 # offset and scale | |
y_center = (y_center + height / 2) / 1200.0 # offset and scale | |
width /= 1920.0 # scale | |
height /= 1200.0 # scale | |
img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']] | |
if not img_dir.exists(): | |
img_dir.mkdir(parents=True, exist_ok=True) | |
k = str(img_dir / img_label_name) | |
if k not in labels: | |
labels[k] = [] | |
labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n") | |
for k in labels: | |
with open(k, "w") as f: | |
f.writelines(labels[k]) | |
# Download 'https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip' (deprecated S3 link) | |
dir = Path(yaml['path']) # dataset root dir | |
urls = ['https://drive.google.com/file/d/1st9qW3BeIwQsnR0t8mRpvbsSWIo16ACi/view?usp=drive_link'] | |
print("\n\nWARNING: Argoverse dataset MUST be downloaded manually, autodownload will NOT work.") | |
print(f"WARNING: Manually download Argoverse dataset '{urls[0]}' to '{dir}' and re-run your command.\n\n") | |
# download(urls, dir=dir) | |
# Convert | |
annotations_dir = 'Argoverse-HD/annotations/' | |
(dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images') # rename 'tracking' to 'images' | |
for d in "train.json", "val.json": | |
argoverse2yolo(dir / annotations_dir / d) # convert Argoverse annotations to YOLO labels | |