Finetuned yolov5 for trucks

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
 
@@ -68,7 +89,7 @@ Parameters:
   from opendr.perception.object_detection_2d import YOLOv5DetectorLearner
   from opendr.perception.object_detection_2d import draw_bounding_boxes
 
-  yolo = YOLOv5DetectorLearner(model_name='yolov5s', device='cpu')
+  yolo = YOLOv5DetectorLearner(model_name='yolov5s-fine', device='cpu')
 
   torch.hub.download_url_to_file('https://ultralytics.com/images/zidane.jpg', 'zidane.jpg')  # download image
   im1 = Image.open('zidane.jpg')  # OpenDR image

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)

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.")