Update my_model/detector/object_detection.py

my_model/detector/object_detection.py

@@ -1,4 +1,4 @@
-
+from typing import Union, Optional, List, Tuple
 import streamlit as st
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
 import torch
@@ -11,18 +11,18 @@ from my_model.utilities.gen_utilities import get_image_path, get_model_path ,sho
 
 class ObjectDetector:
     """
-        A class for detecting objects in images using models like Detic and YOLOv5.
-
-        This class supports loading and using different object detection models to identify objects
-        in images and draw bounding boxes around them.
-
-        Attributes:
-            model (torch.nn.Module): The loaded object detection model.
-            processor (transformers.AutoImageProcessor): Processor for the Detic model.
-            model_name (str): Name of the model used for detection.
-        """
+    A class for detecting objects in images using models like Detic and YOLOv5.
+    This class supports loading and using different object detection models to identify objects
+    in images and draw bounding boxes around them.
+
+    Attributes:
+        model (torch.nn.Module or None): The loaded object detection model.
+        processor (transformers.AutoImageProcessor or None): Processor for the Detic model.
+        model_name (str or None): Name of the model used for detection.
+        device (str): Device to use for computation ('cuda' if available, otherwise 'cpu').
+    """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initializes the ObjectDetector class with default values.
         """
@@ -33,17 +33,17 @@ class ObjectDetector:
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         
 
-    def load_model(self, model_name='detic', pretrained=True, model_version='yolov5s'):
+    def load_model(self, model_name: str = 'detic', pretrained: bool = True, model_version: str = 'yolov5s') -> None:
         """
-                Load the specified object detection model.
+        Load the specified object detection model.
 
-                Args:
-                    model_name (str): Name of the model to load. Options are 'detic' and 'yolov5'.
-                    pretrained (bool): Boolean indicating if a pretrained model should be used.
-                    model_version (str): Version of the YOLOv5 model, applicable only when using YOLOv5.
+        Args:
+            model_name (str): Name of the model to load. Options are 'detic' and 'yolov5'.
+            pretrained (bool): Boolean indicating if a pretrained model should be used.
+            model_version (str): Version of the YOLOv5 model, applicable only when using YOLOv5.
 
-                Raises:
-                    ValueError: If an unsupported model name is provided.
+        Raises:
+            ValueError: If an unsupported model name is provided.
         """
 
         self.model_name = model_name
@@ -55,12 +55,15 @@ class ObjectDetector:
             raise ValueError(f"Unsupported model name: {model_name}")
 
     
-    def _load_detic_model(self, pretrained):
+    def _load_detic_model(self, pretrained: bool) -> None:
         """
         Load the Detic model.
 
         Args:
             pretrained (bool): If True, load a pretrained model.
+
+        Raises:
+            Exception: If an error occurs during model loading.
         """
 
         try:
@@ -72,13 +75,15 @@ class ObjectDetector:
             raise
 
     
-    def _load_yolov5_model(self, pretrained, model_version):
+    def _load_yolov5_model(self, pretrained: bool, model_version: str) -> None:
         """
         Load the YOLOv5 model.
 
         Args:
             pretrained (bool): If True, load a pretrained model.
             model_version (str): Version of the YOLOv5 model.
+        Raises:
+            Exception: If an error occurs during model loading.
         """
 
         try:
@@ -92,13 +97,16 @@ class ObjectDetector:
             raise
 
     
-    def process_image(self, image_input):
+    def process_image(self, image_input: Union[str, io.IOBase, Image.Image]) -> Image.Image:
         """
         Process the image from the given path or file-like object.
+
         Args:
-            image_input (str or file-like object): Path to the image file or a file-like object.
+            image_input (Union[str, io.IOBase, Image.Image]): Path to the image file, a file-like object, or a PIL Image.
+
         Returns:
             Image.Image: Processed image in RGB format.
+
         Raises:
             Exception: If an error occurs during image processing.
         """
@@ -119,16 +127,17 @@ class ObjectDetector:
             raise
 
     
