Update app.py

app.py

@@ -14,6 +14,16 @@ import time
 from typing import Tuple, List, Dict
 import logging
 
+# Configure TensorFlow GPU memory growth to prevent memory conflicts
+try:
+    gpus = tf.config.list_physical_devices('GPU')
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        print(f"Found {len(gpus)} GPU(s), memory growth enabled")
+except Exception as e:
+    print(f"Error configuring GPU: {e}")
+
 # Import spaces correctly for ZeroGPU
 try:
     import spaces
@@ -57,37 +67,56 @@ def load_models():
         
         logger.info("Loading models from Hugging Face Hub...")
         
-        # Load YOLO Card Detection Model
-        card_model_path = hf_hub_download("Oamitai/card-detection", "best.pt")
-        card_detector = YOLO(card_model_path)
-        card_detector.conf = 0.5
-        
-        # Load YOLO Shape Detection Model
-        shape_model_path = hf_hub_download("Oamitai/shape-detection", "best.pt")
-        shape_detector = YOLO(shape_model_path)
-        shape_detector.conf = 0.5
-        
-        # Load Shape Classification Model
-        shape_classifier = tf.keras.models.load_model(
-            hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
-        )
-        
-        # Load Fill Classification Model
-        fill_classifier = tf.keras.models.load_model(
-            hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
-        )
+        # Load TensorFlow models first to better manage GPU memory
+        try:
+            # Load Shape Classification Model
+            logger.info("Loading shape classification model...")
+            shape_classifier = tf.keras.models.load_model(
+                hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
+            )
+            
+            # Load Fill Classification Model
+            logger.info("Loading fill classification model...")
+            fill_classifier = tf.keras.models.load_model(
+                hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
+            )
+            
+            # Cache TensorFlow models
+            _SHAPE_CLASSIFIER = shape_classifier
+            _FILL_CLASSIFIER = fill_classifier
+            
+        except Exception as tf_error:
+            logger.error(f"Error loading TensorFlow models: {tf_error}")
+            raise
         
-        # Use GPU if available
-        if torch.cuda.is_available():
-            logger.info("CUDA is available. Using GPU for inference.")
-            card_detector.to("cuda")
-            shape_detector.to("cuda")
+        # Add a small delay to ensure TensorFlow releases GPU resources
+        time.sleep(0.5)
         
-        # Cache the models
-        _CARD_DETECTOR = card_detector
-        _SHAPE_DETECTOR = shape_detector
-        _SHAPE_CLASSIFIER = shape_classifier
-        _FILL_CLASSIFIER = fill_classifier
+        # Now load YOLO models
+        try:
+            # Load YOLO Card Detection Model
+            logger.info("Loading card detection model...")
+            card_model_path = hf_hub_download("Oamitai/card-detection", "best.pt")
+            card_detector = YOLO(card_model_path)
+            card_detector.conf = 0.5
+            
+            # Load YOLO Shape Detection Model
+            logger.info("Loading shape detection model...")
+            shape_model_path = hf_hub_download("Oamitai/shape-detection", "best.pt")
+            shape_detector = YOLO(shape_model_path)
+            shape_detector.conf = 0.5
+            
+            # Use CPU initially for YOLO models to avoid GPU conflicts
+            card_detector.to("cpu")
+            shape_detector.to("cpu")
+            
+            # Cache PyTorch models
+            _CARD_DETECTOR = card_detector
+            _SHAPE_DETECTOR = shape_detector
+            
+        except Exception as pt_error:
+            logger.error(f"Error loading PyTorch models: {pt_error}")
+            raise
         
         logger.info("All models loaded successfully!")
         return card_detector, shape_detector, shape_classifier, fill_classifier
