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.

68 satır
2.7KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI
  3. # Example usage: python train.py --data Argoverse.yaml
  4. # parent
  5. # ├── yolov5
  6. # └── datasets
  7. # └── Argoverse ← downloads here
  8. # 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, ..]
  9. path: ../datasets/Argoverse # dataset root dir
  10. train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images
  11. val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images
  12. test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview
  13. # Classes
  14. nc: 8 # number of classes
  15. names: ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign'] # class names
  16. # Download script/URL (optional) ---------------------------------------------------------------------------------------
  17. download: |
  18. import json
  19. from tqdm import tqdm
  20. from utils.general import download, Path
  21. def argoverse2yolo(set):
  22. labels = {}
  23. a = json.load(open(set, "rb"))
  24. for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."):
  25. img_id = annot['image_id']
  26. img_name = a['images'][img_id]['name']
  27. img_label_name = img_name[:-3] + "txt"
  28. cls = annot['category_id'] # instance class id
  29. x_center, y_center, width, height = annot['bbox']
  30. x_center = (x_center + width / 2) / 1920.0 # offset and scale
  31. y_center = (y_center + height / 2) / 1200.0 # offset and scale
  32. width /= 1920.0 # scale
  33. height /= 1200.0 # scale
  34. img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']]
  35. if not img_dir.exists():
  36. img_dir.mkdir(parents=True, exist_ok=True)
  37. k = str(img_dir / img_label_name)
  38. if k not in labels:
  39. labels[k] = []
  40. labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n")
  41. for k in labels:
  42. with open(k, "w") as f:
  43. f.writelines(labels[k])
  44. # Download
  45. dir = Path('../datasets/Argoverse') # dataset root dir
  46. urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip']
  47. download(urls, dir=dir, delete=False)
  48. # Convert
  49. annotations_dir = 'Argoverse-HD/annotations/'
  50. (dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images') # rename 'tracking' to 'images'
  51. for d in "train.json", "val.json":
  52. argoverse2yolo(dir / annotations_dir / d) # convert VisDrone annotations to YOLO labels