Update my_model/utilities/gen_utilities.py

my_model/utilities/gen_utilities.py

@@ -10,15 +10,17 @@ import gc
 import streamlit as st
 
 
-    
-def show_image(image):
+def show_image(image: Union[str, Image.Image, np.ndarray, torch.Tensor]) -> None:
     """
     Display an image in various environments (Jupyter, PyCharm, Hugging Face Spaces).
-    Handles different types of image inputs (file path, PIL Image, numpy array, OpenCV, PyTorch tensor).
+    Handles different types of image inputs (file path, PIL Image, numpy array, PyTorch tensor).
 
     Args:
-        image (str or PIL.Image or numpy.ndarray or torch.Tensor): The image to display.
-    """
+        image (Union[str, Image.Image, np.ndarray, torch.Tensor]): The image to display.
+
+    Returns:
+        None
+    """    
     
     in_jupyter = is_jupyter_notebook()
     in_colab = is_google_colab()
@@ -53,7 +55,17 @@ def show_image(image):
 
 
 
-def show_image_with_matplotlib(image):
+def show_image_with_matplotlib(image: Union[str, Image.Image, np.ndarray, torch.Tensor]) -> None:
+    """
+    Display an image using Matplotlib.
+
+    Args:
+        image (Union[str, Image.Image, np.ndarray, torch.Tensor]): The image to display.
+
+    Returns:
+        None
+    """
+    
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, np.ndarray):
@@ -66,7 +78,7 @@ def show_image_with_matplotlib(image):
     plt.show()
 
 
-def is_jupyter_notebook():
+def is_jupyter_notebook() -> bool:
     """
     Check if the code is running in a Jupyter notebook.
 
@@ -85,25 +97,40 @@ def is_jupyter_notebook():
     return False  # Default to False if none of the above conditions are met
 
 
-def is_pycharm():
+def is_pycharm() -> bool:
+    """
+    Check if the code is running in PyCharm.
+
+    Returns:
+        bool: True if running in PyCharm, False otherwise.
+    """
+    
     return 'PYCHARM_HOSTED' in os.environ
 
 
-def is_google_colab():
+def is_google_colab() -> bool:
+    """
+    Check if the code is running in Google Colab.
+
+    Returns:
+        bool: True if running in Google Colab, False otherwise.
+    """
+    
     return 'COLAB_GPU' in os.environ or 'google.colab' in sys.modules
 
 
-def get_image_path(name, path_type):
+def get_image_path(name: str, path_type: str) -> str:
     """
     Generates a path for models, images, or data based on the specified type.
 
     Args:
-    name (str): The name of the model, image, or data folder/file.
-    path_type (str): The type of path needed ('models', 'images', or 'data').
+        name (str): The name of the model, image, or data folder/file.
+        path_type (str): The type of path needed ('models', 'images', or 'data').
 
     Returns:
-    str: The full path to the specified resource.
+        str: The full path to the specified resource.
     """
+    
     # Get the current working directory (assumed to be inside 'code' folder)
     current_dir = os.getcwd()
 
@@ -119,7 +146,7 @@ def get_image_path(name, path_type):
     return full_path
 
 
-def get_model_path(model_name):
+def get_model_path(model_name: str) -> str:
     """
     Get the path to the specified model folder.
     
@@ -129,6 +156,7 @@ def get_model_path(model_name):
     Returns:
         str: Absolute path to the specified model folder.
     """
+    
     # Directory of the current script
     current_script_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -140,11 +168,62 @@ def get_model_path(model_name):
 
     return model_path
 
+def add_detected_objects_to_dataframe(df: pd.DataFrame, detector_type: str, image_directory: str, detector: object) -> pd.DataFrame:
+    """
+    Adds a column to the DataFrame with detected objects for each image specified in the 'image_name' column.
+    Prints a message every 200 images processed.
 
+    Args:
+        df (pd.DataFrame): DataFrame containing a column 'image_name' with image filenames.
+        detector_type (str): The detection model to use ('detic' or 'yolov5').
+        image_directory (str): Path to the directory containing images.
+        detector (object): An instance of the ObjectDetector class.
 
-def free_gpu_resources():
+    Returns:
+        pd.DataFrame: The original DataFrame with an additional column 'detected_objects'.
+    """
+    
+    # Ensure 'image_name' column exists in the DataFrame
+    if 'image_name' not in df.columns:
+        raise ValueError("DataFrame must contain an 'image_name' column.")
+    
+    detector.load_model(detector_type)
+    
+    # Initialize a counter for images processed
+    images_processed = 0
+    
+    # Function to detect objects for a given image filename
+    def detect_objects_for_image(image_name):
+        nonlocal images_processed  # Use the nonlocal keyword to modify the images_processed variable
+        image_path = os.path.join(image_directory, image_name)
+        if os.path.exists(image_path):
+            
+            image = detector.process_image(image_path)
+            detected_objects_str, _ = detector.detect_objects(image, 0.2)
+            images_processed += 1
+            
+            # Print message every 2 images processed
+            if images_processed % 200 == 0:
+                print(f"Completed {images_processed} images detection")
+                
+            return detected_objects_str
+        else:
+            images_processed += 1
+            return "Image not found"
+
+    # Apply the function to each row in the DataFrame
+    df[detector.model_name] = df['image_name'].apply(detect_objects_for_image)
+
+    return df
+
+
+
+def free_gpu_resources() -> None:
     """
     Clears GPU memory.
+
+    Returns:
+        None
     """
 
     if torch.cuda.is_available():