Update app.py

app.py

@@ -4,299 +4,426 @@ import json
 import math
 import random
 import gradio as gr
+from sklearn.cluster import KMeans
 
 # Set random seed for reproducibility
 random.seed(42)
 np.random.seed(42)
 
-# --- Image Processing and Shape Detection ---
+# --- Utility Functions ---
 
-def detect_and_trace_shapes(img_bgr_np, canny_low_threshold, canny_high_threshold, 
-                            shape_detection_tolerance, min_contour_area, output_mode):
+def generate_id(prefix="el"):
+    """Generate a random ID for Excalidraw elements."""
+    return f"{prefix}_{''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=9))}"
+
+def bgr_to_hex(bgr_color):
+    """Convert BGR color to hex string."""
+    b, g, r = [int(c) for c in bgr_color[:3]]
+    return f"#{r:02x}{g:02x}{b:02x}"
+
+def simplify_contour_high_fidelity(contour, epsilon_factor=0.01):
     """
-    Detects shapes and lines of all colors in a BGR numpy image and converts them
-    into Excalidraw elements.
+    Simplifies a contour while preserving important features.
+    
+    Args:
+        contour: OpenCV contour (numpy array of points)
+        epsilon_factor: Factor to multiply with arc length for approximation tolerance
+    
+    Returns:
+        List of [x, y] coordinate pairs
     """
-    if img_bgr_np is None:
+    if len(contour) < 2:
         return []
+    
+    # Calculate epsilon based on contour perimeter
+    epsilon = epsilon_factor * cv2.arcLength(contour, True)
+    
+    # Apply Douglas-Peucker algorithm for polygon approximation
+    simplified = cv2.approxPolyDP(contour, epsilon, True)
+    
+    # Convert from OpenCV format to simple [x, y] list
+    points = []
+    for point in simplified:
+        x, y = point[0]  # OpenCV contours have shape (n, 1, 2)
+        points.append([float(x), float(y)])
+    
+    return points
 
-    # Convert to grayscale for edge detection
-    gray = cv2.cvtColor(img_bgr_np, cv2.COLOR_BGR2GRAY)
+def extract_dominant_colors(image, n_colors=8):
+    """Extract dominant colors from the image using K-means clustering."""
+    # Reshape image to be a list of pixels
+    data = image.reshape((-1, 3))
+    
+    # Remove black pixels (background) for better color detection
+    non_black_mask = np.sum(data, axis=1) > 30  # Threshold to exclude near-black pixels
+    if np.sum(non_black_mask) > 0:
+        data = data[non_black_mask]
     
-    # Apply Gaussian blur to reduce noise and help with edge detection
-    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    if len(data) == 0:
+        return [(0, 0, 0)]  # Return black if no colors found
+    
+    # Apply K-means clustering
+    kmeans = KMeans(n_clusters=min(n_colors, len(data)), random_state=42, n_init=10)
+    kmeans.fit(data)
+    
+    # Get the colors
+    colors = kmeans.cluster_centers_.astype(int)
+    return [tuple(color) for color in colors]
 
-    # Canny Edge Detection
-    edges = cv2.Canny(blurred, canny_low_threshold, canny_high_threshold)
+def create_color_mask(image, target_color, tolerance=50):
+    """Create a mask for pixels close to the target color."""
+    # Convert target color to numpy array
+    target = np.array(target_color, dtype=np.uint8)
+    
+    # Calculate color distance
+    diff = np.abs(image.astype(int) - target.astype(int))
+    distance = np.sqrt(np.sum(diff**2, axis=2))
+    
+    # Create mask where distance is less than tolerance
+    mask = distance < tolerance
+    return mask.astype(np.uint8) * 255
 
-    # Find contours from the edges
-    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+def detect_shapes_by_color(img_bgr_np, canny_low, canny_high, min_area, 
+                          shape_tolerance, output_mode, color_tolerance):
+    """
+    Detect shapes by processing each dominant color separately.
+    """
+    if img_bgr_np is None:
+        return []
     
     elements = []
