Updated app.py with new improvements

app.py

@@ -11,6 +11,27 @@ from itertools import combinations
 from huggingface_hub import hf_hub_download
 import spaces  # Added for ZeroGPU support
 
+# -----------------------------------------------------------------------------
+#   SAFE CUDA CHECK
+# -----------------------------------------------------------------------------
+def safe_cuda_available():
+    try:
+        return torch.cuda.is_available()
+    except Exception as e:
+        print("Error checking CUDA availability:", e)
+        return False
+
+gpu_available = safe_cuda_available()
+if gpu_available:
+    try:
+        # Move YOLO models to GPU later after loading models.
+        gpu_message = "GPU is available. Will move models to GPU."
+    except Exception as e:
+        print("Failed to set GPU mode:", e)
+        gpu_available = False
+else:
+    print("CUDA not available or error encountered. Running on CPU.")
+
 # =============================================================================
 #                           MODEL LOADING
 # =============================================================================
@@ -35,31 +56,34 @@ fill_classification_model = tf.keras.models.load_model(
     hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
 )
 
+# Move YOLO models to GPU if available
+if gpu_available:
+    try:
+        card_detection_model.to("cuda")
+        shape_detection_model.to("cuda")
+        print("YOLO models moved to GPU.")
+    except Exception as e:
+        print("Failed to move YOLO models to GPU:", e)
+
 # =============================================================================
 #                  ORIENTATION CORRECTION FUNCTIONS
 # =============================================================================
-def check_and_rotate_input_image(board_image, card_detection_model):
+def check_and_rotate_input_image(board_image, card_boxes):
     """
-    Checks the orientation of the board image by detecting card bounding boxes.
+    Checks the orientation of the board image based on detected card bounding boxes.
     If the average detected card height is greater than the average card width,
     rotates the image 90° clockwise.
     """
-    card_results = card_detection_model(board_image)
-    card_boxes = card_results[0].boxes.xyxy.cpu().numpy().astype(int)
-    
-    # If no cards are detected, assume no rotation is needed.
     if len(card_boxes) == 0:
         return board_image, False
 
-    total_width = total_height = count = 0
+    total_width = total_height = 0
     for box in card_boxes:
         x1, y1, x2, y2 = box
         total_width += (x2 - x1)
         total_height += (y2 - y1)
-        count += 1
-
-    avg_width = total_width / count
-    avg_height = total_height / count
+    avg_width = total_width / len(card_boxes)
+    avg_height = total_height / len(card_boxes)
 
     if avg_height > avg_width:
         rotated_image = cv2.rotate(board_image, cv2.ROTATE_90_CLOCKWISE)
@@ -69,7 +93,7 @@ def check_and_rotate_input_image(board_image, card_detection_model):
 
 def restore_original_orientation(image, was_rotated):
     """
-    If the image was rotated for processing, rotate it back to the original orientation.
+    Rotates the image back to the original orientation if it was rotated.
     """
     if was_rotated:
         return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
@@ -80,16 +104,14 @@ def restore_original_orientation(image, was_rotated):
 # =============================================================================
 def predict_color(shape_image):
     """
-    Predict the dominant color (green, purple, or red) using HSV thresholds.
+    Predicts the dominant color (green, purple, or red) using HSV thresholds.
     """
     hsv_image = cv2.cvtColor(shape_image, cv2.COLOR_BGR2HSV)
-    # Define HSV ranges
     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)
