Update app.py

app.py

@@ -1,273 +1,726 @@
-# IMPORTANT: Import spaces first to prevent premature CUDA initialization.
-import spaces
-
-import gradio as gr
-from huggingface_hub import hf_hub_download
-import torch
+import os
 import cv2
 import numpy as np
+import pandas as pd
 import tensorflow as tf
 from tensorflow.keras.models import load_model
+import torch
 from ultralytics import YOLO
+from itertools import combinations
+from pathlib import Path
 from PIL import Image
+import gradio as gr
+import functools
 import traceback
-import json
-import pandas as pd
-from itertools import combinations
+import time
+from typing import Tuple, List, Dict
+import logging
 
 # =============================================================================
-#                           MODEL LOADING
+# LOGGING CONFIGURATION
 # =============================================================================
+logging.basicConfig(level=logging.INFO, 
+                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("set_detector")
 
-# Load YOLO Card Detection Model from HuggingFace Hub
-card_model_path = hf_hub_download("Oamitai/card-detection", "best.pt")
-card_detection_model = YOLO(card_model_path)
-card_detection_model.conf = 0.5
+# =============================================================================
+# MODEL PATHS & LOADING
+# =============================================================================
+# For loading models from Hugging Face Hub
+HF_MODEL_REPO_PREFIX = "Omamitai"
+
+# Define model repos and paths
+CARD_DETECTION_REPO = f"{HF_MODEL_REPO_PREFIX}/card-detection"
+SHAPE_DETECTION_REPO = f"{HF_MODEL_REPO_PREFIX}/shape-detection"
+SHAPE_CLASSIFICATION_REPO = f"{HF_MODEL_REPO_PREFIX}/shape-classification"
+FILL_CLASSIFICATION_REPO = f"{HF_MODEL_REPO_PREFIX}/fill-classification"
+
+# Model filenames
+CARD_MODEL_FILENAME = "best.pt"
+SHAPE_MODEL_FILENAME = "best.pt"
+SHAPE_CLASS_MODEL_FILENAME = "shape_model.keras"
+FILL_CLASS_MODEL_FILENAME = "fill_model.keras"
+
+# For local testing: fallback to local models if HF downloads fail
+# Use the local directory structure as fallback
+if os.path.exists("/home/user"):  # Check if we're on HF Spaces
+    local_base_dir = Path("/home/user/app/models")
+else:
+    local_base_dir = Path(os.path.dirname(os.path.abspath(__file__))) / "models"
+
+local_char_path = local_base_dir / "Characteristics" / "11022025"
+local_shape_path = local_base_dir / "Shape" / "15052024"
+local_card_path = local_base_dir / "Card" / "16042024"
+
+# Global variables for model caching
+_MODEL_SHAPE = None
+_MODEL_FILL = None
+_DETECTOR_CARD = None
+_DETECTOR_SHAPE = None
+_MODELS_LOADED = False
+_MODEL_LOADING_ERROR = None
+
+def load_classification_models() -> Tuple[tf.keras.Model, tf.keras.Model]:
+    """
+    Loads the Keras models for 'shape' and 'fill' classification from HuggingFace Hub.
+    Returns (shape_model, fill_model).
+    """
+    global _MODEL_SHAPE, _MODEL_FILL, _MODEL_LOADING_ERROR
+    
+    # If models are already loaded, return them
+    if _MODEL_SHAPE is not None and _MODEL_FILL is not None:
+        return _MODEL_SHAPE, _MODEL_FILL
+    
+    try:
+        from huggingface_hub import hf_hub_download
+        
+        # Try to download from HuggingFace Hub
+        logger.info(f"Downloading shape classification model from {SHAPE_CLASSIFICATION_REPO}...")
+        shape_model_path = hf_hub_download(
+            repo_id=SHAPE_CLASSIFICATION_REPO, 
+            filename=SHAPE_CLASS_MODEL_FILENAME,
+            cache_dir="./hf_cache"
+        )
+        
+        logger.info(f"Downloading fill classification model from {FILL_CLASSIFICATION_REPO}...")
+        fill_model_path = hf_hub_download(
+            repo_id=FILL_CLASSIFICATION_REPO, 
+            filename=FILL_CLASS_MODEL_FILENAME,
+            cache_dir="./hf_cache"
+        )
+        
+        # Load the models
+        logger.info("Loading classification models...")
+        model_shape = load_model(str(shape_model_path))
+        model_fill = load_model(str(fill_model_path))
+        
+        logger.info("Classification models loaded successfully")
+        _MODEL_SHAPE, _MODEL_FILL = model_shape, model_fill
+        return model_shape, model_fill
+    
+    except Exception as e:
+        error_msg = f"Error downloading classification models from HF Hub: {str(e)}"
+        logger.error(error_msg)
+        
+        # Try fallback to local files
+        try:
+            logger.info("Trying fallback to local model files...")
+            shape_model_path = local_char_path / SHAPE_CLASS_MODEL_FILENAME
+            fill_model_path = local_char_path / FILL_CLASS_MODEL_FILENAME
+            
+            if not shape_model_path.exists() or not fill_model_path.exists():
+                raise FileNotFoundError("Local model files not found")
+                
+            # Load the models
+            model_shape = load_model(str(shape_model_path))
+            model_fill = load_model(str(fill_model_path))
+            
+            logger.info("Classification models loaded successfully from local files")
+            _MODEL_SHAPE, _MODEL_FILL = model_shape, model_fill
+            return model_shape, model_fill
+            
+        except Exception as fallback_error:
+            error_msg = f"{error_msg}\nFallback to local files also failed: {str(fallback_error)}"
+            logger.error(error_msg)
+            _MODEL_LOADING_ERROR = error_msg
+            return None, None
 
-# Load YOLO Shape Detection Model from HuggingFace Hub
-shape_model_path = hf_hub_download("Oamitai/shape-detection", "best.pt")
-shape_detection_model = YOLO(shape_model_path)
-shape_detection_model.conf = 0.5
+def load_detection_models() -> Tuple[YOLO, YOLO]:
+    """
+    Loads the YOLO detection models for cards and shapes from HuggingFace Hub.
+    Returns (card_detector, shape_detector).
+    """
+    global _DETECTOR_CARD, _DETECTOR_SHAPE, _MODEL_LOADING_ERROR
+    
+    # If models are already loaded, return them
+    if _DETECTOR_CARD is not None and _DETECTOR_SHAPE is not None:
+        return _DETECTOR_CARD, _DETECTOR_SHAPE
+    
+    try:
+        from huggingface_hub import hf_hub_download
+        
+        # Try to download from HuggingFace Hub
+        logger.info(f"Downloading card detection model from {CARD_DETECTION_REPO}...")
+        card_model_path = hf_hub_download(
+            repo_id=CARD_DETECTION_REPO, 
+            filename=CARD_MODEL_FILENAME,
+            cache_dir="./hf_cache"
+        )
+        
+        logger.info(f"Downloading shape detection model from {SHAPE_DETECTION_REPO}...")
+        shape_model_path = hf_hub_download(
+            repo_id=SHAPE_DETECTION_REPO, 
+            filename=SHAPE_MODEL_FILENAME,
+            cache_dir="./hf_cache"
+        )
+        
+        # Load the models
+        logger.info("Loading detection models...")
+        detector_shape = YOLO(str(shape_model_path))
+        detector_shape.conf = 0.5
+        detector_card = YOLO(str(card_model_path))
+        detector_card.conf = 0.5
+        
+        # Use GPU if available
+        if torch.cuda.is_available():
+            logger.info("CUDA is available. Using GPU for inference.")
+            detector_card.to("cuda")
+            detector_shape.to("cuda")
+        else:
+            logger.info("CUDA is not available. Using CPU for inference.")
+            
+        logger.info("Detection models loaded successfully")
+        _DETECTOR_CARD, _DETECTOR_SHAPE = detector_card, detector_shape
+        return detector_card, detector_shape
+        
+    except Exception as e:
+        error_msg = f"Error downloading detection models from HF Hub: {str(e)}"
+        logger.error(error_msg)
+        
+        # Try fallback to local files
+        try:
+            logger.info("Trying fallback to local model files...")
+            shape_model_path = local_shape_path / SHAPE_MODEL_FILENAME
+            card_model_path = local_card_path / CARD_MODEL_FILENAME
+            
+            if not shape_model_path.exists() or not card_model_path.exists():
+                raise FileNotFoundError("Local model files not found")
+                
+            # Load the models
+            detector_shape = YOLO(str(shape_model_path))
+            detector_shape.conf = 0.5
+            detector_card = YOLO(str(card_model_path))
+            detector_card.conf = 0.5
+            
+            # Use GPU if available
+            if torch.cuda.is_available():
+                logger.info("CUDA is available. Using GPU for inference.")
+                detector_card.to("cuda")
+                detector_shape.to("cuda")
+                
+            logger.info("Detection models loaded successfully from local files")
+            _DETECTOR_CARD, _DETECTOR_SHAPE = detector_card, detector_shape
+            return detector_card, detector_shape
+            
+        except Exception as fallback_error:
+            error_msg = f"{error_msg}\nFallback to local files also failed: {str(fallback_error)}"
+            logger.error(error_msg)
+            _MODEL_LOADING_ERROR = error_msg
+            return None, None
 
-# Load Shape Classification Model (Keras) from HuggingFace Hub
-shape_classification_model = load_model(
-    hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
-)
+def load_all_models() -> bool:
+    """
+    Loads all required models and returns True if successful.
+    """
+    global _MODELS_LOADED, _MODEL_LOADING_ERROR
+    
+    if _MODELS_LOADED:
+        return True
+    
+    try:
+        model_shape, model_fill = load_classification_models()
+        detector_card, detector_shape = load_detection_models()
+        
+        models_loaded = all([model_shape, model_fill, detector_card, detector_shape])
+        _MODELS_LOADED = models_loaded
+        
+        if not models_loaded and _MODEL_LOADING_ERROR is None:
+            _MODEL_LOADING_ERROR = "Unknown error loading models"
+            
+        return models_loaded
+    except Exception as e:
+        error_msg = f"Error loading models: {str(e)}"
+        logger.error(error_msg)
+        _MODEL_LOADING_ERROR = error_msg
+        return False
 
-# Load Fill Classification Model (Keras) from HuggingFace Hub
-fill_classification_model = load_model(
-    hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
-)
+def get_model_status() -> str:
+    """
+    Returns a status message about the models.
+    """
+    if _MODELS_LOADED:
+        return "All models loaded successfully!"
+    elif _MODEL_LOADING_ERROR:
+        return f"Error: {_MODEL_LOADING_ERROR}"
+    else:
+        return "Models not loaded yet. Click 'Load Models' to preload them."
 
 # =============================================================================
-#                     UTILITY & PROCESSING FUNCTIONS
+# UTILITY & DETECTION FUNCTIONS
 # =============================================================================
-
-def check_and_rotate_input_image(board_image: np.ndarray, detector) -> (np.ndarray, bool):
+def verify_and_rotate_image(board_image: np.ndarray, card_detector: YOLO) -> Tuple[np.ndarray, bool]:
     """
-    Detect card regions and determine if the image needs to be rotated.
+    Checks if the detected cards are oriented primarily vertically or horizontally.
+    If they're vertical, rotates the board_image 90 degrees clockwise for consistent processing.
+    Returns (possibly_rotated_image, was_rotated_flag).
     """
-    card_results = detector(board_image)
-    card_boxes = card_results[0].boxes.xyxy.cpu().numpy().astype(int)
-    if card_boxes.size == 0:
+    detection = card_detector(board_image)
+    boxes = detection[0].boxes.xyxy.cpu().numpy().astype(int)
+    if boxes.size == 0:
         return board_image, False
 
-    widths = card_boxes[:, 2] - card_boxes[:, 0]
-    heights = card_boxes[:, 3] - card_boxes[:, 1]
+    widths = boxes[:, 2] - boxes[:, 0]
+    heights = boxes[:, 3] - boxes[:, 1]
+
+    # Rotate if average height > average width
     if np.mean(heights) > np.mean(widths):
         return cv2.rotate(board_image, cv2.ROTATE_90_CLOCKWISE), True
-    return board_image, False
+    else:
+        return board_image, False
 
-def restore_original_orientation(image: np.ndarray, was_rotated: bool) -> np.ndarray:
+def restore_orientation(img: np.ndarray, was_rotated: bool) -> np.ndarray:
     """
-    Restore the original orientation of the image if it was rotated.
+    Restores original orientation if the image was previously rotated.
     """
     if was_rotated:
-        return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
-    return image
+        return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    return img
 
-def predict_color(shape_image: np.ndarray) -> str:
+def predict_color(img_bgr: np.ndarray) -> str:
     """
-    Determine the dominant color in a shape image using HSV thresholds.
+    Rough color classification using HSV thresholds to differentiate 'red', 'green', 'purple'.
     """
-    hsv_image = cv2.cvtColor(shape_image, cv2.COLOR_BGR2HSV)
-    green_mask = cv2.inRange(hsv_image, np.array([40, 50, 50]), np.array([80, 255, 255]))
-    purple_mask = cv2.inRange(hsv_image, np.array([120, 50, 50]), np.array([160, 255, 255]))
-    red_mask1 = cv2.inRange(hsv_image, np.array([0, 50, 50]), np.array([10, 255, 255]))
-    red_mask2 = cv2.inRange(hsv_image, np.array([170, 50, 50]), np.array([180, 255, 255]))
-    red_mask = cv2.bitwise_or(red_mask1, red_mask2)
-    
-    color_counts = {
-        'green': cv2.countNonZero(green_mask),
-        'purple': cv2.countNonZero(purple_mask),
-        'red': cv2.countNonZero(red_mask)
+    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
+    mask_green = cv2.inRange(hsv, np.array([40, 50, 50]), np.array([80, 255, 255]))
+    mask_purple = cv2.inRange(hsv, np.array([120, 50, 50]), np.array([160, 255, 255]))
+
+    # Red can wrap around hue=0, so we combine both ends
+    mask_red1 = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([10, 255, 255]))
+    mask_red2 = cv2.inRange(hsv, np.array([170, 50, 50]), np.array([180, 255, 255]))
+    mask_red = cv2.bitwise_or(mask_red1, mask_red2)
+
+    counts = {
+        "green": cv2.countNonZero(mask_green),
+        "purple": cv2.countNonZero(mask_purple),
+        "red": cv2.countNonZero(mask_red),
     }
-    return max(color_counts, key=color_counts.get)
+    return max(counts, key=counts.get)
 
-def predict_card_features(card_image: np.ndarray, shape_detector, fill_model, shape_model, box: list) -> dict:
+def detect_cards(board_img: np.ndarray, card_detector: YOLO) -> List[Tuple[np.ndarray, List[int]]]:
     """
-    Detect and classify features on a card image.
+    Runs YOLO on the board_img to detect card bounding boxes.
+    Returns a list of (card_image, [x1, y1, x2, y2]) for each detected card.
     """
-    shape_results = shape_detector(card_image)
-    card_h, card_w = card_image.shape[:2]
-    card_area = card_w * card_h
-
-    filtered_boxes = [
-        [int(x1), int(y1), int(x2), int(y2)]
-        for x1, y1, x2, y2 in shape_results[0].boxes.xyxy.cpu().numpy()
-        if (x2 - x1) * (y2 - y1) > 0.03 * card_area
-    ]
-
-    if not filtered_boxes:
-        return {'count': 0, 'color': 'unknown', 'fill': 'unknown', 'shape': 'unknown', 'box': box}
-
-    fill_input_shape = fill_model.input_shape[1:3]
-    shape_input_shape = shape_model.input_shape[1:3]
-    fill_imgs, shape_imgs, color_list = [], [], []
-
-    for fb in filtered_boxes:
-        x1, y1, x2, y2 = fb
-        shape_img = card_image[y1:y2, x1:x2]
-        fill_img = cv2.resize(shape_img, tuple(fill_input_shape)) / 255.0
-        shape_img_resized = cv2.resize(shape_img, tuple(shape_input_shape)) / 255.0
-        fill_imgs.append(fill_img)
-        shape_imgs.append(shape_img_resized)
-        color_list.append(predict_color(shape_img))
-        
-    fill_imgs = np.array(fill_imgs)
-    shape_imgs = np.array(shape_imgs)
-    
-    fill_preds = fill_model.predict(fill_imgs, batch_size=len(fill_imgs))
-    shape_preds = shape_model.predict(shape_imgs, batch_size=len(shape_imgs))
-    
-    fill_labels_list = ['empty', 'full', 'striped']
-    shape_labels_list = ['diamond', 'oval', 'squiggle']
-    
-    predicted_fill = [fill_labels_list[np.argmax(pred)] for pred in fill_preds]
-    predicted_shape = [shape_labels_list[np.argmax(pred)] for pred in shape_preds]
+    result = card_detector(board_img)
+    boxes = result[0].boxes.xyxy.cpu().numpy().astype(int)
+    detected_cards = []
 
-    color_label = max(set(color_list), key=color_list.count)
-    fill_label = max(set(predicted_fill), key=predicted_fill.count)
-    shape_label = max(set(predicted_shape), key=predicted_shape.count)
-    
-    return {'count': len(filtered_boxes), 'color': color_label,
-            'fill': fill_label, 'shape': shape_label, 'box': box}
+    for x1, y1, x2, y2 in boxes:
+        detected_cards.append((board_img[y1:y2, x1:x2], [x1, y1, x2, y2]))
+    return detected_cards
 
-def is_set(cards: list) -> bool:
+def predict_card_features(
+    card_img: np.ndarray,
+    shape_detector: YOLO,
+    fill_model: tf.keras.Model,
+    shape_model: tf.keras.Model,
+    card_box: List[int]
+) -> Dict:
     """
-    Check if a group of cards forms a valid set. For each feature,
-    values must be all identical or all distinct.
+    Predicts the 'count', 'color', 'fill', 'shape' features for a single card.
+    It uses a shape_detector YOLO model to locate shapes, then passes them to fill_model and shape_model.
     """
-    for feature in ['Count', 'Color', 'Fill', 'Shape']:
-        if len({card[feature] for card in cards}) not in [1, 3]:
-            return False
-    return True
+    # Detect shapes on the card
+    shape_detections = shape_detector(card_img)
+    c_h, c_w = card_img.shape[:2]
+    card_area = c_w * c_h
+
+    # Filter out spurious shape detections
+    shape_boxes = []
+    for coords in shape_detections[0].boxes.xyxy.cpu().numpy():
+        x1, y1, x2, y2 = coords.astype(int)
+        if (x2 - x1) * (y2 - y1) > 0.03 * card_area:
+            shape_boxes.append([x1, y1, x2, y2])
+
+    if not shape_boxes:
+        return {
+            'count': 0,
+            'color': 'unknown',
+            'fill': 'unknown',
+            'shape': 'unknown',
+            'box': card_box
+        }
+
+    fill_input_size = fill_model.input_shape[1:3]
+    shape_input_size = shape_model.input_shape[1:3]
+    fill_imgs = []
+    shape_imgs = []
+    color_candidates = []
+
+    # Prepare each detected shape region for classification
+    for sb in shape_boxes:
+        sx1, sy1, sx2, sy2 = sb
+        shape_crop = card_img[sy1:sy2, sx1:sx2]
 
-def find_sets(card_df: pd.DataFrame) -> list:
+        fill_crop = cv2.resize(shape_crop, fill_input_size) / 255.0
+        shape_crop_resized = cv2.resize(shape_crop, shape_input_size) / 255.0
+
+        fill_imgs.append(fill_crop)
+        shape_imgs.append(shape_crop_resized)
+        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)
+
+    fill_labels = ['empty', 'full', 'striped']
+    shape_labels = ['diamond', 'oval', 'squiggle']
+
+    fill_result = [fill_labels[np.argmax(fp)] for fp in fill_preds]
+    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)
+
+    return {
+        'count': len(shape_boxes),
+        'color': final_color,
+        'fill': final_fill,
+        'shape': final_shape,
+        'box': card_box
+    }
+
+def classify_cards_on_board(
+    board_img: np.ndarray,
+    card_detector: YOLO,
+    shape_detector: YOLO,
+    fill_model: tf.keras.Model,
+    shape_model: tf.keras.Model
+) -> pd.DataFrame:
     """
-    Iterate over all combinations of three cards to identify valid sets.
+    Detects cards on the board, then classifies each card's features.
+    Returns a DataFrame with columns: 'Count', 'Color', 'Fill', 'Shape', 'Coordinates'.
     """