-
-    for contour in contours:
-        # Filter out small contours (noise)
-        if cv2.contourArea(contour) < min_contour_area:
+    
+    # Extract dominant colors from the image
+    dominant_colors = extract_dominant_colors(img_bgr_np, n_colors=6)
+    
+    # Process each dominant color
+    for color_bgr in dominant_colors:
+        # Skip very dark colors (likely background)
+        if sum(color_bgr) < 50:
             continue
-
-        # Get bounding box to sample color
-        x, y, w, h = cv2.boundingRect(contour)
-        
-        # Sample average color from the original image within the bounding box
-        # Ensure coordinates are within image bounds
-        x1, y1 = max(0, x), max(0, y)
-        x2, y2 = min(img_bgr_np.shape[1], x + w), min(img_bgr_np.shape[0], y + h)
+            
+        # Create mask for this color
+        color_mask = create_color_mask(img_bgr_np, color_bgr, color_tolerance)
         
-        if x2 <= x1 or y2 <= y1: # Handle empty or invalid bounding boxes
+        # Skip if mask is too small
+        if np.sum(color_mask) < min_area:
             continue
-
-        color_region = img_bgr_np[y1:y2, x1:x2]
-        if color_region.size == 0:
-            avg_color_bgr = [0, 0, 0] # Default to black if region is empty
-        else:
-            avg_color_bgr = cv2.mean(color_region)[:3] # Get BGR values
         
