From 9f4279031f1d515f85dd656945326b3ad0827a99 Mon Sep 17 00:00:00 2001 From: Vivi Nousi Date: Wed, 25 Jan 2023 10:02:26 +0200 Subject: [PATCH 1/2] yolov5 dataset conversion for external training --- .../python/perception/object_detection_2d/yolov5/README.md | 5 ++++- .../perception/object_detection_2d/yolov5/yolov5_learner.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/projects/python/perception/object_detection_2d/yolov5/README.md b/projects/python/perception/object_detection_2d/yolov5/README.md index fd3aa7c4c0..861a2bcf25 100644 --- a/projects/python/perception/object_detection_2d/yolov5/README.md +++ b/projects/python/perception/object_detection_2d/yolov5/README.md @@ -4,4 +4,7 @@ This folder contains minimal code usage examples that showcase the basic inferen provided by OpenDR. Specifically the following examples are provided: 1. inference_demo.py: Perform inference on a single image. Setting `--device cpu` performs inference on CPU. 2. webcam_demo.py: A simple tool that performs live object detection using a webcam. -3. inference_tutorial.ipynb: Perform inference using pretrained or custom models. \ No newline at end of file +3. inference_tutorial.ipynb: Perform inference using pretrained or custom models. +4. convert_detection_dataset.py: An example of how to convert a `DetectionDataset` into the required format +to train a custom model. Training instructions can be found on the original + [YOLOv5 repository](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data#3-train) \ No newline at end of file diff --git a/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py b/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py index 2aede4f848..27ca7f8c7d 100644 --- a/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py +++ b/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py @@ -40,7 +40,7 @@ def __init__(self, model_name, path=None, device='cuda', temp_path='.', force_re force_reload=force_reload) else: self.model = torch.hub.load('ultralytics/yolov5:master', 'custom', path=path, - force_reload=force_reload, skip_validation=True) + force_reload=force_reload) torch.hub.set_dir(default_dir) self.model.to(device) From 70b35ed633a4379efcc68cfd5bd6833a6282bc7a Mon Sep 17 00:00:00 2001 From: Vivi Nousi Date: Wed, 25 Jan 2023 10:02:37 +0200 Subject: [PATCH 2/2] yolov5 dataset conversion for external training --- .../yolov5/convert_detection_dataset.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 projects/python/perception/object_detection_2d/yolov5/convert_detection_dataset.py diff --git a/projects/python/perception/object_detection_2d/yolov5/convert_detection_dataset.py b/projects/python/perception/object_detection_2d/yolov5/convert_detection_dataset.py new file mode 100644 index 0000000000..19e9d44b18 --- /dev/null +++ b/projects/python/perception/object_detection_2d/yolov5/convert_detection_dataset.py @@ -0,0 +1,112 @@ +# Copyright 2020-2023 OpenDR European Project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This example shows one way to convert a DetectionDataset, namely the AGIHumans dataset, to YOLOv5 format, +to be used for training a custom model using the original YOLOv5 implementation: https://github.com/ultralytics/yolov5 +The opendr_datasets package can be installed using: `pip install git+https://github.com/opendr-eu/datasets.git` +""" + +import argparse +import os +import yaml +import cv2 +from opendr_datasets import AGIHumans + + +def main(args): + path = args.original_data_path + train_set = AGIHumans(path, train=True) + val_set = AGIHumans(path, train=False) + + new_path = args.new_data_path + os.makedirs(new_path, exist_ok=True) + + # step 1: write dataset .yml file + # the new data structure is as follows: + # new_path + # ├── train + # │ ├── images + # │ │ ├── im00001.jpg + # │ │ └── ... + # │ └── labels + # │ ├── im00001.txt + # │ └── ... + # ├── test + # │ ├── images + # │ │ ├── im00001.jpg + # │ │ └── ... + # │ └── labels + # │ ├── im00001.txt + # │ └── ... + # └── AGIHumans.yml + + d = { + 'path': new_path, + 'train': 'train', + 'val': 'test', + 'names': {c: c_name for c, c_name in enumerate(train_set.class_names)} + } + + with open('AGIHumans.yaml', 'w') as yaml_file: + yaml.dump(d, yaml_file, default_flow_style=False) + + # step 2: convert annotations to .txt files + # train set + os.makedirs(os.path.join(new_path, 'train', 'images'), exist_ok=True) + os.makedirs(os.path.join(new_path, 'train', 'labels'), exist_ok=True) + for idx, (img, boxes) in enumerate(train_set): + # save img to 'train/images/im{:05d}.jpg' + im_cv = img.opencv() + cv2.imwrite(os.path.join(new_path, 'train', 'images', f'im{idx:05d}.jpg'), im_cv) + im_height, im_width, im_c = im_cv.shape + # save normalized label to 'train/labels/im{:05d}.txt + lines = '' + for box in boxes: + x_center = (box.left + box.width * 0.5) / im_width + y_center = (box.top + box.height * 0.5) / im_height + width = box.width / im_width + height = box.height / im_height + lines += f'{box.name} {x_center} {y_center} {width} {height}\n' + if len(lines) > 0: + with open(os.path.join(new_path, 'train', 'labels', f'im{idx:05d}.txt'), 'w') as f: + f.write(lines) + + # validation/test set + os.makedirs(os.path.join(new_path, 'test', 'images'), exist_ok=True) + os.makedirs(os.path.join(new_path, 'test', 'labels'), exist_ok=True) + for idx, (img, boxes) in enumerate(val_set): + # save img to 'train/images/im{:05d}.jpg' + im_cv = img.opencv() + cv2.imwrite(os.path.join(new_path, 'test', 'images', f'im{idx:05d}.jpg'), im_cv) + im_height, im_width, im_c = im_cv.shape + # save normalized label to 'train/labels/im{:05d}.txt + lines = '' + for box in boxes: + x_center = (box.left + box.width * 0.5) / im_width + y_center = (box.top + box.height * 0.5) / im_height + width = box.width / im_width + height = box.height / im_height + lines += f'{box.name} {x_center} {y_center} {width} {height}\n' + if len(lines) > 0: + with open(os.path.join(new_path, 'test', 'labels', f'im{idx:05d}.txt'), 'w') as f: + f.write(lines) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--original-data-path", help="Dataset root", type=str) + parser.add_argument("--new-data-path", help="Path to converted dataset location", type=str) + + args = parser.parse_args()