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.7 KiB
73 lines
2.7 KiB
# Ultralytics YOLO 🚀, GPL-3.0 license |
|
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI |
|
# Example usage: yolo train data=Argoverse.yaml |
|
# parent |
|
# ├── yolov5 |
|
# └── datasets |
|
# └── Argoverse ← downloads here (31.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/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.yolo.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 |
|
dir = Path(yaml['path']) # dataset root dir |
|
urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip'] |
|
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 VisDrone annotations to YOLO labels
|
|
|