voc_ball.yaml 3.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
  3. # Example usage: python train.py --data VOC.yaml
  4. # parent
  5. # ├── yolov5
  6. # └── datasets
  7. # └── VOC ← 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: ./
  10. train: # train images (relative to 'path') 16551 images
  11. - VOCdevkit/images/train/
  12. val: # val images (relative to 'path') 4952 images
  13. - VOCdevkit/images/val/
  14. test: # test images (optional)
  15. # Classes
  16. nc: 1 # number of classes
  17. names: ['ball'] # class names
  18. # Download script/URL (optional) ---------------------------------------------------------------------------------------
  19. download: |
  20. import xml.etree.ElementTree as ET
  21. from tqdm import tqdm
  22. from utils.general import download, Path
  23. def convert_label(path, lb_path, year, image_id):
  24. def convert_box(size, box):
  25. dw, dh = 1. / size[0], 1. / size[1]
  26. x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
  27. return x * dw, y * dh, w * dw, h * dh
  28. in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
  29. out_file = open(lb_path, 'w')
  30. tree = ET.parse(in_file)
  31. root = tree.getroot()
  32. size = root.find('size')
  33. w = int(size.find('width').text)
  34. h = int(size.find('height').text)
  35. for obj in root.iter('object'):
  36. cls = obj.find('name').text
  37. if cls in yaml['names'] and not int(obj.find('difficult').text) == 1:
  38. xmlbox = obj.find('bndbox')
  39. bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
  40. cls_id = yaml['names'].index(cls) # class id
  41. out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')
  42. # Download
  43. dir = Path(yaml['path']) # dataset root dir
  44. url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
  45. urls = [url + 'VOCtrainval_06-Nov-2007.zip', # 446MB, 5012 images
  46. url + 'VOCtest_06-Nov-2007.zip', # 438MB, 4953 images
  47. url + 'VOCtrainval_11-May-2012.zip'] # 1.95GB, 17126 images
  48. download(urls, dir=dir / 'images', delete=False)
  49. # Convert
  50. path = dir / f'images/VOCdevkit'
  51. for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):
  52. imgs_path = dir / 'images' / f'{image_set}{year}'
  53. lbs_path = dir / 'labels' / f'{image_set}{year}'
  54. imgs_path.mkdir(exist_ok=True, parents=True)
  55. lbs_path.mkdir(exist_ok=True, parents=True)
  56. image_ids = open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt').read().strip().split()
  57. for id in tqdm(image_ids, desc=f'{image_set}{year}'):
  58. f = path / f'VOC{year}/JPEGImages/{id}.jpg' # old img path
  59. lb_path = (lbs_path / f.name).with_suffix('.txt') # new label path
  60. f.rename(imgs_path / f.name) # move image
  61. convert_label(path, lb_path, year, id) # convert labels to YOLO format