app.py

import gradio as gr
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as T
from PIL import Image, ImageDraw
import json
import numpy as np
from huggingface_hub import hf_hub_download


# ──────────────────────────────────────────────────────────────────────────────
# Model definition — must match the architecture used during training
# Output: (B, 8)  →  4 corners (x,y) normalised to [0, 1], order TL TR BR BL
# ──────────────────────────────────────────────────────────────────────────────

class ParkingSpaceDetector(nn.Module):
    def __init__(self):
        super().__init__()
        backbone = models.mobilenet_v2(weights=None)
        self.features = backbone.features
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(1280, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(256, 8),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.head(self.pool(self.features(x)).flatten(1))


# ──────────────────────────────────────────────────────────────────────────────
# Model loading
# Repo:  UmeshAdabala/RectArea_Parkospace  (model repo type)
# File:  best.pt  — plain OrderedDict state_dict
#        (saved with torch.save(model.state_dict(), "best.pt"))
# ──────────────────────────────────────────────────────────────────────────────

MODEL = None


def get_model():
    global MODEL
    if MODEL is not None:
        return MODEL

    print("Loading model from HuggingFace Hub ...")
    try:
        path = hf_hub_download(
            repo_id="UmeshAdabala/RectArea_Parkospace",
            filename="best.pt",
            repo_type="model",
        )
        ckpt = torch.load(path, map_location="cpu")

        m = ParkingSpaceDetector()

        if isinstance(ckpt, dict):
            # Detect the correct key for the weights
            if "model_state" in ckpt:
                state = ckpt["model_state"]          # actual format in this repo
            elif "model_state_dict" in ckpt:
                state = ckpt["model_state_dict"]
            elif "state_dict" in ckpt:
                state = ckpt["state_dict"]
            else:
                state = ckpt                          # bare flat state_dict
            m.load_state_dict(state, strict=True)
            m.eval()
            MODEL = m
        else:
            # Full model object was pickled
            ckpt.eval()
            MODEL = ckpt

        print("Model loaded successfully.")

    except Exception as e:
        print(f"Model load error: {e}  — using random weights (demo mode).")
        MODEL = ParkingSpaceDetector().eval()

    return MODEL


# ──────────────────────────────────────────────────────────────────────────────
# Image preprocessing
# ──────────────────────────────────────────────────────────────────────────────

TRANSFORM = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])


# ──────────────────────────────────────────────────────────────────────────────
# Inference helpers
# ──────────────────────────────────────────────────────────────────────────────

def run_detection(image: Image.Image):
    """Run model on image, return pixel-space corners as JSON string."""
    if image is None:
        return None, "Upload a parking space photo first."

    img = image.convert("RGB")
    W, H = img.size
    tensor = TRANSFORM(img).unsqueeze(0)       # (1, 3, 224, 224)

    with torch.no_grad():
        raw = get_model()(tensor)[0].tolist()  # 8 values in [0, 1]

    # corners TL TR BR BL — normalised → pixel
    corners_px = [[raw[i] * W, raw[i + 1] * H] for i in range(0, 8, 2)]
    payload = json.dumps({"corners": corners_px, "width": W, "height": H})
    return payload, ""


def compute_area_md(corners_json: str) -> str:
    """Estimate dimensions + pricing from pixel-space corners JSON."""
    if not corners_json:
        return ""
    try:
        data    = json.loads(corners_json)
        corners = data["corners"]           # [[x,y], ...]
        xs = [p[0] for p in corners]
        ys = [p[1] for p in corners]
        bw_px = max(xs) - min(xs)
        bh_px = max(ys) - min(ys)
        asp   = bw_px / (bh_px + 1e-6)

        L    = 5.0 if asp >= 1 else round(min(max(2.5 / asp, 1), 15), 2)
        B    = round(min(max(5.0 / asp, 1), 10), 2) if asp >= 1 else 2.5
        area = round(L * B, 2)
        price_mo = int(area * 10)

        return f"""## Detection Results

| Dimension | Value |
|-----------|-------|
| Length    | **{L} m** |
| Breadth   | **{B} m** |
| Area      | **{area} m²** |

### Suggested Pricing

| Type    | Price |
|---------|-------|
| Hourly  | Rs. 50 |
| Daily   | Rs. 300 |
| Monthly | **Rs. {price_mo}** (area x Rs. 10 / m²) |

*Drag the corner handles to fine-tune the detected region.*"""
    except Exception as e:
        return f"Error computing area: {e}"


