Updated app.py with new improvements

app.py

@@ -11,10 +11,10 @@ from itertools import combinations
 from huggingface_hub import hf_hub_download
 import spaces  # For ZeroGPU support
 
-# -----------------------------------------------------------------------------
-# Global variables for models (loaded lazily inside the GPU function)
-# Keras models can be loaded globally if they don't trigger CUDA initialization.
-# We'll load the Keras models here since they typically run on CPU and use GPU later.
+# =============================================================================
+#                           MODEL LOADING (Keras Models on CPU)
+# =============================================================================
+# These models can be loaded globally.
 shape_classification_model = tf.keras.models.load_model(
     hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
 )
@@ -22,7 +22,7 @@ fill_classification_model = tf.keras.models.load_model(
     hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
 )
 
-# Global YOLO models will be loaded inside the GPU function.
+# Global YOLO models will be loaded lazily inside the GPU function.
 global_card_detection_model = None
 global_shape_detection_model = None
 
@@ -55,10 +55,10 @@ def restore_original_orientation(image, was_rotated):
 # =============================================================================
 def predict_color(shape_image):
     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]))
+    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),
@@ -71,14 +71,17 @@ def predict_card_features(card_image, shape_detection_model, fill_model, shape_m
     shape_results = shape_detection_model(card_image)
     card_height, card_width = card_image.shape[:2]
     card_area = card_width * card_height
+
     filtered_boxes = []
     for detected_box in shape_results[0].boxes.xyxy.cpu().numpy():
         x1, y1, x2, y2 = detected_box.astype(int)
         shape_area = (x2 - x1) * (y2 - y1)
         if shape_area > 0.03 * card_area:
             filtered_boxes.append([x1, y1, x2, y2])
+    
     if len(filtered_boxes) == 0:
         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 = []
@@ -94,16 +97,19 @@ def predict_card_features(card_image, shape_detection_model, fill_model, shape_m
         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]
+    
     count = min(len(filtered_boxes), 3)
     final_color = max(set(color_list), key=color_list.count)
     final_fill = max(set(predicted_fill), key=predicted_fill.count)
     final_shape = max(set(predicted_shape), key=predicted_shape.count)
+    
     return {'count': count, 'color': final_color, 'fill': final_fill, 'shape': final_shape, 'box': box}
 
 def is_set(cards):
@@ -194,7 +200,7 @@ def detect_and_display_sets_interface(input_image):
     try:
         image_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
         global global_card_detection_model, global_shape_detection_model
-        # Load YOLO models on GPU only once after allocation.
+        # Lazy load YOLO models on GPU after allocation.
         if global_card_detection_model is None:
             card_model_path = hf_hub_download("Oamitai/card-detection", "best.pt")
             global_card_detection_model = YOLO(card_model_path)