Skip to content

Commit

Permalink
CVPR 2021 Argoverse-HD dataset autodownload support (ultralytics#2400)
Browse files Browse the repository at this point in the history
* added argoverse-download ability

* bugfix

* add support for Argoverse dataset

* Refactored code

* renamed to argoverse-HD

* unzip -q and YOLOv5

small cleanup items

* add image counts

Co-authored-by: Kartikeya Sharma <kartikes@trinity.vision.cs.cmu.edu>
Co-authored-by: Kartikeya Sharma <kartikes@trinity-0-32.eth>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
4 people authored Mar 10, 2021
1 parent 6d04607 commit 280c243
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 0 deletions.
21 changes: 21 additions & 0 deletions data/argoverse_hd.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
# Train command: python train.py --data argoverse_hd.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /argoverse
# /yolov5


# download command/URL (optional)
download: bash data/scripts/get_argoverse_hd.sh

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../argoverse/Argoverse-1.1/images/train/ # 39384 images
val: ../argoverse/Argoverse-1.1/images/val/ # 15062 iamges
test: ../argoverse/Argoverse-1.1/images/test/ # Submit to: https://eval.ai/web/challenges/challenge-page/800/overview

# number of classes
nc: 8

# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign' ]
65 changes: 65 additions & 0 deletions data/scripts/get_argoverse_hd.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
#!/bin/bash
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
# Download command: bash data/scripts/get_argoverse_hd.sh
# Train command: python train.py --data argoverse_hd.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /argoverse
# /yolov5

# Download/unzip images
d='../argoverse/' # unzip directory
mkdir $d
url=https://argoverse-hd.s3.us-east-2.amazonaws.com/
f=Argoverse-HD-Full.zip
wget $url$f -O $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
wait # finish background tasks

cd ../argoverse/Argoverse-1.1/
ln -s tracking images

cd ../Argoverse-HD/annotations/

python3 - "$@" <<END
import json
from pathlib import Path
annotation_files = ["train.json", "val.json"]
print("Converting annotations to YOLOv5 format...")
for val in annotation_files:
a = json.load(open(val, "rb"))
label_dict = {}
for annot in a['annotations']:
img_id = annot['image_id']
img_name = a['images'][img_id]['name']
img_label_name = img_name[:-3] + "txt"
obj_class = annot['category_id']
x_center, y_center, width, height = annot['bbox']
x_center = x_center / 1920.0
width = width / 1920.0
y_center = y_center / 1200.0
height = height / 1200.0
img_dir = "./labels/" + a['seq_dirs'][a['images'][annot['image_id']]['sid']]
Path(img_dir).mkdir(parents=True, exist_ok=True)
if img_dir + "/" + img_label_name not in label_dict:
label_dict[img_dir + "/" + img_label_name] = []
label_dict[img_dir + "/" + img_label_name].append(f"{obj_class} {x_center} {y_center} {width} {height}\n")
for filename in label_dict:
with open(filename, "w") as file:
for string in label_dict[filename]:
file.write(string)
END

mv ./labels ../../Argoverse-1.1/




0 comments on commit 280c243

Please sign in to comment.