-    # Count non-zero pixels in each mask
     color_counts = {
         'green': cv2.countNonZero(green_mask),
         'purple': cv2.countNonZero(purple_mask),
@@ -99,8 +121,8 @@ def predict_color(shape_image):
 
 def predict_card_features(card_image, shape_detection_model, fill_model, shape_model, box):
     """
-    For a given card image, detect shapes and predict fill and shape attributes.
-    Returns a dictionary of features.
+    For a given card image, detects shapes and batches the cropped shape images for
+    classification. Returns a dictionary of features.
     """
     shape_results = shape_detection_model(card_image)
     card_height, card_width = card_image.shape[:2]
@@ -112,47 +134,47 @@ def predict_card_features(card_image, shape_detection_model, fill_model, shape_m
         x1, y1, x2, y2 = detected_box.astype(int)
         shape_area = (x2 - x1) * (y2 - y1)
         if shape_area > 0.03 * card_area:
-            filtered_boxes.append(detected_box)
+            filtered_boxes.append([x1, y1, x2, y2])
     
-    count = min(len(filtered_boxes), 3)
-    color_labels, fill_labels, shape_labels = [], [], []
+    if len(filtered_boxes) == 0:
+        return {'count': 0, 'color': 'unknown', 'fill': 'unknown', 'shape': 'unknown', 'box': box}
 
-    for shape_box in filtered_boxes:
-        shape_box = shape_box.astype(int)
-        shape_img = card_image[shape_box[1]:shape_box[3], shape_box[0]:shape_box[2]]
-        # Preprocess images for classification models
-        fill_input_shape = fill_model.input_shape[1:3]
-        shape_input_shape = shape_model.input_shape[1:3]
-        fill_img = cv2.resize(shape_img, fill_input_shape) / 255.0
-        shape_img_resized = cv2.resize(shape_img, shape_input_shape) / 255.0
-        fill_img = np.expand_dims(fill_img, axis=0)
-        shape_img_resized = np.expand_dims(shape_img_resized, axis=0)
-        # Make predictions
-        fill_pred = fill_model.predict(fill_img)
-        shape_pred = shape_model.predict(shape_img_resized)
-        fill_labels.append(['empty', 'full', 'striped'][np.argmax(fill_pred)])
-        shape_labels.append(['diamond', 'oval', 'squiggle'][np.argmax(shape_pred)])
-        color_labels.append(predict_color(shape_img))
+    # Precompute input shapes for classification models
+    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)
     
-    if count > 0:
-        color_label = max(set(color_labels), key=color_labels.count)
-        fill_label = max(set(fill_labels), key=fill_labels.count)
-        shape_label = max(set(shape_labels), key=shape_labels.count)
-    else:
-        color_label = fill_label = shape_label = 'unknown'
+    # Batch predictions (GPU or CPU will be used based on model configuration)
+    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]
     
-    return {
-        'count': count,
-        'color': color_label,
-        'fill': fill_label,
-        'shape': shape_label,
-        'box': box
-    }
+    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):
     """
-    Check if a combination of cards forms a valid set.
-    For each feature, values must be either all the same or all different.
+    Returns True if for each feature (Count, Color, Fill, Shape) the values are
+    all the same or all different; otherwise, returns False.
     """
     for feature in ['Count', 'Color', 'Fill', 'Shape']:
         if len({card[feature] for card in cards}) not in [1, 3]:
@@ -161,7 +183,7 @@ def is_set(cards):
 
 def find_sets(card_df):
     """
-    Examine every combination of three cards from the DataFrame and return valid sets.
+    Examines every combination of three cards from the DataFrame and returns valid sets.
     """
     sets_found = []
     for combo in combinations(card_df.iterrows(), 3):
@@ -176,7 +198,7 @@ def find_sets(card_df):
 
 def detect_cards_from_image(board_image, card_detection_model):
     """
-    Use the YOLO card detection model to detect cards on the board image.
+    Uses the YOLO card detection model to detect cards on the board image.
     Returns a list of tuples: (cropped card image, bounding box).
     """
     card_results = card_detection_model(board_image)
@@ -186,17 +208,18 @@ def detect_cards_from_image(board_image, card_detection_model):
         x1, y1, x2, y2 = box
         card_img = board_image[y1:y2, x1:x2]
         cards.append((card_img, box))
-    return cards
+    return cards, card_boxes
 
-def classify_cards_from_board_image(board_image, card_detection_model, shape_detection_model, fill_model, shape_model):
+def classify_cards_from_board_image(board_image, card_boxes, shape_detection_model, fill_model, shape_model):
     """
-    For each detected card on the board image, predict its features.
+    For each detected card (using precomputed bounding boxes), predicts its features.
     Returns a pandas DataFrame of card feature data.
     """
