app1.py

import base64
import io
import cv2
import os
import numpy as np
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

# -------------------
# FastAPI app
# -------------------
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# -------------------
# Health check
# -------------------
@app.get("/")
def home():
    return {"status": "running"}


# -------------------------------
# Helper functions
# -------------------------------
def _largest_contour(mask, use_chain_approx_none=True):
    mode = cv2.CHAIN_APPROX_NONE if use_chain_approx_none else cv2.CHAIN_APPROX_SIMPLE
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, mode)
    if not contours:
        return None
    return max(contours, key=cv2.contourArea)


def _min_area_rect_to_poly(cnt):
    rect = cv2.minAreaRect(cnt)
    box = cv2.boxPoints(rect)
    return box.astype(np.float32).reshape(-1, 1, 2)


def _smooth_closed_poly(pts, k=7):
    """Simple circular moving-average smoothing for a closed contour."""
    if k < 3 or len(pts) < k:
        return pts.astype(np.float32)

    pts = pts.reshape(-1, 2).astype(np.float32)
    pad = k // 2
    pts_pad = np.vstack([pts[-pad:], pts, pts[:pad]])

    kernel = np.ones((k,), dtype=np.float32) / k
    xs = np.convolve(pts_pad[:, 0], kernel, mode="valid")
    ys = np.convolve(pts_pad[:, 1], kernel, mode="valid")

    smoothed = np.stack([xs, ys], axis=1)
    return smoothed.astype(np.float32)


def mask_to_polygon_no_holes(mask, epsilon_factor=0.01, min_area=100):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None

    contour = max(contours, key=cv2.contourArea)
    if cv2.contourArea(contour) < min_area:
        return None

    epsilon = epsilon_factor * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, epsilon, True)
    return approx


def clean_polygon_strict(mask, epsilon_factor=0.005, smooth_k=7,
                         rect_iou_thresh=0.88, min_area=100):
    """Return either a snapped rectangle or a clean simplified polygon, with smoothing."""
    if mask.dtype != np.uint8:
        mask = mask.astype(np.uint8)

    bw = (mask > 127).astype(np.uint8)
    cnt = _largest_contour(bw, use_chain_approx_none=True)
    if cnt is None or len(cnt) < 4:
        return None, "No contour"

    # optional smoothing
    if smooth_k and smooth_k >= 3:
        sm = _smooth_closed_poly(cnt, k=smooth_k)
        cnt_for_approx = sm.reshape(-1, 1, 2)
    else:
        cnt_for_approx = cnt.astype(np.float32)

    # Candidate A: rectangle
    rect_poly = _min_area_rect_to_poly(cnt_for_approx)

    # Candidate B: polygon
    polyB = mask_to_polygon_no_holes(bw, epsilon_factor=0.005, min_area=min_area)

    # Decision: rectangle vs polygon
    if rect_poly is not None and polyB is not None:
        rect_area = cv2.contourArea(rect_poly)
        contour_area = cv2.contourArea(cnt)
        area_ratio = rect_area / contour_area if contour_area > 0 else 0

        if 0.85 < area_ratio < 1.15:
            return rect_poly, "Candidate A (Rectangle, eps=0.005)"
        else:
            return polyB, "Candidate B (Polygon, eps=0.003)"

    elif rect_poly is not None:
        return rect_poly, "Candidate A (Rectangle, eps=0.005)"

    elif polyB is not None:
        return polyB, "Candidate B (Polygon, eps=0.005)"

    else:
        return None, "No polygon"


# -------------------------------
# API Endpoint
# -------------------------------
@app.post("/polygon")
async def polygon_endpoint(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("L")  # grayscale mask
    mask = np.array(image)

    # Run polygonization
    poly, chosen = clean_polygon_strict(mask)

    if poly is None:
        return JSONResponse(content={"chosen": chosen, "polygon": None, "image": None})

    # Draw polygon for preview
    overlay = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
    cv2.polylines(
        overlay, [poly.astype(np.int32)], isClosed=True,
        color=(0, 0, 255), thickness=2
    )

    # Encode preview as base64
    _, buffer = cv2.imencode(".png", overlay)
    img_b64 = base64.b64encode(buffer).decode("utf-8")

    return {
        "chosen": chosen,
        "polygon": poly.reshape(-1, 2).tolist(),
        "image": img_b64
    }