app.py

import gradio as gr
import cv2
import numpy as np
import io
import os
import zipfile
import tempfile
from PIL import Image
import matplotlib
matplotlib.use("Agg")

# ─── Cellpose model (lazy) ────────────────────────────────────────────────────
_model = None

def get_model():
    global _model
    if _model is None:
        from cellpose import models
        from huggingface_hub import hf_hub_download
        fpath = hf_hub_download(repo_id="mouseland/cellpose-sam", filename="cpsam")
        _model = models.CellposeModel(gpu=False, pretrained_model=fpath)
    return _model

# ─── Image helpers ────────────────────────────────────────────────────────────
def normalize99(img):
    X = img.copy().astype(np.float32)
    p1, p99 = np.percentile(X, 1), np.percentile(X, 99)
    return (X - p1) / (1e-10 + p99 - p1)

def image_resize(img, resize=1000):
    ny, nx = img.shape[:2]
    if max(ny, nx) > resize:
        if ny > nx:
            nx = int(nx / ny * resize); ny = resize
        else:
            ny = int(ny / nx * resize); nx = resize
        img = cv2.resize(img, (nx, ny))
    return img.astype(np.uint8)

def run_cellpose(img, model, flow_threshold=0.4, cellprob_threshold=0.0):
    masks, flows, _ = model.eval(
        img, niter=250,
        flow_threshold=flow_threshold,
        cellprob_threshold=cellprob_threshold,
    )
    return masks