-    cards = detect_cards_from_image(board_image, card_detection_model)
     card_data = []
-    for card_image, box in cards:
-        features = predict_card_features(card_image, shape_detection_model, fill_model, shape_model, box)
+    for box in card_boxes:
+        x1, y1, x2, y2 = box
+        card_img = board_image[y1:y2, x1:x2]
+        features = predict_card_features(card_img, shape_detection_model, fill_model, shape_model, box)
         card_data.append({
             "Count": features['count'],
             "Color": features['color'],
@@ -208,12 +231,20 @@ def classify_cards_from_board_image(board_image, card_detection_model, shape_det
 
 def classify_and_find_sets_from_array(board_image, card_detection_model, shape_detection_model, fill_model, shape_model):
     """
-    Processes a board image (in BGR format), corrects its orientation, detects cards,
-    classifies them, finds valid sets, and finally restores the original orientation.
-    Returns a tuple: (sets_found, annotated image).
+    Processes a board image by:
+      1. Detecting cards and checking/correcting orientation.
+      2. Re-detecting cards if the image was rotated.
+      3. Classifying each card.
+      4. Finding valid sets and annotating the image.
+    Returns a tuple: (sets_found, final_image).
     """
-    board_image, was_rotated = check_and_rotate_input_image(board_image, card_detection_model)
-    card_df = classify_cards_from_board_image(board_image, card_detection_model, shape_detection_model, fill_model, shape_model)
+    # First detection for orientation check
+    _, card_boxes = detect_cards_from_image(board_image, card_detection_model)
+    board_image, was_rotated = check_and_rotate_input_image(board_image, card_boxes)
+    # Re-run detection if image was rotated to get correct bounding boxes.
+    if was_rotated:
+        _, card_boxes = detect_cards_from_image(board_image, card_detection_model)
+    card_df = classify_cards_from_board_image(board_image, card_boxes, shape_detection_model, fill_model, shape_model)
     sets_found = find_sets(card_df)
     annotated_image = draw_sets_on_image(board_image.copy(), sets_found)
     final_image = restore_original_orientation(annotated_image, was_rotated)
@@ -224,7 +255,7 @@ def classify_and_find_sets_from_array(board_image, card_detection_model, shape_d
 # =============================================================================
 def draw_sets_on_image(board_image, sets_info):
     """
-    Draw bounding boxes and set labels on the board image for each detected set.
+    Draws bounding boxes and set labels on the board image for each detected set.
     """
     colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255),
               (255, 255, 0), (255, 0, 255), (0, 255, 255)]
@@ -248,19 +279,18 @@ def draw_sets_on_image(board_image, sets_info):
     return board_image
 
 # =============================================================================
-#                        GRADIO INTERFACE FUNCTION
+#                     GRADIO INTERFACE FUNCTION (ZeroGPU)
 # =============================================================================
-@spaces.GPU  # Decorator added for ZeroGPU dynamic GPU allocation
+@spaces.GPU
 def detect_and_display_sets_interface(input_image):
     """
-    Gradio interface function:
-      - Accepts a PIL image (board image)
-      - Converts it to a cv2 BGR image
-      - Processes it for set detection
-      - Returns the annotated image (as PIL) and a status message.
+    Gradio interface function that:
+      - Accepts a PIL image.
+      - Converts it to OpenCV (BGR) format.
+      - Processes it for set detection.
+      - Returns the annotated image (as a PIL image) along with a status message.
     """
     try:
-        # Convert input PIL image to OpenCV (BGR) format.
         image_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
         sets_found, final_image = classify_and_find_sets_from_array(
             image_cv, 
@@ -269,12 +299,10 @@ def detect_and_display_sets_interface(input_image):
             fill_classification_model, 
             shape_classification_model
         )
-        # Convert back to RGB for display.
         final_image_rgb = cv2.cvtColor(final_image, cv2.COLOR_BGR2RGB)
         return Image.fromarray(final_image_rgb), "Sets detected successfully."
     except Exception as e:
         err = f"❌ Error: {str(e)}\n{traceback.format_exc()}"
-        # Return a blank image with error details.
         return Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8)), err
 
 # =============================================================================