@@ -213,8 +242,31 @@ def predict_card_features(
         color_candidates.append(predict_color(shape_crop))
 
     # Use verbose=0 to suppress progress bar
-    fill_preds = fill_model.predict(np.array(fill_imgs), batch_size=len(fill_imgs), verbose=0)
-    shape_preds = shape_model.predict(np.array(shape_imgs), batch_size=len(shape_imgs), verbose=0)
+    # Add error handling for TensorFlow prediction
+    try:
+        fill_preds = fill_model.predict(np.array(fill_imgs), batch_size=len(fill_imgs), verbose=0)
+        shape_preds = shape_model.predict(np.array(shape_imgs), batch_size=len(shape_imgs), verbose=0)
+    except Exception as e:
+        logger.error(f"Error during TensorFlow prediction: {e}")
+        # Try with batch size of 1 as fallback
+        fill_preds = []
+        shape_preds = []
+        
+        for img in fill_imgs:
+            try:
+                pred = fill_model.predict(np.array([img]), verbose=0)
+                fill_preds.append(pred[0])
+            except Exception as e2:
+                logger.error(f"Fill prediction error: {e2}")
+                fill_preds.append(np.array([0.33, 0.33, 0.34]))  # Fallback with uniform probabilities
+        
+        for img in shape_imgs:
+            try:
+                pred = shape_model.predict(np.array([img]), verbose=0)
+                shape_preds.append(pred[0])
+            except Exception as e2:
+                logger.error(f"Shape prediction error: {e2}")
+                shape_preds.append(np.array([0.33, 0.33, 0.34]))  # Fallback with uniform probabilities
 
     fill_labels = ['empty', 'full', 'striped']
     shape_labels = ['diamond', 'oval', 'squiggle']
@@ -223,9 +275,20 @@ def predict_card_features(
     shape_result = [shape_labels[np.argmax(sp)] for sp in shape_preds]
 
     # Take the most common color/fill/shape across all shape detections for the card
-    final_color = max(set(color_candidates), key=color_candidates.count)
-    final_fill = max(set(fill_result), key=fill_result.count)
-    final_shape = max(set(shape_result), key=shape_result.count)
+    if color_candidates:
+        final_color = max(set(color_candidates), key=color_candidates.count)
+    else:
+        final_color = "unknown"
+        
+    if fill_result:
+        final_fill = max(set(fill_result), key=fill_result.count)
+    else:
+        final_fill = "unknown"
+        
+    if shape_result:
+        final_shape = max(set(shape_result), key=shape_result.count)
+    else:
+        final_shape = "unknown"
 
     return {
         'count': len(shape_boxes),
@@ -349,11 +412,11 @@ def optimize_image_size(image_array: np.ndarray, max_dim=1200) -> np.ndarray:
         return cv2.resize(image_array, (new_width, new_height), interpolation=cv2.INTER_AREA)
     return image_array
 
-@spaces.GPU
+# Make the process_image function use CPU for PyTorch models
 def process_image(input_image):
     """
-    Main processing function for SET detection.
-    Takes an input image, processes it, and returns the annotated image and status.
+    Main processing function for SET detection using CPU for both PyTorch 
+    and TensorFlow to avoid GPU conflicts.
     """
     if input_image is None:
         return None, "Please upload an image."
@@ -364,6 +427,10 @@ def process_image(input_image):
         # Load models
         card_detector, shape_detector, shape_model, fill_model = load_models()
         
+        # Force CPU mode for YOLO models (PyTorch)
+        card_detector.to("cpu")
+        shape_detector.to("cpu")
+        
         # Optimize image size
         optimized_img = optimize_image_size(input_image)
         
@@ -407,6 +474,15 @@ def process_image(input_image):
         logger.error(traceback.format_exc())
         return input_image, error_message
 
+# Wrap the CPU-based function with spaces.GPU
+@spaces.GPU
+def process_image_with_gpu(input_image):
+    """
+    Wrapper function that uses spaces.GPU decorator but internally
+    uses CPU processing to avoid GPU conflicts.
+    """
+    return process_image(input_image)
+
 # =============================================================================
 # GRADIO INTERFACE
 # =============================================================================
@@ -442,7 +518,7 @@ with gr.Blocks(title="SET Game Detector") as demo:
     
     # Function bindings inside the Blocks context
     find_sets_btn.click(
-        fn=process_image,
+        fn=process_image_with_gpu,  # Use the wrapper function
         inputs=[input_image],
         outputs=[output_image, status]
     )