# ─── YOLO Annotation Exporter ─────────────────────────────────────────────────
def export_yolo_annotations(masks, img_shape, class_id=0):
    """
    Converts Cellpose masks → YOLO segmentation format.

    YOLO segmentation line format:
        class_id  x1 y1 x2 y2 ... (all normalized 0–1)

    class_id = 0 → 'grain'  (you will split into broken/whole on Roboflow)
    """
    h, w = img_shape[:2]
    lines = []
    num_grains = int(masks.max())

    for i in range(1, num_grains + 1):
        # Binary mask for this single grain
        single = (masks == i).astype(np.uint8)
        contours, _ = cv2.findContours(single, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if not contours:
            continue

        # Pick the largest contour (in case of tiny noise)
        c = max(contours, key=cv2.contourArea)
        c = c.squeeze()

        if c.ndim < 2 or len(c) < 4:
            continue

        # Normalize each point to [0, 1]
        norm_pts = []
        for x, y in c:
            norm_pts.append(round(float(x) / w, 6))
            norm_pts.append(round(float(y) / h, 6))

        pts_str = " ".join(map(str, norm_pts))
        lines.append(f"{class_id} {pts_str}")

    return "\n".join(lines), num_grains


def make_preview(img_np, masks):
    """Draw red outlines of all grain masks on the image for preview."""
    preview = img_np.copy()
    for i in range(1, int(masks.max()) + 1):
        single = (masks == i).astype(np.uint8)
        contours, _ = cv2.findContours(single, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(preview, contours, -1, (220, 38, 38), 2)
    return Image.fromarray(preview)


# ─── Main batch processor ─────────────────────────────────────────────────────
def process_batch(image_files, flow_threshold, cellprob_threshold, progress=gr.Progress()):
    """
    Takes a list of uploaded image file paths.
    Returns:
      - Gallery of preview images (with outlines)
      - Summary text
      - Path to downloadable ZIP
    """
    if not image_files:
        return [], "⚠️ No images uploaded.", None

    model = get_model()

    previews    = []   # (PIL image, caption) for gallery
    log_lines   = []
    total_grains = 0
    failed       = []

    # Temp folder to collect annotation files
    tmp_dir = tempfile.mkdtemp()
    images_dir = os.path.join(tmp_dir, "images")
    labels_dir = os.path.join(tmp_dir, "labels")
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    for idx, file_obj in enumerate(progress.tqdm(image_files, desc="Processing images")):
        # file_obj is a filepath string when using gr.File with type="filepath"
        filepath = file_obj if isinstance(file_obj, str) else file_obj.name
        fname    = os.path.splitext(os.path.basename(filepath))[0]

        try:
            pil_img    = Image.open(filepath).convert("RGB")
            img_np     = np.array(pil_img)
            img_np     = image_resize(img_np, resize=1000)

            masks = run_cellpose(img_np, model,
                                  flow_threshold=float(flow_threshold),
                                  cellprob_threshold=float(cellprob_threshold))

            num_grains = int(masks.max())

            if num_grains == 0:
                log_lines.append(f"⚠️  [{idx+1}] {fname} — No grains detected, skipped.")
                failed.append(fname)
                continue

            # Export YOLO annotation txt
            annotation_txt, _ = export_yolo_annotations(masks, img_np.shape, class_id=0)
            txt_path = os.path.join(labels_dir, f"{fname}.txt")
            with open(txt_path, "w") as f:
                f.write(annotation_txt)

            # Save image to images/
            img_save_path = os.path.join(images_dir, f"{fname}.jpg")
            Image.fromarray(img_np).save(img_save_path, quality=95)

            # Make preview
            preview_pil = make_preview(img_np, masks)
            previews.append((preview_pil, f"{fname} — {num_grains} grains"))

            total_grains += num_grains
            log_lines.append(f"✅  [{idx+1}] {fname} — {num_grains} grains annotated.")

        except Exception as e:
            log_lines.append(f"❌  [{idx+1}] {fname} — Error: {str(e)}")
            failed.append(fname)

    # ── Write data.yaml ───────────────────────────────────────────────────────
    yaml_content = (
        "# YOLO Dataset — Rice Grain Segmentation\n"
        "# Generated by MLBench Annotation Tool\n\n"
        "path: ./dataset\n"
        "train: images/train\n"
        "val:   images/val\n\n"
        "nc: 2\n"
        "names:\n"
        "  0: whole_grain\n"
        "  1: broken_grain\n\n"
        "# NOTE: All grains are currently class 0 (whole_grain).\n"
        "# Upload to Roboflow and re-label broken grains as class 1.\n"
    )
    with open(os.path.join(tmp_dir, "data.yaml"), "w") as f:
        f.write(yaml_content)

    # ── Write README ──────────────────────────────────────────────────────────
    readme = (
        "# Rice Grain YOLO Dataset\n\n"
        "## Folder Structure\n"
        "```\n"
        "dataset/\n"
        "  images/   ← your rice photos (.jpg)\n"
        "  labels/   ← YOLO polygon annotations (.txt)\n"
        "  data.yaml ← class config for YOLO training\n"
        "```\n\n"
        "## Label Format (YOLO Segmentation)\n"
        "Each .txt file has one line per grain:\n"
        "```\n"
        "class_id  x1 y1 x2 y2 x3 y3 ...  (normalized 0–1)\n"
        "```\n\n"
        "## Classes\n"
        "| ID | Name         |\n"
        "|----|-------------|\n"
        "| 0  | whole_grain  |\n"
        "| 1  | broken_grain |\n\n"
        "## Next Steps\n"
        "1. Upload this zip to **Roboflow** (Import > YOLOv8 Segmentation format)\n"
        "2. Re-label broken grains as class `1` in Roboflow\n"
        "3. Export from Roboflow as YOLOv8 format\n"
        "4. Train: `yolo segment train data=data.yaml model=yolov8n-seg.pt epochs=100`\n"
    )
    with open(os.path.join(tmp_dir, "README.md"), "w") as f:
        f.write(readme)

    # ── Package as ZIP ────────────────────────────────────────────────────────
    zip_path = os.path.join(tempfile.mkdtemp(), "rice_yolo_dataset.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for root, _, files in os.walk(tmp_dir):
            for file in files:
                full_path = os.path.join(root, file)
                arcname   = os.path.relpath(full_path, tmp_dir)
                zf.write(full_path, arcname)

    # ── Summary ───────────────────────────────────────────────────────────────
    ok_count = len(image_files) - len(failed)
    summary = (
        f"### ✅ Done!\n"
        f"- **{ok_count} / {len(image_files)}** images processed\n"
        f"- **{total_grains}** total grains annotated\n"
        f"- **{len(failed)}** failed: {', '.join(failed) if failed else 'none'}\n\n"
        "**Download the ZIP below → upload to Roboflow → label broken grains → train YOLO!**\n\n"
        "---\n" + "\n".join(log_lines)
    )

    return previews, summary, zip_path


# ─── UI ───────────────────────────────────────────────────────────────────────
CSS = """
body { font-family: 'IBM Plex Mono', monospace; }
#header { 
    background: #0F172A; 
    padding: 20px 24px 14px; 
    border-radius: 10px; 
    margin-bottom: 12px; 
}
#run-btn { margin-top: 8px; background: #7C3AED !important; }
#dl-btn  { margin-top: 6px; }
.gr-gallery-item img { border-radius: 6px; }
"""

THEME = gr.themes.Soft(
    primary_hue="violet",
    secondary_hue="indigo",
    neutral_hue="slate",
)

with gr.Blocks(theme=THEME, css=CSS, title="Rice YOLO Annotator") as demo:

    gr.HTML("""
    <div id="header">
      <span style="font-size:1.9rem;font-weight:900;color:#F1F5F9;font-family:monospace;">
        ML<span style="color:#EF4444;">Bench</span>
        <span style="font-size:1rem;font-weight:400;color:#94A3B8;margin-left:12px;">
          Rice Grain → YOLO Annotation Exporter
        </span>
      </span>
      <p style="color:#64748B;font-size:0.85rem;margin-top:6px;font-family:monospace;">
        Upload up to 50 images · Cellpose segments each grain · 
        Download ZIP with YOLO labels ready for Roboflow
      </p>
    </div>
    """)

    with gr.Row():
        # ── LEFT ──────────────────────────────────────────────────────────────
        with gr.Column(scale=1):
            gr.Markdown("### 📂 Upload Images")
            image_input = gr.File(
                file_count="multiple",
                file_types=["image"],
                label="Drop up to 50 rice images here",
                height=180,
            )

            with gr.Accordion("⚙️ Cellpose Settings", open=False):
                flow_thresh = gr.Slider(
                    0.0, 1.0, value=0.4, step=0.05,
                    label="Flow Threshold",
                    info="Higher = stricter (fewer false grains)"
                )
                cellprob_thresh = gr.Slider(
                    -4.0, 4.0, value=0.0, step=0.5,
                    label="Cell Probability Threshold",
                    info="Lower = detect more grains"
                )

            run_btn = gr.Button(
                "🚀  Run Cellpose & Export Annotations",
                variant="primary", size="lg", elem_id="run-btn"
            )

            gr.Markdown("""
            ### 📋 Workflow
            1. Upload 50 images here  
            2. Click **Run** — Cellpose segments every grain  
            3. Download the ZIP  
            4. Upload ZIP to **Roboflow** (format: YOLOv8 Segmentation)  
            5. Re-label broken grains as `broken_grain` class  
            6. Export & train YOLOv8!  
            """)

            download_btn = gr.File(
                label="⬇️ Download YOLO Dataset ZIP",
                interactive=False,
                elem_id="dl-btn",
            )

        # ── RIGHT ─────────────────────────────────────────────────────────────
        with gr.Column(scale=2):
            gr.Markdown("### 🔍 Segmentation Previews")
            gallery = gr.Gallery(
                label="",
                show_label=False,
                columns=3,
                height=460,
                object_fit="contain",
            )
            summary_box = gr.Markdown(
                value="_Results will appear here after processing..._"
            )

    run_btn.click(
        fn=process_batch,
        inputs=[image_input, flow_thresh, cellprob_thresh],
        outputs=[gallery, summary_box, download_btn],
    )

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