-    def detect_objects(self, image, threshold=0.4):
+   def detect_objects(self, image: Image.Image, threshold: float = 0.4) -> Tuple[str, List[Tuple[str, List[float], float]]]:
         """
         Detect objects in the given image using the loaded model.
 
         Args:
             image (Image.Image): Image in which to detect objects.
-            threshold (float): Model detection confidence.
+            threshold (float): Model detection confidence threshold.
 
         Returns:
-            tuple: A tuple containing a string representation and a list of detected objects.
+            Tuple[str, List[Tuple[str, List[float], float]]]: A tuple containing a string representation and a list of detected objects.
+                                                               Each object in the list is represented as a tuple (label_name, box_rounded, certainty).
 
         Raises:
             ValueError: If the model is not loaded or the model name is unsupported.
@@ -142,7 +151,7 @@ class ObjectDetector:
             raise ValueError("Model not loaded or unsupported model name")
 
     
-    def _detect_with_detic(self, image, threshold):
+    def _detect_with_detic(self, image: Image.Image, threshold: float) -> Tuple[str, List[Tuple[str, List[float], float]]]:
         """
         Detect objects using the Detic model.
 
@@ -151,8 +160,8 @@ class ObjectDetector:
             threshold (float): The confidence threshold for detections.
 
         Returns:
-            tuple: A tuple containing a string representation and a list of detected objects.
-                   Each object in the list is represented as a tuple (label_name, box_rounded, certainty).
+            Tuple[str, List[Tuple[str, List[float], float]]]: A tuple containing a string representation and a list of detected objects.
+                                                               Each object in the list is represented as a tuple (label_name, box_rounded, certainty).
         """
 
         inputs = self.processor(images=image, return_tensors="pt")
@@ -171,7 +180,7 @@ class ObjectDetector:
         return detected_objects_str, detected_objects_list
 
     
-    def _detect_with_yolov5(self, image, threshold):
+    def _detect_with_yolov5(self, image: Image.Image, threshold: float) -> Tuple[str, List[Tuple[str, List[float], float]]]:
         """
         Detect objects using the YOLOv5 model.
 
@@ -180,8 +189,8 @@ class ObjectDetector:
             threshold (float): The confidence threshold for detections.
 
         Returns:
-            tuple: A tuple containing a string representation and a list of detected objects.
-                   Each object in the list is represented as a tuple (label_name, box_rounded, certainty).
+            Tuple[str, List[Tuple[str, List[float], float]]]: A tuple containing a string representation and a list of detected objects.
+                                                               Each object in the list is represented as a tuple (label_name, box_rounded, certainty).
         """
 
         cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
@@ -198,13 +207,13 @@ class ObjectDetector:
         return detected_objects_str, detected_objects_list
 
     
-    def draw_boxes(self, image, detected_objects, show_confidence=True):
+    def draw_boxes(self, image: Image.Image, detected_objects: List[Tuple[str, List[float], float]], show_confidence: bool = True) -> Image.Image:
         """
         Draw bounding boxes around detected objects in the image.
 
         Args:
             image (Image.Image): Image on which to draw.
-            detected_objects (list): List of detected objects.
+            detected_objects (List[Tuple[str, List[float], float]]): List of detected objects.
             show_confidence (bool): Whether to show confidence scores.
 
         Returns:
@@ -232,7 +241,7 @@ class ObjectDetector:
         return image
 
 
-def detect_and_draw_objects(image_path, model_type='yolov5', threshold=0.2, show_confidence=True):
+def detect_and_draw_objects(image_path: str, model_type: str = 'yolov5', threshold: float = 0.2, show_confidence: bool = True) -> Tuple[Image.Image, str]:
     """
     Detects objects in an image, draws bounding boxes around them, and returns the processed image and a string description.
 
@@ -243,7 +252,7 @@ def detect_and_draw_objects(image_path, model_type='yolov5', threshold=0.2, show
         show_confidence (bool): Whether to show confidence scores on the output image.
 
     Returns:
-        tuple: A tuple containing the processed Image.Image and a string of detected objects.
+        Tuple[Image.Image, str]: A tuple containing the processed Image.Image and a string of detected objects.
     """
 
     detector = ObjectDetector()
@@ -252,9 +261,3 @@ def detect_and_draw_objects(image_path, model_type='yolov5', threshold=0.2, show
     detected_objects_string, detected_objects_list = detector.detect_objects(image, threshold=threshold)
     image_with_boxes = detector.draw_boxes(image, detected_objects_list, show_confidence=show_confidence)
     return image_with_boxes, detected_objects_string
-
-
-
-if __name__ == "__main__":
-    pass
-