# ──────────────────────────────────────────────────────────────────────────────
# Gradio callbacks
# ──────────────────────────────────────────────────────────────────────────────

def on_detect(image):
    """Button click: run model, return image + corners JSON + info markdown."""
    if image is None:
        return None, "", "Upload an image to begin."

    pixel_json, err = run_detection(image)
    if err:
        return image, "", err

    return image, pixel_json, compute_area_md(pixel_json)


def on_corners_change(corners_json):
    """Corners were updated by JS drag — recompute and return area markdown."""
    return compute_area_md(corners_json)


# ──────────────────────────────────────────────────────────────────────────────
# Warm up model at startup
# ──────────────────────────────────────────────────────────────────────────────

get_model()


# ──────────────────────────────────────────────────────────────────────────────
# CSS
# ──────────────────────────────────────────────────────────────────────────────

CSS = """
body { background: #0d0d1a !important; }
.gradio-container {
    background: #0d0d1a !important;
    max-width: 960px !important;
    margin: 0 auto;
}
.gr-button-primary {
    background: #2ED8DF !important;
    color: #0d0d1a !important;
    font-weight: 800 !important;
    border: none !important;
    border-radius: 8px !important;
    letter-spacing: 0.04em !important;
}
.gr-button-primary:hover { background: #12EF86 !important; }
footer { display: none !important; }
"""

# ──────────────────────────────────────────────────────────────────────────────
# Canvas HTML + JS
#
# Flow:
#   Python → hidden_img updated  → JS observer fires → canvas reloads image
#   Python → corners_state updated → JS observer fires → canvas draws handles
#   User drags handle             → JS pushes JSON → corners_state.change()
#                                                   → Python recomputes area
# ──────────────────────────────────────────────────────────────────────────────

