Update app.py

app.py

@@ -7,24 +7,17 @@ from tensorflow.keras.models import load_model
 import torch
 from ultralytics import YOLO
 from itertools import combinations
-from pathlib import Path
 import gradio as gr
 import traceback
 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
+# Force CPU mode for TensorFlow
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+tf.config.set_visible_devices([], 'GPU')
+
+# Import spaces for ZeroGPU wrapper
 try:
     import spaces
 except ImportError:
@@ -52,7 +45,7 @@ _FILL_CLASSIFIER = None
 
 def load_models():
     """
-    Load all models needed for SET detection.
+    Load all models needed for SET detection in CPU-only mode.
     Returns tuple of (card_detector, shape_detector, shape_classifier, fill_classifier)
     """
     global _CARD_DETECTOR, _SHAPE_DETECTOR, _SHAPE_CLASSIFIER, _FILL_CLASSIFIER
@@ -67,58 +60,42 @@ def load_models():
         
         logger.info("Loading models from Hugging Face Hub...")
         
-        # 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
+        # Load Shape Classification Model (TensorFlow)
+        logger.info("Loading shape classification model...")
+        shape_classifier = load_model(
+            hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
+        )
         
-        # Add a small delay to ensure TensorFlow releases GPU resources
-        time.sleep(0.5)
+        # Load Fill Classification Model (TensorFlow)
+        logger.info("Loading fill classification model...")
+        fill_classifier = load_model(
+            hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
+        )
         
-        # 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
+        # Load YOLO Card Detection Model (PyTorch)
+        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 (PyTorch)
+        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
         
-        logger.info("All models loaded successfully!")
+        # Explicitly set to CPU mode
+        logger.info("Setting models to CPU mode...")
+        card_detector.to("cpu")
+        shape_detector.to("cpu")
+        
+        # Cache the models
+        _CARD_DETECTOR = card_detector
+        _SHAPE_DETECTOR = shape_detector
+        _SHAPE_CLASSIFIER = shape_classifier
+        _FILL_CLASSIFIER = fill_classifier
+        
+        logger.info("All models loaded successfully in CPU mode!")
         return card_detector, shape_detector, shape_classifier, fill_classifier
         
     except Exception as e:
@@ -241,32 +218,25 @@ def predict_card_features(
         shape_imgs.append(shape_crop_resized)
         color_candidates.append(predict_color(shape_crop))
 
-    # Use verbose=0 to suppress progress bar
-    # 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
+    # Handle TensorFlow prediction - process one image at a time to avoid memory issues
+    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 e:
+            logger.error(f"Fill prediction error: {e}")
+            fill_preds.append(np.array([0.33, 0.33, 0.34]))  # Fallback
+    
+    for img in shape_imgs:
+        try:
+            pred = shape_model.predict(np.array([img]), verbose=0)
+            shape_preds.append(pred[0])
+        except Exception as e:
+            logger.error(f"Shape prediction error: {e}")
+            shape_preds.append(np.array([0.33, 0.33, 0.34]))  # Fallback
 
     fill_labels = ['empty', 'full', 'striped']
     shape_labels = ['diamond', 'oval', 'squiggle']
@@ -412,11 +382,9 @@ 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
 
-# Make the process_image function use CPU for PyTorch models
 def process_image(input_image):
     """
-    Main processing function for SET detection using CPU for both PyTorch 
-    and TensorFlow to avoid GPU conflicts.
+    CPU-only processing function for SET detection.
     """
     if input_image is None:
         return None, "Please upload an image."
@@ -424,17 +392,13 @@ def process_image(input_image):
     try:
         start_time = time.time()
         
-        # Load models
+        # Load models (CPU-only)
         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)
         
-        # Convert to BGR if needed (OpenCV format)
+        # Convert to BGR (OpenCV format)
         if len(optimized_img.shape) == 3 and optimized_img.shape[2] == 4:  # RGBA
             optimized_img = cv2.cvtColor(optimized_img, cv2.COLOR_RGBA2BGR)
         elif len(optimized_img.shape) == 3 and optimized_img.shape[2] == 3:
@@ -474,17 +438,17 @@ def process_image(input_image):
         logger.error(traceback.format_exc())
         return input_image, error_message
 
-# Wrap the CPU-based function with spaces.GPU
+# Keep the spaces.GPU decorator for ZeroGPU API but use CPU internally
 @spaces.GPU
-def process_image_with_gpu(input_image):
+def process_image_wrapper(input_image):
     """
-    Wrapper function that uses spaces.GPU decorator but internally
-    uses CPU processing to avoid GPU conflicts.
+    Wrapper for process_image that uses the spaces.GPU decorator
+    but internally works in CPU-only mode.
     """
     return process_image(input_image)
 
 # =============================================================================
-# GRADIO INTERFACE
+# SIMPLIFIED GRADIO INTERFACE
 # =============================================================================
 with gr.Blocks(title="SET Game Detector") as demo:
     gr.HTML("""
@@ -516,9 +480,9 @@ with gr.Blocks(title="SET Game Detector") as demo:
                 interactive=False
             )
     
-    # Function bindings inside the Blocks context
+    # Function bindings
     find_sets_btn.click(
-        fn=process_image_with_gpu,  # Use the wrapper function
+        fn=process_image_wrapper,
         inputs=[input_image],
         outputs=[output_image, status]
     )