Updated app.py with new improvements

app.py

@@ -0,0 +1,291 @@
+import cv2
+import numpy as np
+import tensorflow as tf
+import torch
+from ultralytics import YOLO
+from PIL import Image
+import gradio as gr
+import traceback
+import pandas as pd
+from itertools import combinations
+from huggingface_hub import hf_hub_download
+
+# =============================================================================
+#                           MODEL LOADING
+# =============================================================================
+
+# Load YOLO Card Detection Model
+card_model_path = hf_hub_download("Oamitai/card-detection", "best.pt")
+card_detection_model = YOLO(card_model_path)
+card_detection_model.conf = 0.5
+
+# Load YOLO Shape Detection Model
+shape_model_path = hf_hub_download("Oamitai/shape-detection", "best.pt")
+shape_detection_model = YOLO(shape_model_path)
+shape_detection_model.conf = 0.5
+
+# Load Shape Classification Model (Keras)
+shape_classification_model = tf.keras.models.load_model(
+    hf_hub_download("Oamitai/shape-classification", "shape_model.keras")
+)
+
+# Load Fill Classification Model (Keras)
+fill_classification_model = tf.keras.models.load_model(
+    hf_hub_download("Oamitai/fill-classification", "fill_model.keras")
+)
+
+# =============================================================================
+#                  ORIENTATION CORRECTION FUNCTIONS
+# =============================================================================
+def check_and_rotate_input_image(board_image, card_detection_model):
+    """
+    Checks the orientation of the board image by detecting 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
+    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
+
+    if avg_height > avg_width:
+        rotated_image = cv2.rotate(board_image, cv2.ROTATE_90_CLOCKWISE)
+        return rotated_image, True
+    else:
+        return board_image, False
+
+def restore_original_orientation(image, was_rotated):
+    """
+    If the image was rotated for processing, rotate it back to the original orientation.
+    """
+    if was_rotated:
+        return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    return image
+
+# =============================================================================
+#                        PREDICTION FUNCTIONS
+# =============================================================================
+def predict_color(shape_image):
+    """
+    Predict 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),
+        'red': cv2.countNonZero(red_mask)
+    }
+    return max(color_counts, key=color_counts.get)
+
+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.
+    """
+    shape_results = shape_detection_model(card_image)
+    card_height, card_width = card_image.shape[:2]
+    card_area = card_width * card_height
+
+    # Filter detections that are too small (less than 3% of card area)
+    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(detected_box)
+    
+    count = min(len(filtered_boxes), 3)
+    color_labels, fill_labels, shape_labels = [], [], []
+
+    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))
+    
+    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'
+    
+    return {
+        'count': count,
+        'color': color_label,
+        'fill': fill_label,
+        'shape': shape_label,
+        '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.
+    """
+    for feature in ['Count', 'Color', 'Fill', 'Shape']:
+        if len({card[feature] for card in cards}) not in [1, 3]:
+            return False
+    return True
+
+def find_sets(card_df):
+    """
+    Examine every combination of three cards from the DataFrame and return valid sets.
+    """
+    sets_found = []
+    for combo in combinations(card_df.iterrows(), 3):
+        cards = [entry[1] for entry in combo]
+        if is_set(cards):
+            set_info = {
+                'set_indices': [entry[0] for entry in combo],
+                'cards': [{feature: card[feature] for feature in ['Count', 'Color', 'Fill', 'Shape', 'Coordinates']} for card in cards]
+            }
+            sets_found.append(set_info)
+    return sets_found
+
+def detect_cards_from_image(board_image, card_detection_model):
+    """
+    Use 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)
+    card_boxes = card_results[0].boxes.xyxy.cpu().numpy().astype(int)
+    cards = []
+    for box in card_boxes:
+        x1, y1, x2, y2 = box
+        card_img = board_image[y1:y2, x1:x2]
+        cards.append((card_img, box))
+    return cards
+
+def classify_cards_from_board_image(board_image, card_detection_model, shape_detection_model, fill_model, shape_model):
+    """
+    For each detected card on the board image, predict 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)
+        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)
+
+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).
+    """
+    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)
+    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)
+    return sets_found, final_image
+
+# =============================================================================
+#                          DRAWING FUNCTIONS
+# =============================================================================
+def draw_sets_on_image(board_image, sets_info):
+    """
+    Draw 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)]
+    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_expanded = max(0, x1 - expansion)
+            y1_expanded = max(0, y1 - expansion)
+            x2_expanded = min(board_image.shape[1], x2 + expansion)
+            y2_expanded = min(board_image.shape[0], y2 + expansion)
+            cv2.rectangle(board_image, (x1_expanded, y1_expanded), (x2_expanded, y2_expanded), color, thickness)
+            if i == 0:
+                cv2.putText(board_image, f"Set {index + 1}", (x1_expanded, y1_expanded - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, thickness)
+    return board_image
+
+# =============================================================================
+#                        GRADIO INTERFACE FUNCTION
+# =============================================================================
+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.
+    """
+    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, 
+            card_detection_model, 
+            shape_detection_model, 
+            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
+
+# =============================================================================
+#                             LAUNCH GRADIO
+# =============================================================================
+iface = gr.Interface(
+    fn=detect_and_display_sets_interface,
+    inputs=gr.Image(type="pil", label="Upload Board Image"),
+    outputs=[gr.Image(type="pil", label="Annotated Image"), gr.Textbox(label="Status")],
+    title="Set Game Detector",
+    description=("Upload an image of a Set game board to detect cards, "
+                 "classify their features, and highlight valid sets.")
+)
+
+if __name__ == "__main__":
+    iface.launch()

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+tensorflow
+torch
+ultralytics
+opencv-python-headless
+numpy
+Pillow
+huggingface_hub
+pandas