my_model/utilities/gen_utilities.py

import pandas as pd
import os
from PIL import Image
import numpy as np
import torch
import matplotlib.pyplot as plt
from IPython import get_ipython
import sys
import gc
import streamlit as st
from typing import Tuple, Dict, List, Union


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, PyTorch tensor).

    Args:
        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()

    # Convert image to PIL Image if it's a file path, numpy array, or PyTorch tensor
    if isinstance(image, str):

        if os.path.isfile(image):
            image = Image.open(image)
        else:
            raise ValueError("File path provided does not exist.")
    elif isinstance(image, np.ndarray):

        if image.ndim == 3 and image.shape[2] in [3, 4]:

            image = Image.fromarray(image[..., ::-1] if image.shape[2] == 3 else image)
        else:

            image = Image.fromarray(image)
    elif torch.is_tensor(image):

        image = Image.fromarray(image.permute(1, 2, 0).numpy().astype(np.uint8))

    # Display the image
    if in_jupyter or in_colab:

        from IPython.display import display
        display(image)
    else:

        image.show()



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):
        image = Image.fromarray(image)
    elif torch.is_tensor(image):
        image = Image.fromarray(image.permute(1, 2, 0).numpy().astype(np.uint8))

    plt.imshow(image)
    plt.axis('off')  # Turn off axis numbers
    plt.show()


def is_jupyter_notebook() -> bool:
    """
    Check if the code is running in a Jupyter notebook.

    Returns:
        bool: True if running in a Jupyter notebook, False otherwise.
    """
    try:
        from IPython import get_ipython
        if 'IPKernelApp' not in get_ipython().config:
            return False
        if 'ipykernel' in str(type(get_ipython())):
            return True  # Running in Jupyter Notebook
    except (NameError, AttributeError):
        return False  # Not running in Jupyter Notebook

    return False  # Default to False if none of the above conditions are met


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() -> 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: 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').

    Returns:
        str: The full path to the specified resource.
    """
    
    # Get the current working directory (assumed to be inside 'code' folder)
    current_dir = os.getcwd()

    # Get the directory one level up (the parent directory)
    parent_dir = os.path.dirname(current_dir)

    # Construct the path to the specified folder
    folder_path = os.path.join(parent_dir, path_type)

    # Construct the full path to the specific resource
    full_path = os.path.join(folder_path, name)

    return full_path


def get_model_path(model_name: str) -> str:
    """
    Get the path to the specified model folder.
    
    Args:
        model_name (str): Name of the model folder.

    Returns:
        str: Absolute path to the specified model folder.
    """
    
    # Directory of the current script
    current_script_dir = os.path.dirname(os.path.abspath(__file__))

    # Directory of the 'app' folder (parent of the 'my_model' folder)
    app_dir = os.path.dirname(os.path.dirname(current_script_dir))

    # Path to the 'models/{model_name}' folder
    model_path = os.path.join(app_dir, "models", 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.

    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():
        torch.cuda.empty_cache()
        torch.cuda.empty_cache()
        gc.collect()
        gc.collect()