# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license

# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI

# Example usage: python train.py --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 utils.general import download, 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, delete=False)

  # 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