diff --git a/docs/reference/object-detection-2d-yolov5.md b/docs/reference/object-detection-2d-yolov5.md index 58420328eb..60af097428 100644 --- a/docs/reference/object-detection-2d-yolov5.md +++ b/docs/reference/object-detection-2d-yolov5.md @@ -58,6 +58,27 @@ Parameters: - **size**: *int, default=640*\ Size of image for inference. The image is resized to this in both sides before being fed to the model. + +#### `YOLOv5DetectorLearner.download` +```python +YOLOv5DetectorLearner.download(self, path, mode, verbose, url, model_name, img_name) +``` + +Downloads the pretrained weights of a YOLOv5s model fine-tuned for truck detection, along with sample truck images for inference, stored in .pt and image files respectively. + +Parameters: + +- **path**: *str, default=None*\ + Specifies the folder where data will be downloaded. If *None*, the *self.temp_path* directory is used instead. +- **mode**: *{'pretrained', 'images', 'test_data'}, default='pretrained'*\ + If *'pretrained'*, downloads a pretrained detector model. If *'images'*, downloads an image to perform inference on. If + *'test_data'* downloads a dummy dataset for testing purposes. +- **verbose**: *bool default=True*\ + If True, enables maximum verbosity. +- **url**: *str, default=OpenDR FTP URL*\ + URL of the FTP server. +- **model_name**: name of model ftp server, *default = 'yolov5_finetuned_in_trucks.pt'.*\ +- **image_name**: name of image in ftp server, *default = 'truck1.png'.*\ #### Examples diff --git a/projects/python/perception/object_detection_2d/yolov5/inference_demo_agriculture_vehicles.py b/projects/python/perception/object_detection_2d/yolov5/inference_demo_agriculture_vehicles.py new file mode 100644 index 0000000000..cdf37fef7e --- /dev/null +++ b/projects/python/perception/object_detection_2d/yolov5/inference_demo_agriculture_vehicles.py @@ -0,0 +1,41 @@ +# 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. + +import argparse +from opendr.engine.data import Image +from opendr.perception.object_detection_2d import YOLOv5DetectorLearner +from opendr.perception.object_detection_2d import draw_bounding_boxes + + +if __name__ == '__main__': + # Parse command line arguments + parser = argparse.ArgumentParser() + parser.add_argument("--model_name", help="Model name or path", type=str, default='yolov5s_trucks') + parser.add_argument("--device", help="Device to use (cpu, cuda)", type=str, default="cuda", choices=["cuda", "cpu"]) + parser.add_argument("--model_dir", help="Model directory", type=str, default="./yolov5s_finetuned_in_trucks.pt") + args = parser.parse_args() + + # Initialize the YOLOv5 detector with the given model and device + yolo = YOLOv5DetectorLearner(model_name=args.model_name, device=args.device, path=args.model_dir) + yolo.download(".", mode="images", verbose=True, img_name="truck4.jpg") + yolo.download(".", mode="images", verbose=True, img_name="truck7.jpg") + + im1 = Image.open('truck4.jpg') + im2 = Image.open('truck7.jpg') + + results = yolo.infer(im1) + draw_bounding_boxes(im1.opencv(), results, yolo.classes, show=True, line_thickness=3) + + results = yolo.infer(im2) + draw_bounding_boxes(im2, results, yolo.classes, show=True, line_thickness=3) 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 27ca7f8c7d..249d188c0a 100644 --- a/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py +++ b/src/opendr/perception/object_detection_2d/yolov5/yolov5_learner.py @@ -11,11 +11,16 @@ # 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. +# General imports +import os +from urllib.request import urlretrieve # OpenDR engine imports from opendr.engine.learners import Learner from opendr.engine.data import Image from opendr.engine.target import BoundingBox, BoundingBoxList +from opendr.engine.constants import OPENDR_SERVER_URL + # yolov5 imports import torch @@ -28,19 +33,31 @@ class YOLOv5DetectorLearner(Learner): def __init__(self, model_name, path=None, device='cuda', temp_path='.', force_reload=False): super(YOLOv5DetectorLearner, self).__init__(device=device, temp_path=temp_path) - if model_name not in self.available_models: - model_name = 'yolov5s' - print('Unrecognized model name, defaulting to "yolov5s"') + self.device = device + self.model_directory = temp_path if path is None else path + self.model_name = model_name default_dir = torch.hub.get_dir() torch.hub.set_dir(temp_path) - if path is None: - self.model = torch.hub.load('ultralytics/yolov5:master', 'custom', f'{temp_path}/{model_name}', - force_reload=force_reload) - else: + # Downloading and loading the fine-tuned yolov5s model in trucks + if model_name == 'yolov5s_trucks': + self.download(path='./', mode="pretrained", verbose=True) self.model = torch.hub.load('ultralytics/yolov5:master', 'custom', path=path, force_reload=force_reload) + # Getting a generic model + else: + if model_name not in self.available_models: + model_name = 'yolov5s' + print('Unrecognized model name, defaulting to "yolov5s"') + + if path is None: + self.model = torch.hub.load('ultralytics/yolov5:master', 'custom', + f'{temp_path}/{model_name}', + force_reload=force_reload) + else: + self.model = torch.hub.load('ultralytics/yolov5:master', 'custom', path=path, + force_reload=force_reload) torch.hub.set_dir(default_dir) self.model.to(device) @@ -86,3 +103,57 @@ def load(self): def save(self): """This method is not used in this implementation.""" return NotImplementedError + + def download(self, path=None, mode="pretrained", verbose=False, + url=OPENDR_SERVER_URL + "/perception/object_detection_2d/yolov5/", + model_name='yolov5s_finetuned_in_trucks.pt', img_name='truck1.jpg'): + """ + Downloads all files necessary for inference, evaluation and training. Valid mode options are: ["pretrained", + "images", "test_data"]. + :param path: folder to which files will be downloaded, if None self.temp_path will be used + :type path: str, optional + :param mode: one of: ["pretrained", "images", "test_data"], where "pretrained" downloads a pretrained + network depending on the self.backbone type, "images" downloads example inference data, "backbone" downloads a + pretrained resnet backbone for training, and "annotations" downloads additional annotation files for training + :type mode: str, optional + :param verbose: if True, additional information is printed on stdout + :type verbose: bool, optional + :param model_name: the name of the model file to download (e.g., 'yolov5s.pt') + :type model_name: str, optional + :param url: URL to file location on FTP server + :type url: str, optional + """ + valid_modes = ["pretrained", "images", "test_data"] + if mode not in valid_modes: + raise ValueError("Invalid mode. Currently, only 'pretrained' mode is supported.") + + if path is None: + path = self.temp_path + + if not os.path.exists(path): + os.makedirs(path) + + if mode == "pretrained": + model_path = os.path.join(path, model_name) + if not os.path.exists(model_path): + if verbose: + print("Downloading pretrained model...") + file_url = os.path.join(url, "pretrained", model_name) + urlretrieve(file_url, model_path) + if verbose: + print(f"Downloaded model to {model_path}.") + else: + if verbose: + print("Model already exists.") + elif mode == "images": + image_path = os.path.join(path, img_name) + if not os.path.exists(image_path): + if verbose: + print("Downloading example image...") + file_url = os.path.join(url, "images", img_name) + urlretrieve(file_url, image_path) + if verbose: + print(f"Downloaded example image to {image_path}.") + else: + if verbose: + print("Example image already exists.")