-    sets_found = []
-    for combo in combinations(card_df.iterrows(), 3):
-        cards = [entry[1] for entry in combo]
-        if is_set(cards):
-            sets_found.append({
-                'set_indices': [entry[0] for entry in combo],
-                'cards': [{feature: card[feature] for feature in 
-                           ['Count', 'Color', 'Fill', 'Shape', 'Coordinates']} for card in cards]
-            })
-    return sets_found
+    detected_cards = detect_cards(board_img, card_detector)
+    card_rows = []
 
-def detect_cards_from_image(board_image: np.ndarray, detector) -> list:
+    for (card_img, box) in detected_cards:
+        card_feats = predict_card_features(card_img, shape_detector, fill_model, shape_model, box)
+        card_rows.append({
+            "Count": card_feats['count'],
+            "Color": card_feats['color'],
+            "Fill": card_feats['fill'],
+            "Shape": card_feats['shape'],
+            "Coordinates": card_feats['box']
+        })
+
+    return pd.DataFrame(card_rows)
+
+def valid_set(cards: List[dict]) -> bool:
     """
-    Extract card regions from the board image using the YOLO card detection model.
+    Checks if the given 3 cards collectively form a valid SET.
     """
-    card_results = detector(board_image)
-    card_boxes = card_results[0].boxes.xyxy.cpu().numpy().astype(int)
-    return [(board_image[y1:y2, x1:x2], [x1, y1, x2, y2]) for x1, y1, x2, y2 in card_boxes]
+    for feature in ["Count", "Color", "Fill", "Shape"]:
+        if len({card[feature] for card in cards}) not in (1, 3):
+            return False
+    return True
 
