likingood/Skin_Lesion_Segmentation

1import json2from pathlib import Path3 4import cv25import gradio as gr6import numpy as np7from ultralytics import YOLO8 9APP_DIR = Path(__file__).parent10MODEL_PATH = APP_DIR / "models" / "yolo11n_seg_best.pt"11# bundled copy travels with the HF Space; fall back to the project Results/ folder for local dev12METRICS_PATH = APP_DIR / "model_metrics.json"13if not METRICS_PATH.exists():14    METRICS_PATH = APP_DIR.parent / "Results" / "test_metrics.json"15 16model = YOLO(str(MODEL_PATH))17 18DISCLAIMER = (19    "Educational demonstration only. Not for diagnosis or malignancy classification. "20    "Area/diameter are pixel-based estimates with no physical scale reference."21)22 23 24def run_segmentation(image_bgr, confidence):25    results = model.predict(source=image_bgr, conf=confidence, verbose=False)26    result = results[0]27    annotated = cv2.cvtColor(result.plot(), cv2.COLOR_BGR2RGB)28 29    if result.masks is None or len(result.masks.data) == 0:30        return annotated, None, None31 32    area_px = float(result.masks.data[0].sum())33    diameter_px = 2 * (area_px / np.pi) ** 0.534    return annotated, area_px, diameter_px35 36 37def detect_single(image: np.ndarray, confidence: float):38    if image is None:39        return None, "Upload a dermoscopy image to begin."40 41    # gr.Image(type="numpy") decodes to RGB; Ultralytics expects BGR for raw arrays.42    annotated, area_px, diameter_px = run_segmentation(cv2.cvtColor(image, cv2.COLOR_RGB2BGR), confidence)43    if area_px is None:44        lines = ["**No lesion boundary detected at this confidence threshold.**"]45    else:46        lines = [47            f"**Estimated lesion area: {area_px:,.0f} px**",48            f"**Estimated equivalent diameter: {diameter_px:,.0f} px**",49        ]50    lines += ["", f"_{DISCLAIMER}_"]51    return annotated, "\n".join(lines)52 53 54def detect_batch(files, confidence: float):55    if not files:56        return [], "Upload one or more dermoscopy images to begin."57 58    gallery = []59    rows = ["| Image | Area (px) | Diameter (px) |", "|---|---|---|"]60    areas = []61 62    for f in files:63        path = f if isinstance(f, str) else f.name64        image_bgr = cv2.imread(path)  # already BGR, matches what Ultralytics expects65        annotated, area_px, diameter_px = run_segmentation(image_bgr, confidence)66        gallery.append((annotated, Path(path).name))67 68        if area_px is None:69            rows.append(f"| {Path(path).name} | - | - |")70        else:71            areas.append(area_px)72            rows.append(f"| {Path(path).name} | {area_px:,.0f} | {diameter_px:,.0f} |")73 74    if areas:75        rows.append(f"| **Mean ({len(areas)}/{len(files)} detected)** | **{np.mean(areas):,.0f}** | - |")76    table = "\n".join(rows) + f"\n\n_{DISCLAIMER}_"77    return gallery, table78 79 80def load_stats_markdown():81    if not METRICS_PATH.exists():82        return "Model performance stats not available yet."83    m = json.load(open(METRICS_PATH))84    return "\n".join([85        "| Metric | Value | What it means |",86        "|---|---|---|",87        f"| Mask precision (mean) | {m['mask_precision']:.3f} | Of all predicted lesion masks, the fraction that were correct — higher means fewer false boundaries |",88        f"| Mask recall (mean) | {m['mask_recall']:.3f} | Of all real lesions present, the fraction the model segmented — higher means fewer missed lesions |",89        f"| Mask mAP@50 | {m['mask_mAP50']:.3f} | Segmentation accuracy when the predicted mask only needs 50% overlap with the true lesion — a lenient pass/fail bar |",90        f"| Mask mAP@50-95 | {m['mask_mAP50-95']:.3f} | Same idea averaged over stricter overlap requirements (50-95%) — the harder, headline metric |",91        f"| Box mAP@50 | {m['box_mAP50']:.3f} | Same lenient measure but for the lesion's bounding box rather than its exact pixel boundary |",92        f"| Box mAP@50-95 | {m['box_mAP50-95']:.3f} | Stricter version of the bounding-box measure above |",93        "",94        "_Evaluated on the held-out ISIC2016 test split (379 images), never seen during training "95        "or model selection._",96    ])97 98 99assets_dir = APP_DIR / "assets"100single_examples = sorted(str(p) for p in assets_dir.glob("*.jpg")) if assets_dir.exists() else []101demo_batch_dir = assets_dir / "demo_batch"102batch_example = sorted(str(p) for p in demo_batch_dir.glob("*.jpg")) if demo_batch_dir.exists() else []103 104with gr.Blocks(title="Skin Lesion Segmentation — BN4601A Individual Assignment") as demo:105    gr.Markdown(106        "# Skin Lesion Segmentation — BN4601A Individual Assignment\n\n"107        "YOLO11n-seg fine-tuned on ISIC 2016 to segment the lesion boundary in a dermoscopy image, "108        "as a visual aid for border delineation.\n\n" + DISCLAIMER109    )110 111    with gr.Tabs():112        with gr.Tab("Single image"):113            with gr.Row():114                single_input = gr.Image(type="numpy", label="Upload dermoscopy image")115                single_output = gr.Image(label="Segmented lesion (YOLO11n-seg)")116            single_conf = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold")117            single_measure = gr.Markdown(label="Measurements")118            single_btn = gr.Button("Segment", variant="primary")119            single_btn.click(detect_single, [single_input, single_conf], [single_output, single_measure])120            if single_examples:121                gr.Examples([[e, 0.25] for e in single_examples], [single_input, single_conf])122 123        with gr.Tab("Multiple images (batch)"):124            batch_input = gr.File(file_count="multiple", file_types=["image"], label="Upload dermoscopy images")125            batch_conf = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold")126            batch_btn = gr.Button("Segment all", variant="primary")127            batch_gallery = gr.Gallery(label="Segmented lesion per image", columns=3)128            batch_measure = gr.Markdown(label="Per-image and mean measurements")129            batch_btn.click(detect_batch, [batch_input, batch_conf], [batch_gallery, batch_measure])130            if batch_example:131                gr.Examples([[batch_example, 0.25]], [batch_input, batch_conf], label="Demo sample set (unseen images)")132 133        with gr.Tab("Model performance"):134            gr.Markdown(load_stats_markdown())135 136if __name__ == "__main__":137    import os138    demo.launch(server_port=int(os.environ.get("PORT", 7862)), allowed_paths=[str(APP_DIR)])139