Update app.py

app.py

@@ -8,9 +8,7 @@ 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 time
 from typing import Tuple, List, Dict
@@ -34,78 +32,50 @@ logging.basicConfig(level=logging.INFO,
 logger = logging.getLogger("set_detector")
 
 # =============================================================================
-# MODEL PATHS
+# MODEL LOADING
 # =============================================================================
-# For loading models from Hugging Face Hub
-HF_MODEL_REPO_PREFIX = "Omamitai"
-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"
-
-CARD_MODEL_FILENAME = "best.pt"
-SHAPE_MODEL_FILENAME = "best.pt"
-SHAPE_CLASS_MODEL_FILENAME = "shape_model.keras"
-FILL_CLASS_MODEL_FILENAME = "fill_model.keras"
-
 # Global variables for model caching
-_MODEL_SHAPE = None
-_MODEL_FILL = None
-_DETECTOR_CARD = None
-_DETECTOR_SHAPE = None
+_CARD_DETECTOR = None
+_SHAPE_DETECTOR = None
+_SHAPE_CLASSIFIER = None
+_FILL_CLASSIFIER = None
 
 def load_models():
     """
     Load all models needed for SET detection.
-    Returns tuple of (card_detector, shape_detector, shape_model, fill_model)
+    Returns tuple of (card_detector, shape_detector, shape_classifier, fill_classifier)
     """
-    global _MODEL_SHAPE, _MODEL_FILL, _DETECTOR_CARD, _DETECTOR_SHAPE
+    global _CARD_DETECTOR, _SHAPE_DETECTOR, _SHAPE_CLASSIFIER, _FILL_CLASSIFIER
     
     # Return cached models if already loaded
-    if all([_MODEL_SHAPE, _MODEL_FILL, _DETECTOR_CARD, _DETECTOR_SHAPE]):
-        return _DETECTOR_CARD, _DETECTOR_SHAPE, _MODEL_SHAPE, _MODEL_FILL
+    if all([_CARD_DETECTOR, _SHAPE_DETECTOR, _SHAPE_CLASSIFIER, _FILL_CLASSIFIER]):
+        logger.info("Using cached models")
+        return _CARD_DETECTOR, _SHAPE_DETECTOR, _SHAPE_CLASSIFIER, _FILL_CLASSIFIER
     
     try:
         from huggingface_hub import hf_hub_download
         
-        # Download and load YOLO models
-        logger.info("Downloading detection models...")
-        card_model_path = hf_hub_download(
-            repo_id=CARD_DETECTION_REPO, 
-            filename=CARD_MODEL_FILENAME,
-            cache_dir="./hf_cache"
-        )
-        
-        shape_model_path = hf_hub_download(
-            repo_id=SHAPE_DETECTION_REPO, 
-            filename=SHAPE_MODEL_FILENAME,
-            cache_dir="./hf_cache"
-        )
-        
-        # Download and load classification models
-        logger.info("Downloading classification models...")
-        shape_class_path = hf_hub_download(
-            repo_id=SHAPE_CLASSIFICATION_REPO, 
-            filename=SHAPE_CLASS_MODEL_FILENAME,
-            cache_dir="./hf_cache"
-        )
-        
-        fill_class_path = hf_hub_download(
-            repo_id=FILL_CLASSIFICATION_REPO, 
-            filename=FILL_CLASS_MODEL_FILENAME,
-            cache_dir="./hf_cache"
-        )
+        logger.info("Loading models from Hugging Face Hub...")
         
-        # Load all models
-        logger.info("Loading models...")
-        card_detector = YOLO(str(card_model_path))
+        # 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
         
-        shape_detector = YOLO(str(shape_model_path))
+        # 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
         
-        shape_classifier = load_model(str(shape_class_path))
-        fill_classifier = load_model(str(fill_class_path))
+        # 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")
+        )
         
         # Use GPU if available
         if torch.cuda.is_available():
@@ -114,10 +84,10 @@ def load_models():
             shape_detector.to("cuda")
         
         # Cache the models
-        _DETECTOR_CARD = card_detector
-        _DETECTOR_SHAPE = shape_detector
-        _MODEL_SHAPE = shape_classifier
-        _MODEL_FILL = fill_classifier
+        _CARD_DETECTOR = card_detector
+        _SHAPE_DETECTOR = shape_detector
+        _SHAPE_CLASSIFIER = shape_classifier
+        _FILL_CLASSIFIER = fill_classifier
         
         logger.info("All models loaded successfully!")
         return card_detector, shape_detector, shape_classifier, fill_classifier
@@ -437,10 +407,6 @@ def process_image(input_image):
         logger.error(traceback.format_exc())
         return input_image, error_message
 
-# Create examples directory
-examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
-os.makedirs(examples_dir, exist_ok=True)
-
 # =============================================================================
 # GRADIO INTERFACE
 # =============================================================================
@@ -474,15 +440,6 @@ with gr.Blocks(title="SET Game Detector") as demo:
                 interactive=False
             )
     
-    # Examples - simplified
-    gr.Examples(
-        examples=[
-            os.path.join(examples_dir, "set_example1.jpg"),
-            os.path.join(examples_dir, "set_example2.jpg")
-        ],
-        inputs=input_image
-    )
-    
     # Function bindings inside the Blocks context
     find_sets_btn.click(
         fn=process_image,