-def classify_cards_from_board_image(board_image: np.ndarray, card_detector, shape_detector, fill_model, shape_model) -> pd.DataFrame:
+def locate_all_sets(cards_df: pd.DataFrame) -> List[dict]:
     """
-    Detect cards from the board image and classify their features.
+    Finds all possible SETs from the card DataFrame.
+    Each SET is a dictionary with 'set_indices' and 'cards' fields.
     """
-    cards = detect_cards_from_image(board_image, card_detector)
-    card_data = []
-    for card_image, box in cards:
-        features = predict_card_features(card_image, shape_detector, fill_model, shape_model, box)
-        card_data.append({
-            "Count": features['count'],
-            "Color": features['color'],
-            "Fill": features['fill'],
-            "Shape": features['shape'],
-            "Coordinates": f"{box[0]}, {box[1]}, {box[2]}, {box[3]}"
-        })
-    return pd.DataFrame(card_data)
+    found_sets = []
+    for combo in combinations(cards_df.iterrows(), 3):
+        cards = [c[1] for c in combo]  # c is (index, row)
+        if valid_set(cards):
+            found_sets.append({
+                'set_indices': [c[0] for c in combo],
+                'cards': [
+                    {f: card[f] for f in ['Count', 'Color', 'Fill', 'Shape', 'Coordinates']}
+                    for card in cards
+                ]
+            })
+    return found_sets
 