-        # Convert BGR to hex color string
-        stroke_color = f"#{int(avg_color_bgr[2]):02x}{int(avg_color_bgr[1]):02x}{int(avg_color_bgr[0]):02x}"
+        # Apply morphological operations to clean up the mask
+        kernel = np.ones((3,3), np.uint8)
+        color_mask = cv2.morphologyEx(color_mask, cv2.MORPH_CLOSE, kernel)
+        color_mask = cv2.morphologyEx(color_mask, cv2.MORPH_OPEN, kernel)
+        
+        # Find contours in this color mask
+        contours, _ = cv2.findContours(color_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Convert color to hex for Excalidraw
+        stroke_color = bgr_to_hex(color_bgr)
         
-        # Default style properties (from user's example)
+        # Default style properties
         default_style = {
             "fillStyle": "cross-hatch",
-            "strokeWidth": 1,
+            "strokeWidth": 2,
             "strokeStyle": "solid",
             "roughness": 1,
             "opacity": 100,
-            "strokeColor": stroke_color, # Use detected color
-            "backgroundColor": "transparent", # Default to transparent for shapes unless filled
-            "roundness": {"type": 3}, # Default for rectangles
+            "strokeColor": stroke_color,
+            "backgroundColor": "transparent",
+            "roundness": {"type": 3},
             "seed": random.randint(1000, 1000000),
             "version": 1,
             "versionNonce": random.randint(1000, 1000000)
         }
-
-        if output_mode == "Plain Lines":
-            # Just create a line element from the simplified contour
-            simplified_points = simplify_contour_high_fidelity(contour)
-            if len(simplified_points) < 2:
+        
+        # Process each contour
+        for contour in contours:
+            if cv2.contourArea(contour) < min_area:
                 continue
-
-            # Excalidraw line points are relative to its x,y
-            start_x, start_y = simplified_points[0][0], simplified_points[0][1]
-            relative_points = [[p[0] - start_x, p[1] - start_y] for p in simplified_points]
-
-            line_element = {
-                "id": generate_id("line"), "type": "line",
-                "x": float(start_x), "y": float(start_y), "angle": 0,
-                "points": relative_points,
-                **default_style,
-                "backgroundColor": "transparent", # Lines don't have background
-                "roundness": {"type": 2} # Lines often use type 2
-            }
-            elements.append(line_element)
-
-        elif output_mode == "Geometric Shapes":
-            # Try to approximate shapes
-            epsilon = shape_detection_tolerance * cv2.arcLength(contour, True)
-            approx_polygon = cv2.approxPolyDP(contour, epsilon, True)
             
-            element_added = False
-
-            # Check for rectangle
-            if len(approx_polygon) == 4:
-                # Calculate angles to verify it's roughly a rectangle
-                # (This is a simplified check, more robust checks involve dot products of vectors)
-                is_rectangle = True
-                for i in range(4):
-                    p1 = approx_polygon[i][0]
-                    p2 = approx_polygon[(i + 1) % 4][0]
-                    p3 = approx_polygon[(i + 2) % 4][0]
-                    
-                    v1 = np.array(p2) - np.array(p1)
-                    v2 = np.array(p3) - np.array(p2)
-                    
-                    dot_product = np.dot(v1, v2)
-                    len_v1 = np.linalg.norm(v1)
-                    len_v2 = np.linalg.norm(v2)
-                    
-                    if len_v1 == 0 or len_v2 == 0:
-                        is_rectangle = False
-                        break
-
-                    angle_rad = np.arccos(np.clip(dot_product / (len_v1 * len_v2), -1.0, 1.0))
-                    angle_deg = math.degrees(angle_rad)
-                    
-                    # Check if angle is roughly 90 degrees (with tolerance)
-                    if not (80 <= angle_deg <= 100):
-                        is_rectangle = False
-                        break
-                
-                if is_rectangle:
-                    # Get rotated bounding box for better fit
-                    rect = cv2.minAreaRect(contour)
-                    box = cv2.boxPoints(rect)
-                    box = box.astype(int) # Corrected: Use .astype(int) instead of np.int0
-                    
-                    # Calculate width, height and angle
-                    width = np.linalg.norm(box[0] - box[1])
-                    height = np.linalg.norm(box[1] - box[2])
-                    angle = rect[2] # angle in degrees
-                    
-                    # Excalidraw angles are in radians, counter-clockwise from positive x-axis
-                    # OpenCV's minAreaRect angle is usually in range [-90, 0)
-                    # Adjust angle for Excalidraw:
-                    if width < height: # If height is the longer side (portrait)
-                        angle += 90
-                        width, height = height, width # Swap to make width the longer side
-                    
-                    angle_rad = math.radians(angle)
-                    
-                    # Center of the rectangle
-                    center_x, center_y = rect[0]
-                    
-                    rect_element = {
-                        "id": generate_id("rect"), "type": "rectangle",
-                        "x": float(center_x - width/2), "y": float(center_y - height/2),
-                        "width": float(width), "height": float(height),
-                        "angle": angle_rad,
-                        **default_style,
-                        "backgroundColor": stroke_color, # Use detected color for background
-                        "fillStyle": "solid" # Default to solid for detected shapes
-                    }
-                    elements.append(rect_element)
-                    element_added = True
-
-            # Check for ellipse/circle if not a rectangle
-            if not element_added and len(contour) >= 5: # min points for ellipse fitting
-                (x_center, y_center), (minor_axis, major_axis), angle = cv2.fitEllipse(contour)
-                
-                # Simple check for circularity (aspect ratio close to 1)
-                aspect_ratio = minor_axis / major_axis if major_axis > 0 else 0
-                if 0.8 <= aspect_ratio <= 1.2: # Treat as ellipse/circle
-                    ellipse_element = {
-                        "id": generate_id("ellipse"), "type": "ellipse",
-                        "x": float(x_center - major_axis/2), "y": float(y_center - major_axis/2),
-                        "width": float(major_axis), "height": float(major_axis), # Use major axis for both for circle
-                        "angle": math.radians(angle),
-                        **default_style,
-                        "backgroundColor": stroke_color, # Use detected color for background
-                        "fillStyle": "solid", # Default to solid for detected shapes
-                        "roundness": {"type": 2} # Ellipses often use type 2
-                    }
-                    elements.append(ellipse_element)
-                    element_added = True
-            
-            # Fallback to line if no specific shape detected
-            if not element_added:
-                simplified_points = simplify_contour_high_fidelity(contour)
+            if output_mode == "Plain Lines":
+                # Create line elements
+                simplified_points = simplify_contour_high_fidelity(contour, epsilon_factor=0.005)
                 if len(simplified_points) < 2:
                     continue
+                
+                # Create line element
                 start_x, start_y = simplified_points[0][0], simplified_points[0][1]
                 relative_points = [[p[0] - start_x, p[1] - start_y] for p in simplified_points]
-
+                
                 line_element = {
-                    "id": generate_id("line"), "type": "line",
-                    "x": float(start_x), "y": float(start_y), "angle": 0,
+                    "id": generate_id("line"),
+                    "type": "line",
+                    "x": float(start_x),
+                    "y": float(start_y),
+                    "angle": 0,
                     "points": relative_points,
                     **default_style,
                     "backgroundColor": "transparent",
                     "roundness": {"type": 2}
                 }
                 elements.append(line_element)
-
-    return elements
-
-def simplify_contour_high_fidelity(contour, epsilon_factor=0.01):
-    """
-    Simplifies a contour while preserving important features.
-    
-    Args:
-        contour: OpenCV contour (numpy array of points)
-        epsilon_factor: Factor to multiply with arc length for approximation tolerance
-                       Lower values = more points preserved (higher fidelity)
-                       Higher values = fewer points (more simplified)
-    
-    Returns:
-        List of [x, y] coordinate pairs
-    """
-    if len(contour) < 2:
-        return []
-    
-    # Calculate epsilon based on contour perimeter
-    epsilon = epsilon_factor * cv2.arcLength(contour, True)
-    
-    # Apply Douglas-Peucker algorithm for polygon approximation
-    simplified = cv2.approxPolyDP(contour, epsilon, True)
-    
-    # Convert from OpenCV format to simple [x, y] list
-    points = []
-    for point in simplified:
-        x, y = point[0]  # OpenCV contours have shape (n, 1, 2)
-        points.append([float(x), float(y)])
+                
+            elif output_mode == "Geometric Shapes":
+                # Try to detect geometric shapes
+                epsilon = shape_tolerance * cv2.arcLength(contour, True)
+                approx_polygon = cv2.approxPolyDP(contour, epsilon, True)
+                
+                element_added = False
+                
+                # Check for rectangle (4 corners)
+                if len(approx_polygon) == 4:
+                    # Verify it's roughly rectangular
+                    is_rectangle = True
+                    angles = []
+                    
+                    for i in range(4):
+                        p1 = approx_polygon[i][0]
+                        p2 = approx_polygon[(i + 1) % 4][0]
+                        p3 = approx_polygon[(i + 2) % 4][0]
+                        
+                        v1 = np.array(p2) - np.array(p1)
+                        v2 = np.array(p3) - np.array(p2)
+                        
+                        # Calculate angle between vectors
+                        if np.linalg.norm(v1) > 0 and np.linalg.norm(v2) > 0:
+                            cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
+                            cos_angle = np.clip(cos_angle, -1.0, 1.0)
+                            angle = math.degrees(math.acos(cos_angle))
+                            angles.append(angle)
+                    
+                    # Check if angles are close to 90 degrees
+                    if len(angles) == 4 and all(75 <= angle <= 105 for angle in angles):
+                        # Create rectangle
+                        rect = cv2.minAreaRect(contour)
+                        box = cv2.boxPoints(rect)
+                        
+                        width = np.linalg.norm(box[0] - box[1])
+                        height = np.linalg.norm(box[1] - box[2])
+                        angle = rect[2]
+                        
+                        if width < height:
+                            angle += 90
+                            width, height = height, width
+                        
+                        center_x, center_y = rect[0]
+                        
+                        rect_element = {
+                            "id": generate_id("rect"),
+                            "type": "rectangle",
+                            "x": float(center_x - width/2),
+                            "y": float(center_y - height/2),
+                            "width": float(width),
+                            "height": float(height),
+                            "angle": math.radians(angle),
+                            **default_style,
+                            "backgroundColor": stroke_color,
+                            "fillStyle": "solid"
+                        }
+                        elements.append(rect_element)
+                        element_added = True
+                
+                # Check for circle/ellipse
+                if not element_added and len(contour) >= 5:
+                    try:
+                        (x_center, y_center), (minor_axis, major_axis), angle = cv2.fitEllipse(contour)
+                        
+                        # Check if it's roughly circular
+                        if major_axis > 0:
+                            aspect_ratio = minor_axis / major_axis
+                            if 0.7 <= aspect_ratio <= 1.3:  # Allow some tolerance for circles
+                                avg_radius = (minor_axis + major_axis) / 2
+                                
+                                ellipse_element = {
+                                    "id": generate_id("ellipse"),
+                                    "type": "ellipse",
+                                    "x": float(x_center - avg_radius/2),
+                                    "y": float(y_center - avg_radius/2),
+                                    "width": float(avg_radius),
+                                    "height": float(avg_radius),
+                                    "angle": math.radians(angle),
+                                    **default_style,
+                                    "backgroundColor": stroke_color,
+                                    "fillStyle": "solid",
+                                    "roundness": {"type": 2}
+                                }
+                                elements.append(ellipse_element)
+                                element_added = True
+                    except:
+                        pass  # Skip if ellipse fitting fails
+                
+                # Fallback to line if no shape detected
+                if not element_added:
+                    simplified_points = simplify_contour_high_fidelity(contour, epsilon_factor=0.005)
+                    if len(simplified_points) >= 2:
+                        start_x, start_y = simplified_points[0][0], simplified_points[0][1]
+                        relative_points = [[p[0] - start_x, p[1] - start_y] for p in simplified_points]
+                        
+                        line_element = {
+                            "id": generate_id("line"),
+                            "type": "line",
+                            "x": float(start_x),
+                            "y": float(start_y),
+                            "angle": 0,
+                            "points": relative_points,
+                            **default_style,
+                            "backgroundColor": "transparent",
+                            "roundness": {"type": 2}
+                        }
+                        elements.append(line_element)
     
-    return points
-
-def generate_id(prefix="el"):
-    """Generate a random ID for Excalidraw elements."""
-    return f"{prefix}_{''.join(random.choices('abcdefghijklmnopqrstuvwxyz013456789', k=9))}"
-
-# --- Gradio UI Function ---
+    return elements
 
-def generate_excalidraw_json_from_image_new(
-    image_np, 
-    canny_low_threshold, 
-    canny_high_threshold, 
-    shape_detection_tolerance, 
-    min_contour_area, 
-    output_mode
+def generate_excalidraw_json_from_image(
+    image_np,
+    canny_low_threshold,
+    canny_high_threshold,
+    shape_detection_tolerance,
+    min_contour_area,
+    output_mode,
+    color_tolerance
 ):
     """
-    Main function to generate Excalidraw JSON from an uploaded image with new controls.
+    Main function to generate Excalidraw JSON from an uploaded image.
     """
     if image_np is None:
         return "Please upload an image to generate Excalidraw elements."
-
-    # Gradio's image input is typically RGB, OpenCV expects BGR
+    
+    # Convert RGB to BGR for OpenCV
     img_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
-
-    final_elements = detect_and_trace_shapes(
-        img_bgr, 
-        canny_low_threshold, 
-        canny_high_threshold, 
-        shape_detection_tolerance, 
-        min_contour_area, 
-        output_mode
+    
+    # Detect shapes by color
+    final_elements = detect_shapes_by_color(
+        img_bgr,
+        canny_low_threshold,
+        canny_high_threshold,
+        min_contour_area,
+        shape_detection_tolerance,
+        output_mode,
+        color_tolerance
     )
-
+    
+    # Create final Excalidraw structure
     final_excalidraw_structure = {
         "type": "excalidraw/clipboard",
         "elements": final_elements,
         "files": {}
     }
+    
     return json.dumps(final_excalidraw_structure, indent=2)
 
-# Gradio interface definition
-iface = gr.Interface(
-    fn=generate_excalidraw_json_from_image_new,
-    inputs=[
-        gr.Image(type="numpy", label="Upload Image"),
-        gr.Slider(minimum=0, maximum=255, value=50, label="Canny Low Threshold", info="Lower values detect more edges."),
-        gr.Slider(minimum=0, maximum=255, value=150, label="Canny High Threshold", info="Higher values require stronger edges."),
-        gr.Slider(minimum=0.001, maximum=0.1, value=0.04, label="Shape Detection Tolerance", info="Lower values for more precise polygon approximation. Affects rectangle/ellipse detection."),
-        gr.Slider(minimum=10, maximum=1000, value=100, label="Minimum Contour Area", info="Filter out small noise contours."),
-        gr.Dropdown(
-            label="Output Mode",
-            choices=["Plain Lines", "Geometric Shapes"],
-            value="Geometric Shapes",
-            info="Choose to output raw lines or try to detect geometric shapes."
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="Excalidraw Color Shape Detector", theme=gr.themes.Soft()) as iface:
+        gr.Markdown("# 🎨 Excalidraw Color Shape Detector")
+        gr.Markdown("Upload an image with colorful line art, and this tool will detect shapes by color and convert them to Excalidraw elements.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(type="numpy", label="Upload Image")
+                
+                gr.Markdown("### Detection Settings")
+                
+                output_mode = gr.Dropdown(
+                    label="Output Mode",
+                    choices=["Plain Lines", "Geometric Shapes"],
+                    value="Geometric Shapes",
+                    info="Choose between raw lines or geometric shape detection"
+                )
+                
+                color_tolerance = gr.Slider(
+                    minimum=10,
+                    maximum=100,
+                    value=40,
+                    label="Color Tolerance",
+                    info="How similar colors are grouped together (lower = more precise)"
+                )
+                
+                min_contour_area = gr.Slider(
+                    minimum=50,
+                    maximum=2000,
+                    value=200,
+                    label="Minimum Shape Area",
+                    info="Filter out small noise (higher = fewer small shapes)"
+                )
+                
+                shape_detection_tolerance = gr.Slider(
+                    minimum=0.005,
+                    maximum=0.05,
+                    value=0.02,
+                    label="Shape Detection Tolerance",
+                    info="How precise shape detection should be (lower = more precise)"
+                )
+                
+                gr.Markdown("### Advanced Settings")
+                
+                canny_low_threshold = gr.Slider(
+                    minimum=10,
+                    maximum=100,
+                    value=30,
+                    label="Edge Detection - Low Threshold"
+                )
+                
+                canny_high_threshold = gr.Slider(
+                    minimum=50,
+                    maximum=200,
+                    value=100,
+                    label="Edge Detection - High Threshold"
+                )
+                
+                process_btn = gr.Button("🔄 Process Image", variant="primary")
+            
+            with gr.Column(scale=2):
+                output_json = gr.Textbox(
+                    label="Excalidraw JSON Output",
+                    info="Copy this JSON and paste it into Excalidraw (Ctrl+V)",
+                    lines=20,
+                    max_lines=30
+                )
+                
+                gr.Markdown("### Instructions:")
+                gr.Markdown("""
+                1. Upload an image with colorful shapes or line art
+                2. Adjust the settings if needed:
+                   - **Color Tolerance**: Lower values for more precise color matching
+                   - **Minimum Shape Area**: Higher values to filter out noise
+                   - **Shape Detection**: Lower values for more precise geometric shapes
+                3. Click "Process Image"
+                4. Copy the JSON output and paste it into Excalidraw
+                """)
+        
+        # Set up the processing function
+        process_btn.click(
+            fn=generate_excalidraw_json_from_image,
+            inputs=[
+                image_input,
+                canny_low_threshold,
+                canny_high_threshold,
+                shape_detection_tolerance,
+                min_contour_area,
+                output_mode,
+                color_tolerance
+            ],
+            outputs=output_json
         )
-    ], 
-    outputs=gr.Textbox(label="Excalidraw JSON Output (Copy & Paste into Excalidraw)"),
-    title="Excalidraw Universal Shape Detector & Tracer",
-    description="Upload an image. This tool will detect shapes and lines of all colors, "
-                "and convert them into Excalidraw elements. You can control the detection "
-                "sensitivity and choose between raw lines or approximated geometric shapes."
-)
+        
+        # Example processing on image upload
+        image_input.change(
+            fn=generate_excalidraw_json_from_image,
+            inputs=[
+                image_input,
+                canny_low_threshold,
+                canny_high_threshold,
+                shape_detection_tolerance,
+                min_contour_area,
+                output_mode,
+                color_tolerance
+            ],
+            outputs=output_json
+        )
+    
+    return iface
 
+# Launch the app
 if __name__ == "__main__":
-    iface.launch()
-
+    app = create_interface()
+    app.launch()