CANVAS_BLOCK = """
<link href="https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&display=swap" rel="stylesheet">

<div id="canvas-wrap"
     style="position:relative;width:100%;touch-action:none;user-select:none;">
  <canvas id="parking-canvas"
    style="width:100%;height:auto;display:block;border-radius:10px;
           border:1px solid #1e1e3a;cursor:crosshair;background:#111127;">
  </canvas>
  <p id="canvas-hint"
     style="color:#4b5563;font-size:11px;font-family:'DM Mono',monospace;
            text-align:center;margin:6px 0 0;">
    Detected corners will appear here after clicking Detect.
  </p>
</div>

<script>
(function () {
  'use strict';

  const canvas = document.getElementById('parking-canvas');
  const ctx    = canvas.getContext('2d');
  const hint   = document.getElementById('canvas-hint');

  let imgObj   = null;
  let corners  = [];      // pixel coords in native image resolution
  let origW    = 1;
  let origH    = 1;
  let dragging = -1;

  // Hit radius in CSS pixels — generous for touch
  const HIT_CSS = 24;

  // ── Render ─────────────────────────────────────────────────────────────────

  function draw() {
    ctx.clearRect(0, 0, canvas.width, canvas.height);
    if (imgObj) ctx.drawImage(imgObj, 0, 0);
    if (corners.length < 4) return;

    // Polygon fill + stroke
    ctx.beginPath();
    corners.forEach(([x, y], i) => (i === 0 ? ctx.moveTo(x, y) : ctx.lineTo(x, y)));
    ctx.closePath();
    ctx.fillStyle   = 'rgba(46,216,223,0.18)';
    ctx.strokeStyle = 'rgba(46,216,223,0.85)';
    ctx.lineWidth   = Math.max(2, origW / 220);
    ctx.fill();
    ctx.stroke();

    // Handles — scaled so they look the same regardless of image resolution
    const scale = origW / canvas.getBoundingClientRect().width;
    const r     = HIT_CSS * scale * 0.72;
    corners.forEach(([x, y]) => {
      ctx.beginPath();
      ctx.arc(x, y, r, 0, Math.PI * 2);
      ctx.fillStyle = 'rgba(46,216,223,0.9)';
      ctx.fill();
      ctx.beginPath();
      ctx.arc(x, y, r * 0.42, 0, Math.PI * 2);
      ctx.fillStyle = '#ffffff';
      ctx.fill();
    });
  }

  // ── Pointer helpers ────────────────────────────────────────────────────────

  function canvasXY(e) {
    const rect   = canvas.getBoundingClientRect();
    const scaleX = canvas.width  / rect.width;
    const scaleY = canvas.height / rect.height;
    const src    = e.touches ? e.touches[0] : e;
    return [
      (src.clientX - rect.left) * scaleX,
      (src.clientY - rect.top)  * scaleY,
    ];
  }

  function findHandle(pos) {
    const rect    = canvas.getBoundingClientRect();
    const hitPx   = HIT_CSS * (canvas.width / rect.width) * 1.3;
    let best = -1, bestD = Infinity;
    corners.forEach(([x, y], i) => {
      const d = Math.hypot(pos[0] - x, pos[1] - y);
      if (d < hitPx && d < bestD) { bestD = d; best = i; }
    });
    return best;
  }

  // ── Event listeners ────────────────────────────────────────────────────────

  canvas.addEventListener('mousedown',  down,  { passive: false });
  canvas.addEventListener('mousemove',  move,  { passive: false });
  canvas.addEventListener('mouseup',    up);
  canvas.addEventListener('mouseleave', up);
  canvas.addEventListener('touchstart', down,  { passive: false });
  canvas.addEventListener('touchmove',  move,  { passive: false });
  canvas.addEventListener('touchend',   up);

  function down(e) {
    e.preventDefault();
    dragging = findHandle(canvasXY(e));
  }

  function move(e) {
    if (dragging < 0) return;
    e.preventDefault();
    const [cx, cy] = canvasXY(e);
    corners[dragging] = [
      Math.max(0, Math.min(origW, cx)),
      Math.max(0, Math.min(origH, cy)),
    ];
    draw();
  }

  function up() {
    if (dragging >= 0) pushCorners();
    dragging = -1;
  }

  // ── Push corners → Gradio hidden Textbox ──────────────────────────────────

  function pushCorners() {
    const ta = document.querySelector('#corners-state-box textarea');
    if (!ta) return;
    const payload = JSON.stringify({
      corners: corners.map(([x, y]) => [x, y]),
      width:   origW,
      height:  origH,
    });
    const setter = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value').set;
    setter.call(ta, payload);
    ta.dispatchEvent(new Event('input', { bubbles: true }));
  }

  // ── Bridge: observe DOM for Gradio updates ─────────────────────────────────

  let lastImgSrc   = '';
  let lastCornersVal = '';

  const obs = new MutationObserver(() => {

    // 1. Hidden image updated → reload canvas image
    const imgEl = document.querySelector('#hidden-img-out img');
    if (imgEl && imgEl.src && imgEl.src !== lastImgSrc) {
      lastImgSrc = imgEl.src;
      const img  = new Image();
      img.crossOrigin = 'anonymous';
      img.onload = () => {
        imgObj = img;
        origW  = img.naturalWidth;
        origH  = img.naturalHeight;
        canvas.width  = origW;
        canvas.height = origH;
        corners = [];
        draw();
        hint.textContent = 'Image loaded. Click Detect to find corners.';
      };
      img.src = lastImgSrc;
    }

    // 2. Corners state textbox updated by Python → parse and render
    const ta = document.querySelector('#corners-state-box textarea');
    if (ta && ta.value && ta.value !== lastCornersVal) {
      lastCornersVal = ta.value;
      try {
        const data = JSON.parse(ta.value);
        if (data && Array.isArray(data.corners) && data.corners.length === 4) {
          origW   = data.width  || origW;
          origH   = data.height || origH;
          corners = data.corners.map(([x, y]) => [x, y]);
          draw();
          hint.textContent = 'Drag the handles to fine-tune the detected area.';
        }
      } catch (_) {}
    }
  });

  obs.observe(document.body, { childList: true, subtree: true, attributes: true });

})();
</script>
"""

