ShapeContourApp.py

import streamlit as st
import cv2
import numpy as np
from math import cos, sin, pi

st.title("🔷 Contours & Shape Detection (Draw ➜ Detect)")


# -----------------------------
# Drawing utilities
# -----------------------------
def regular_polygon_points(center, radius, sides, rotation_deg=0):
    cx, cy = center
    rot = np.deg2rad(rotation_deg)
    angles = np.linspace(0, 2 * np.pi, sides, endpoint=False) + rot
    pts = np.column_stack(
        [cx + radius * np.cos(angles), cy + radius * np.sin(angles)]
    ).astype(np.int32)
    return pts.reshape((-1, 1, 2))


def star_points(center, r_outer, r_inner, rotation_deg=0, spikes=5):
    """Return 10 points (outer/inner alternating) for a star."""
    cx, cy = center
    rot = np.deg2rad(rotation_deg)
    pts = []
    for i in range(spikes * 2):
        r = r_outer if i % 2 == 0 else r_inner
        ang = rot + i * np.pi / spikes
        x = cx + r * np.cos(ang)
        y = cy + r * np.sin(ang)
        pts.append([x, y])
    pts = np.array(pts, dtype=np.int32).reshape((-1, 1, 2))
    return pts


def draw_shape_on_canvas(shape, canvas_size=400, rotation=0, **kwargs):
    """Returns a grayscale image with the chosen shape filled in white."""
    img = np.zeros((canvas_size, canvas_size), dtype=np.uint8)
    c = (canvas_size // 2, canvas_size // 2)

    if shape == "Circle":
        r = kwargs.get("radius", min(canvas_size // 3, 120))
        cv2.circle(img, c, r, 255, -1)

    elif shape == "Ellipse":
        ax = kwargs.get("axis_x", canvas_size // 4)
        ay = kwargs.get("axis_y", canvas_size // 6)
        cv2.ellipse(img, c, (ax, ay), rotation, 0, 360, 255, -1)

    elif shape == "Rectangle":
        w = kwargs.get("width", canvas_size // 3)
        h = kwargs.get("height", canvas_size // 5)
        rect = ((c[0], c[1]), (float(w), float(h)), float(rotation))
        box = cv2.boxPoints(rect).astype(np.int32).reshape((-1, 1, 2))
        cv2.fillPoly(img, [box], 255)

    elif shape == "Square":
        s = kwargs.get("side", canvas_size // 3)
        rect = ((c[0], c[1]), (float(s), float(s)), float(rotation))
        box = cv2.boxPoints(rect).astype(np.int32).reshape((-1, 1, 2))
        cv2.fillPoly(img, [box], 255)

    elif shape in {"Triangle", "Pentagon", "Hexagon"}:
        sides_map = {"Triangle": 3, "Pentagon": 5, "Hexagon": 6}
        sides = sides_map[shape]
        r = kwargs.get("radius", min(canvas_size // 3, 120))
        pts = regular_polygon_points(c, r, sides, rotation)
        cv2.fillPoly(img, [pts], 255)

    elif shape == "Regular Polygon (custom)":
        sides = kwargs.get("sides", 7)
        r = kwargs.get("radius", min(canvas_size // 3, 120))
        pts = regular_polygon_points(c, r, sides, rotation)
        cv2.fillPoly(img, [pts], 255)

    elif shape == "Star (5-point)":
        r_outer = kwargs.get("r_outer", min(canvas_size // 3, 120))
        r_inner = kwargs.get("r_inner", int(r_outer * 0.45))
        pts = star_points(c, r_outer, r_inner, rotation, spikes=5)
        cv2.fillPoly(img, [pts], 255)

    else:
        # fallback: draw nothing
        pass

    return img


# -----------------------------
# Detection utilities
# -----------------------------
def classify_shape(
    cnt, approx, circularity_thresh_circle=0.82, circularity_thresh_ellipse=0.65
):
    """Robust-ish classification:
    - 3 → Triangle
    - 4 → Square vs Rectangle via minAreaRect aspect ratio
    - 5 → Pentagon
    - 6 → Hexagon
    - >6 → Circle / Ellipse / Polygon(n) using circularity + convexity
    """
    area = cv2.contourArea(cnt)
    if area <= 0:
        return "Unknown"

    peri = cv2.arcLength(cnt, True)
    corners = len(approx)
    circularity = float(4 * np.pi * area / (peri * peri + 1e-6))

    if corners == 3:
        return "Triangle"
    elif corners == 4:
        rect = cv2.minAreaRect(cnt)
        (w, h) = rect[1]
        if min(w, h) == 0:
            return "Quadrilateral"
        ratio = min(w, h) / max(w, h)
        return "Square" if ratio > 0.92 else "Rectangle"
    elif corners == 5:
        return "Pentagon"
    elif corners == 6:
        return "Hexagon"
    else:
        # Concavity check (e.g., stars)
        if not cv2.isContourConvex(approx):
            return f"Concave polygon ({corners})"

        # Use circularity to distinguish circle/ellipse vs generic polygon
        if circularity >= circularity_thresh_circle:
            return "Circle"
        elif circularity >= circularity_thresh_ellipse:
            return "Ellipse"
        else:
            return f"Polygon ({corners})"


def detect_shapes(binary_img, eps_percent=1.5, min_area=100):
    """Detect, approximate and label shapes on a grayscale binary image."""
    contours, _ = cv2.findContours(
        binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    annotated = cv2.cvtColor(binary_img, cv2.COLOR_GRAY2BGR)

    for cnt in contours:
        if cv2.contourArea(cnt) < min_area:
            continue

        epsilon = (eps_percent / 100.0) * cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, epsilon, True)
        x, y, w, h = cv2.boundingRect(approx)

        shape_name = classify_shape(cnt, approx)

        cv2.drawContours(annotated, [approx], 0, (0, 255, 0), 2)
        cv2.putText(
            annotated,
            shape_name,
            (x, max(15, y - 8)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (255, 0, 0),
            2,
        )

    return annotated


# -----------------------------
# Sidebar controls
# -----------------------------
st.sidebar.header("🛠️ Draw Settings")

shape = st.sidebar.selectbox(
    "Shape to draw",
    [
        "Triangle",
        "Square",
        "Rectangle",
        "Pentagon",
        "Hexagon",
        "Circle",
        "Ellipse",
        "Star (5-point)",
        "Regular Polygon (custom)",
    ],
    index=5,
)

canvas_size = st.sidebar.slider("Canvas size (px)", 256, 800, 400, step=16)

# Per-shape parameters
rotation = 0
params = {}

if shape in {"Triangle", "Pentagon", "Hexagon", "Regular Polygon (custom)"}:
    params["radius"] = st.sidebar.slider(
        "Radius", 20, canvas_size // 2 - 10, min(120, canvas_size // 3)
    )
    rotation = st.sidebar.slider("Rotation (deg)", 0, 179, 0)
    if shape == "Regular Polygon (custom)":
        params["sides"] = st.sidebar.slider("Number of sides", 3, 14, 7)

elif shape == "Square":
    params["side"] = st.sidebar.slider(
        "Side length", 20, canvas_size - 40, min(160, canvas_size // 2)
    )
    rotation = st.sidebar.slider("Rotation (deg)", 0, 179, 0)

elif shape == "Rectangle":
    params["width"] = st.sidebar.slider(
        "Width", 20, canvas_size - 40, min(200, canvas_size // 2 + 40)
    )
    params["height"] = st.sidebar.slider(
        "Height", 20, canvas_size - 40, min(120, canvas_size // 3)
    )
    rotation = st.sidebar.slider("Rotation (deg)", 0, 179, 0)

elif shape == "Circle":
    params["radius"] = st.sidebar.slider(
        "Radius", 10, canvas_size // 2 - 10, min(120, canvas_size // 3)
    )

elif shape == "Ellipse":
    params["axis_x"] = st.sidebar.slider(
        "Axis X (semi-major)",
        10,
        canvas_size // 2 - 10,
        min(130, canvas_size // 3 + 10),
    )
    params["axis_y"] = st.sidebar.slider(
        "Axis Y (semi-minor)", 10, canvas_size // 2 - 10, min(80, canvas_size // 4)
    )
    rotation = st.sidebar.slider("Rotation (deg)", 0, 179, 30)

elif shape == "Star (5-point)":
    params["r_outer"] = st.sidebar.slider(
        "Outer radius", 20, canvas_size // 2 - 10, min(140, canvas_size // 3 + 20)
    )
    inner_ratio = st.sidebar.slider("Inner/Outer radius ratio", 20, 90, 45) / 100.0
    params["r_inner"] = int(params["r_outer"] * inner_ratio)
    rotation = st.sidebar.slider("Rotation (deg)", 0, 179, 0)

# Detection parameters
st.sidebar.header("🔎 Detection Settings")
eps_percent = st.sidebar.slider(
    "approxPolyDP epsilon (% of perimeter)", 0.5, 5.0, 1.5, 0.1
)
min_area = st.sidebar.slider("Min contour area (px²)", 20, 2000, 100, 10)

# -----------------------------
# Generate image and detect
# -----------------------------
img = draw_shape_on_canvas(shape, canvas_size=canvas_size, rotation=rotation, **params)

st.subheader("Processed Result")
# Ensure clean binary
_, thresh = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
result = detect_shapes(thresh, eps_percent=eps_percent, min_area=min_area)

col1, col2 = st.columns(2)
with col1:
    st.image(img, caption="Drawn Shape", use_container_width=True, channels="GRAY")
with col2:
    st.image(
        result,
        caption="Detected Contours & Labels",
        use_container_width=True,
        channels="BGR",
    )