-def draw_sets_on_image(board_image: np.ndarray, sets_info: list) -> np.ndarray:
+def draw_detected_sets(board_img: np.ndarray, sets_detected: List[dict]) -> np.ndarray:
     """
-    Draw bounding boxes and labels for each detected set on the board image.
+    Annotates the board image with bounding boxes for each detected SET.
+    Each SET is drawn in a different color and offset (thickness & expansion) 
+    so that overlapping sets are visible.
     """
-    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255),
-              (255, 255, 0), (255, 0, 255), (0, 255, 255)]
+    # Some distinct BGR colors
+    colors = [
+        (255, 0, 0), (0, 255, 0), (0, 0, 255),
+        (255, 255, 0), (255, 0, 255), (0, 255, 255)
+    ]
     base_thickness = 8
     base_expansion = 5
-    for index, set_info in enumerate(sets_info):
-        color = colors[index % len(colors)]
-        thickness = base_thickness + 2 * index
-        expansion = base_expansion + 15 * index
-        for i, card in enumerate(set_info['cards']):
-            coordinates = list(map(int, card['Coordinates'].split(',')))
-            x1, y1, x2, y2 = coordinates
-            x1_exp = max(0, x1 - expansion)
-            y1_exp = max(0, y1 - expansion)
-            x2_exp = min(board_image.shape[1], x2 + expansion)
-            y2_exp = min(board_image.shape[0], y2 + expansion)
-            cv2.rectangle(board_image, (x1_exp, y1_exp), (x2_exp, y2_exp), color, thickness)
+
+    for idx, single_set in enumerate(sets_detected):
+        color = colors[idx % len(colors)]
+        thickness = base_thickness + 2 * idx
+        expansion = base_expansion + 15 * idx
+
+        for i, card_info in enumerate(single_set["cards"]):
+            x1, y1, x2, y2 = card_info["Coordinates"]
+            # Expand the bounding box slightly
+            x1e = max(0, x1 - expansion)
+            y1e = max(0, y1 - expansion)
+            x2e = min(board_img.shape[1], x2 + expansion)
+            y2e = min(board_img.shape[0], y2 + expansion)
+
+            cv2.rectangle(board_img, (x1e, y1e), (x2e, y2e), color, thickness)
+
+            # Label only the first card's box with "Set <number>"
             if i == 0:
