wip

age_estimation/age_estimation.py

@@ -5,6 +5,8 @@ from utils.image_utils import preprocess_image, get_image_from_input
 from utils.face_detector import load_face_detector
 from .predict import predict_age
 
+model = load_model("cpu")  # Load the model on CPU by default
+
 
 def age_estimation(input_type, uploaded_image, image_url, base64_string):
     """
@@ -32,9 +34,10 @@ def age_estimation(input_type, uploaded_image, image_url, base64_string):
 
     try:
         face_detector = load_face_detector()
-
+        global model
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = load_model(device)
+        print(f"Using device: {device}")  # Debug print
+        model = model.to(device)
 
         # Preprocess the image (convert PIL to numpy, ensure RGB)
         processed_image = preprocess_image(image)

app.py

@@ -1,24 +1,13 @@
 # Standard library imports
 import os
 import gradio as gr
-import spaces
-
-
 
 # Local imports
-from age_estimation.age_estimation import age_estimation as _age_estimation
-from detection.face_detection import face_detection 
-from detection.object_detection import object_detection as _object_detection
+from detection.object_detection import object_detection
+from age_estimation.age_estimation import age_estimation
+from detection.face_detection import face_detection
 from utils.ui_utils import update_input_visibility
 
-@spaces.GPU
-def age_estimation(input_type, uploaded_image, image_url, base64_string):
-    return _age_estimation(input_type, uploaded_image, image_url, base64_string)
-
-@spaces.GPU
-def object_detection(input_type, uploaded_image, image_url, base64_string):
-    return _object_detection(input_type, uploaded_image, image_url, base64_string)
-
 with gr.Blocks() as demo:
     # Add a title to the interface
     gr.Markdown("# Computer Vision Tools")

detection/object_detection.py

@@ -11,12 +11,25 @@ YOLO_MODEL = "yolo11n.pt"
 
 # Load the YOLO model globally to avoid reloading on each function call
 # Using a common pre-trained YOLO nano model ('yolov8n.pt')
-try:
-    model = YOLO(YOLO_MODEL)
-    print("YOLO model loaded successfully.")
-except Exception as e:
-    print(f"Error loading YOLO model: {e}")
-    model = None  # Set model to None if loading fails
+model = None
+
+
+def load_model():
+    """
+    Load the YOLO model (YOLOv11 nano) for object detection.
+    This function is called once at the start to avoid reloading the model multiple times.
+    """
+    try:
+        global model
+        if model is not None:
+            print("YOLO model already loaded.")
+            return
+        model = YOLO(YOLO_MODEL)
+        print("YOLO model loaded successfully.")
+        return
+    except Exception as e:
+        print(f"Error loading YOLO model: {e}")
+        return
 
 
 def object_detection(input_type, uploaded_image, image_url, base64_string):
@@ -34,6 +47,8 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
             - numpy.ndarray: The image with detected objects drawn on it, or None if an error occurred or model not loaded.
             - dict: A dictionary containing the raw detection data (bounding boxes, classes, scores), or None.
     """
+    load_model()  # Load the model if not already loaded
+
     if model is None:
         print("YOLO model is not loaded. Cannot perform object detection.")
         return None, None  # Return None for both outputs