Browse Source

CVPR 2021 Argoverse-HD dataset autodownload support (#2400)

* 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>
5.0
Kartikeya Sharma GitHub 3 years ago
parent
commit
c51dfec8ea
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 86 additions and 0 deletions
  1. +21
    -0
      data/argoverse_hd.yaml
  2. +65
    -0
      data/scripts/get_argoverse_hd.sh

+ 21
- 0
data/argoverse_hd.yaml View File

@@ -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
- 0
data/scripts/get_argoverse_hd.sh View File

@@ -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/





Loading…
Cancel
Save