-                cv2.putText(board_image, f"Set {index + 1}", (x1_exp, y1_exp - 10),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, thickness)
-    return board_image
+                cv2.putText(
+                    board_img,
+                    f"Set {idx + 1}",
+                    (x1e, y1e - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.9,
+                    color,
+                    thickness
+                )
+    return board_img
 
-def classify_and_find_sets_from_array(board_image: np.ndarray, card_detector, shape_detector, fill_model, shape_model) -> (list, np.ndarray):
+def identify_sets_from_image(
+    board_img: np.ndarray
+) -> Tuple[List[dict], np.ndarray, str]:
     """
-    Process the input image: adjust orientation, classify card features, detect sets, and annotate the image.
+    End-to-end pipeline to classify cards on the board and detect valid sets.
+    Returns a tuple of (list of sets, annotated image, status message).
     """
-    processed_image, was_rotated = check_and_rotate_input_image(board_image, card_detector)
-    card_df = classify_cards_from_board_image(processed_image, card_detector, shape_detector, fill_model, shape_model)
-    sets_found = find_sets(card_df)
-    annotated_image = draw_sets_on_image(processed_image.copy(), sets_found)
-    final_image = restore_original_orientation(annotated_image, was_rotated)
-    return sets_found, final_image
+    # Load models
+    if not load_all_models():
+        error_msg = _MODEL_LOADING_ERROR or "Error: Could not load models."
+        return [], board_img, error_msg
+    
+    card_detector, shape_detector = _DETECTOR_CARD, _DETECTOR_SHAPE
+    model_shape, model_fill = _MODEL_SHAPE, _MODEL_FILL
+    
+    # Convert image to BGR if needed (OpenCV format)
+    if len(board_img.shape) == 3 and board_img.shape[2] == 4:  # RGBA
+        board_img = cv2.cvtColor(board_img, cv2.COLOR_RGBA2BGR)
+    elif len(board_img.shape) == 3 and board_img.shape[2] == 3:
+        # We assume the image is already in BGR format (OpenCV standard)
+        # If it's in RGB format (common from web uploads), we'll convert it
+        board_img = cv2.cvtColor(board_img, cv2.COLOR_RGB2BGR)
+    else:
+        return [], board_img, "Error: Unsupported image format. Please upload a color image."
+    
+    # 1. Check and fix orientation if needed
+    processed, was_rotated = verify_and_rotate_image(board_img, card_detector)
 
-# =============================================================================
-#                      GRADIO INFERENCE FUNCTION
-# =============================================================================
+    # 2. Verify that cards are present
+    cards = detect_cards(processed, card_detector)
+    if not cards:
+        return [], cv2.cvtColor(board_img, cv2.COLOR_BGR2RGB), "No cards detected in the image. Please check that it's a SET game board."
+
+    # 3. Classify each card's features, then find sets
+    df_cards = classify_cards_on_board(processed, card_detector, shape_detector, model_fill, model_shape)
+    found_sets = locate_all_sets(df_cards)
 
-@spaces.GPU()
-def detect_sets(input_image: Image.Image):
+    if not found_sets:
+        return [], cv2.cvtColor(processed, cv2.COLOR_BGR2RGB), "Cards detected, but no valid SETs found. You may need to add more cards to the table!"
+
+    # 4. Draw sets on a copy of the image
+    annotated = draw_detected_sets(processed.copy(), found_sets)
+
+    # 5. Restore orientation if we rotated earlier
+    final_output = restore_orientation(annotated, was_rotated)
+    
+    # Convert back to RGB for display
+    final_output_rgb = cv2.cvtColor(final_output, cv2.COLOR_BGR2RGB)
+    
+    return found_sets, final_output_rgb, f"Found {len(found_sets)} SET(s) in the image."
+
+def optimize_image_size(image_array: np.ndarray, max_dim=1200) -> np.ndarray:
     """
-    Process an uploaded image and return the annotated image along with detected sets info.
+    Resizes an image if its largest dimension exceeds max_dim, to reduce processing time.
     """
-    try:
-        # Convert the PIL image to OpenCV BGR format
-        image_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
-        # Run the detection pipeline
-        sets_info, annotated_image = classify_and_find_sets_from_array(
-            image_cv,
-            card_detection_model,
-            shape_detection_model,
-            fill_classification_model,
-            shape_classification_model
-        )
-        # Convert annotated image back to RGB for display
-        annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
-        return annotated_image_rgb, json.dumps(sets_info, indent=2)
-    except Exception:
-        return None, f"Error occurred: {traceback.format_exc()}"
+    if image_array is None:
+        return None
+        
+    height, width = image_array.shape[:2]
+    if max(width, height) > max_dim:
+        if width > height:
+            new_width = max_dim
+            new_height = int(height * (max_dim / width))
+        else:
+            new_height = max_dim
+            new_width = int(width * (max_dim / height))
+
+        return cv2.resize(image_array, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    return image_array
 
 # =============================================================================
-#                           GRADIO INTERFACE
+# MAIN PROCESSING FUNCTIONS FOR GRADIO
 # =============================================================================
+def preload_models():
+    """
+    Function to preload models and return status.
+    """
+    try:
+        if load_all_models():
+            return "Models loaded successfully! Ready to detect SETs."
+        else:
+            return f"Error loading models: {_MODEL_LOADING_ERROR or 'Unknown error'}"
+    except Exception as e:
+        return f"Error loading models: {str(e)}"
 
-with gr.Blocks(css="#col-container { margin: 0 auto; max-width: 800px; }") as demo:
-    gr.Markdown("# Set Game Detector\nUpload an image of a Set game board to detect valid sets.")
-    
-    with gr.Row(elem_id="col-container"):
-        image_input = gr.Image(label="Upload Set Game Board", type="pil")
-        detect_button = gr.Button("Detect Sets")
+@spaces.GPU
+def process_set_image(input_image):
+    """
+    Main processing function for the Gradio interface.
+    Takes an input image, processes it to find SETs, and returns the output image and status.
     
-    with gr.Row():
-        result_image = gr.Image(label="Annotated Image")
-        result_info = gr.JSON(label="Detected Sets Info")
+    Uses @spaces.GPU for Hugging Face Spaces zero-GPU optimization.
+    """
+    if input_image is None:
+        return None, "Please upload an image."
     
-    detect_button.click(
-        detect_sets,
-        inputs=[image_input],
-        outputs=[result_image, result_info]
-    )
+    try:
+        start_time = time.time()
+        logger.info("Processing image...")
+        
+        # Check if image needs to be optimized (resized)
+        optimized_image = optimize_image_size(input_image)
+        
+        # Identify sets
+        found_sets, annotated_image, status_message = identify_sets_from_image(optimized_image)
+        
+        process_time = time.time() - start_time
+        logger.info(f"Image processed in {process_time:.2f} seconds.")
+        
+        return annotated_image, status_message
+        
+    except Exception as e:
+        error_message = f"Error processing image: {str(e)}"
+        logger.error(error_message)
+        logger.error(traceback.format_exc())
+        return input_image, error_message
 
 # =============================================================================
-#                          LAUNCH THE APP
+# GRADIO INTERFACE
 # =============================================================================
-
-demo.launch()
+def create_gradio_interface():
+    """
+    Creates and returns the Gradio interface for the SET Game Detector.
+    """
+    # CSS for styling the Gradio interface
+    css = """
+    @import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;600;700&display=swap');
+    
+    .gradio-container {
+        font-family: 'Poppins', sans-serif;
+    }
+    .app-header {
+        text-align: center;
+        margin-bottom: 20px;
+        background: linear-gradient(90deg, rgba(124, 58, 237, 0.1) 0%, rgba(236, 72, 153, 0.1) 100%);
+        padding: 1rem;
+        border-radius: 12px;
+    }
+    .app-header h1 {
+        font-size: 2.5rem;
+        background: linear-gradient(90deg, #8B5CF6 0%, #7C3AED 50%, #EC4899 100%);
+        -webkit-background-clip: text;
+        background-clip: text;
+        -webkit-text-fill-color: transparent;
+        margin-bottom: 5px;
+    }
+    .app-header p {
+        font-size: 1.1rem;
+        opacity: 0.8;
+        margin-top: 0;
+    }
+    .footer {
+        text-align: center;
+        margin-top: 20px;
+        padding: 10px;
+        background: linear-gradient(90deg, rgba(124, 58, 237, 0.05) 0%, rgba(236, 72, 153, 0.05) 100%);
+        border-radius: 12px;
+    }
+    
+    /* Responsive design for mobile */
+    @media (max-width: 600px) {
+        .app-header h1 {
+            font-size: 1.8rem;
+        }
+        .app-header p {
+            font-size: 0.9rem;
+        }
+    }
+    
+    /* Custom styling for buttons */
+    #find-sets-btn {
+        background: linear-gradient(90deg, #7C3AED 0%, #EC4899 100%);
+        color: white !important;
+    }
+    #find-sets-btn:hover {
+        opacity: 0.9;
+    }
+    
+    /* Image containers */
+    .image-container {
+        border-radius: 12px;
+        overflow: hidden;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    }
+    
+    /* Status box styling */
+    #status-box {
+        font-weight: 500;
+        border-radius: 8px;
+    }
+    """
+    
+    # Create the Gradio interface
+    with gr.Blocks(css=css, title="SET Game Detector") as demo:
+        # Header
+        gr.HTML("""
+        <div class="app-header">
+            <h1>🎴 SET Game Detector</h1>
+            <p>Upload an image of a SET board to find all valid sets</p>
+        </div>
+        """)
+        
+        # Model status display
+        model_status = gr.Textbox(
+            label="Model Status",
+            value=get_model_status(),
+            interactive=False
+        )
+        load_models_btn = gr.Button("🔄 Load Models", visible=not _MODELS_LOADED)
+        
+        # Main layout
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    label="Upload SET Board Image",
+                    tool="upload", 
+                    type="numpy",
+                    elem_id="input-image",
+                    elem_classes="image-container"
+                )
+                
+                with gr.Row():
+                    process_btn = gr.Button(
+                        "🔎 Find Sets", 
+                        variant="primary",
+                        elem_id="find-sets-btn",
+                        interactive=_MODELS_LOADED
+                    )
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
+            
+            with gr.Column():
+                output_image = gr.Image(
+                    label="Detected Sets",
+                    elem_id="output-image",
+                    elem_classes="image-container",
+                    interactive=False
+                )
+                status = gr.Textbox(
+                    label="Status", 
+                    placeholder="Upload an image and click 'Find Sets'",
+                    elem_id="status-box",
+                    interactive=False
+                )
+        
+        # Example images section - Create an examples directory for deployment
+        examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
+        os.makedirs(examples_dir, exist_ok=True)