# ──────────────────────────────────────────────────────────────────────────────
# Gradio layout
# ──────────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="ParkoSpace — AI Area Detector") as demo:

    # Inline CSS — works across all Gradio versions
    gr.HTML(f"<style>{CSS}</style>")

    gr.HTML("""
    <link href="https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&display=swap" rel="stylesheet">
    <div style="text-align:center;padding:32px 16px 16px">
      <div style="font-size:11px;letter-spacing:0.25em;color:#2ED8DF;
                  font-family:'DM Mono',monospace;text-transform:uppercase;margin-bottom:6px">
        AI Parking Area Detector
      </div>
      <h1 style="font-family:'DM Mono',monospace;font-size:26px;font-weight:500;
                 color:#f0f4f8;margin:0;letter-spacing:-0.02em">ParkoSpace</h1>
      <p style="color:#4b5563;font-size:12px;margin-top:8px;
                font-family:'DM Mono',monospace;letter-spacing:0.03em">
        Upload a parking space photo — AI detects corners — drag to refine — area updates instantly
      </p>
    </div>
    """)

    with gr.Row(equal_height=False):

        # Left column: upload + detect button
        with gr.Column(scale=1):
            inp = gr.Image(
                type="pil",
                label="Upload Parking Space Photo",
                height=360,
            )
            btn = gr.Button("Detect Parking Area", variant="primary", size="lg")
            gr.HTML("""
            <div style="background:#111127;border:1px solid #1e1e3a;border-radius:10px;
                        padding:12px 14px;margin-top:8px">
              <p style="color:#4b5563;font-size:11px;font-family:'DM Mono',monospace;
                        margin:0;line-height:1.9">
                <span style="color:#2ED8DF">Tips for best results</span><br>
                Stand at one corner and shoot diagonally<br>
                Include all 4 corners in the frame<br>
                Good lighting — avoid strong shadows<br>
                Works with car parks, open spaces, garages
              </p>
            </div>
            """)

        # Right column: canvas + results
        with gr.Column(scale=1):
            gr.HTML(CANVAS_BLOCK)
            out_md = gr.Markdown("*Upload a photo and click Detect to see results.*")

    # ── Hidden bridge components ───────────────────────────────────────────────

    # Holds pixel-space corners JSON; written by Python after detection
    # and by JS after a drag.
    corners_state = gr.Textbox(
        value="",
        visible=False,
        label="corners_state",
        elem_id="corners-state-box",
    )

    # Holds the PIL image; Python writes here so JS can grab the img src URL.
    hidden_img = gr.Image(
        type="pil",
        visible=False,
        label="hidden_img",
        elem_id="hidden-img-out",
    )

    # ── Event wiring ───────────────────────────────────────────────────────────

    btn.click(
        fn=on_detect,
        inputs=[inp],
        outputs=[hidden_img, corners_state, out_md],
    )

    corners_state.change(
        fn=on_corners_change,
        inputs=[corners_state],
        outputs=[out_md],
    )

    gr.HTML("""
    <div style="margin-top:24px;padding:14px;background:#111127;border-radius:10px;
                border:1px solid #1e1e3a">
      <p style="color:#2e3650;font-size:11px;font-family:'DM Mono',monospace;
                text-align:center;margin:0">
        Model: UmeshAdabala/RectArea_Parkospace &nbsp;|&nbsp;
        MobileNetV2 + Keypoint Regression &nbsp;|&nbsp; ~13 MB &nbsp;|&nbsp; MIT License
      </p>
    </div>
    """)


if __name__ == "__main__":
    demo.launch(share = True)