"""
HemaVision Gradio Interface — Tier 1
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Production-quality Gradio demo with clean, clinical-grade design.
Features:
• Auto-segments multi-cell blood smear fields into individual crops
• Analyzes each cell independently with Grad-CAM overlays
• Rich HTML result cards — no monospace text boxes
• Input validation with clear user guidance
Author: Firoj
"""
import io
import base64
import logging
from pathlib import Path
from typing import List, Optional, Tuple
import gradio as gr
import numpy as np
import torch
from PIL import Image
from core.model import DualStreamFusionModel
from core.gradcam import GradCAM
from core.dataset import get_eval_transforms
from core.train import AMLTrainer
from core.morphology import (
extract_single_image_features,
MORPHOLOGY_FEATURE_NAMES,
NUM_MORPHOLOGY_FEATURES,
)
from core.cell_segmenter import segment_cells, SegmentationResult, CellCrop
from utils.config import get_config, AugmentationConfig
logger = logging.getLogger(__name__)
# ── Global State ─────────────────────────────────────────────
MODEL: Optional[DualStreamFusionModel] = None
GRADCAM: Optional[GradCAM] = None
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TRANSFORM = get_eval_transforms()
CONFIG = get_config()
OPTIMAL_THRESHOLD: float = 0.5 # Updated when checkpoint is loaded
# Feature configuration — morphological features extracted from images
TABULAR_FEATURE_NAMES = list(MORPHOLOGY_FEATURE_NAMES) # 20 features
def load_model(checkpoint_path: Optional[str] = None) -> DualStreamFusionModel:
"""Load the trained model or create a fresh one for demo."""
global MODEL, GRADCAM, OPTIMAL_THRESHOLD
# For Hugging Face Spaces: prefer final_model.pt in current directory
space_final_path = Path("final_model.pt")
if space_final_path.exists():
checkpoint_path = str(space_final_path)
logger.info(f"[HF Space] Using final_model.pt from current directory: {checkpoint_path}")
# Original auto-discovery logic (used locally or if final_model.pt missing)
if not checkpoint_path:
default_ckpt = CONFIG.paths.checkpoints_dir / "best_model.pt"
if default_ckpt.exists():
checkpoint_path = str(default_ckpt)
logger.info(f"Auto-discovered checkpoint: {checkpoint_path}")
if checkpoint_path and Path(checkpoint_path).exists():
# Peek at checkpoint for optimal threshold
ckpt = torch.load(checkpoint_path, map_location=DEVICE, weights_only=False)
OPTIMAL_THRESHOLD = ckpt.get("optimal_threshold", 0.5)
# If best_model.pt has default threshold, try final_model.pt for the optimized one
if OPTIMAL_THRESHOLD == 0.5:
final_ckpt_path = Path(checkpoint_path).parent / "final_model.pt"
if final_ckpt_path.exists():
final_ckpt = torch.load(str(final_ckpt_path), map_location=DEVICE, weights_only=False)
opt_thresh = final_ckpt.get("optimal_threshold", 0.5)
if opt_thresh != 0.5:
OPTIMAL_THRESHOLD = opt_thresh
logger.info(f"Loaded optimized threshold from final_model.pt: {OPTIMAL_THRESHOLD:.4f}")
logger.info(f"Using optimal threshold: {OPTIMAL_THRESHOLD:.4f}")
# Get num_tabular_features from checkpoint config
saved_config = ckpt.get("config", {})
num_tab = saved_config.get("num_tabular_features", len(TABULAR_FEATURE_NAMES))
model = AMLTrainer.load_checkpoint(
checkpoint_path,
num_tabular_features=num_tab,
device=DEVICE,
)
else:
# Demo mode — create untrained model
logger.info("No checkpoint found. Running in demo mode.")
model = DualStreamFusionModel(
num_tabular_features=len(TABULAR_FEATURE_NAMES),
)
model = model.to(DEVICE)
model.eval()
MODEL = model
GRADCAM = GradCAM(model, target_layers=["layer3", "layer4"])
return model
def _predict_single_cell(
image_pil: Image.Image,
) -> Tuple[float, Image.Image]:
"""
Run prediction + Grad-CAM on a single cell crop.
Returns:
(probability, gradcam_overlay_pil)
"""
image_tensor = TRANSFORM(image_pil).unsqueeze(0).to(DEVICE)
morph_vec = extract_single_image_features(image_pil, normalize=True)
tabular = torch.tensor(morph_vec, dtype=torch.float32).unsqueeze(0).to(DEVICE)
expected = MODEL.num_tabular_features
if tabular.shape[1] < expected:
pad = torch.zeros(1, expected - tabular.shape[1], device=DEVICE)
tabular = torch.cat([tabular, pad], dim=1)
elif tabular.shape[1] > expected:
tabular = tabular[:, :expected]
MODEL.eval()
heatmap, prob = GRADCAM.generate(image_tensor, tabular)
original_np = np.array(image_pil.resize((224, 224)))
overlay = GRADCAM.create_overlay(original_np, heatmap, alpha=0.45)
return prob, Image.fromarray(overlay)
def predict(
image: Image.Image,
) -> Tuple[str, Optional[Image.Image]]:
"""
Run AML prediction — auto-segments multi-cell images.
If the uploaded image contains multiple cells (e.g. a whole
blood smear field), each cell is cropped and analyzed separately.
Single-cell crops are analyzed directly.
Returns:
(result_html, composite_gradcam_image)
"""
global MODEL, GRADCAM
if MODEL is None:
load_model()
if image is None:
return "⚠️ Please upload a cell microscopy image.", None
image_pil = image.convert("RGB") if isinstance(image, Image.Image) else Image.fromarray(image).convert("RGB")
# ── Auto-segment ─────────────────────────────────────────
seg = segment_cells(image_pil, max_cells=20, annotate=True)
cells = seg.cells
n_cells = len(cells)
if n_cells == 0:
return _empty_state_html(), None
# ── Run predictions on each cell ─────────────────────────
cell_results = [] # list of (cell_crop, prob, gradcam_pil)
for cell in cells:
prob, gcam = _predict_single_cell(cell.image)
cell_results.append((cell, prob, gcam))
# ── Build composite Grad-CAM image ───────────────────────
if n_cells == 1:
composite = cell_results[0][2]
else:
composite = _build_composite_image(cell_results, seg)
# ── Build HTML ───────────────────────────────────────────
expected = MODEL.num_tabular_features
result_html = _build_result_html(cell_results, seg, expected)
return result_html, composite
def _build_composite_image(
cell_results: list,
seg: SegmentationResult,
) -> Image.Image:
"""Build a grid of Grad-CAM overlays, plus the annotated source image."""
n = len(cell_results)
# Decide grid layout: annotated full image on left, cell grid on right
cell_size = 160
cols = min(n, 4)
rows = (n + cols - 1) // cols
grid_w = cols * cell_size
grid_h = rows * cell_size
# Scale annotated image to fit beside the grid
ann = seg.annotated_image or Image.new("RGB", (400, 400), (240, 240, 240))
ann_h = grid_h if grid_h > 0 else 400
ann_w = int(ann.width * ann_h / max(ann.height, 1))
ann_w = min(ann_w, 500)
ann = ann.resize((ann_w, ann_h), Image.LANCZOS)
total_w = ann_w + 16 + grid_w
total_h = max(ann_h, grid_h)
canvas = Image.new("RGB", (total_w, total_h), (15, 25, 35))
# Paste annotated image on left
canvas.paste(ann, (0, (total_h - ann_h) // 2))
# Paste Grad-CAM cells in a grid on the right
for i, (cell, prob, gcam) in enumerate(cell_results):
r, c = divmod(i, cols)
gcam_resized = gcam.resize((cell_size, cell_size), Image.LANCZOS)
x_off = ann_w + 16 + c * cell_size
y_off = r * cell_size
canvas.paste(gcam_resized, (x_off, y_off))
# Draw a thin colored border (blast=red, normal=green)
try:
import cv2
arr = np.array(canvas)
color = (220, 38, 38) if prob > OPTIMAL_THRESHOLD else (5, 150, 105)
cv2.rectangle(arr, (x_off, y_off),
(x_off + cell_size - 1, y_off + cell_size - 1),
color, 2)
# Cell number label
label = f"#{i+1}"
cv2.putText(arr, label, (x_off + 4, y_off + 16),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
cv2.putText(arr, label, (x_off + 4, y_off + 16),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
canvas = Image.fromarray(arr)
except ImportError:
pass
return canvas
def _build_result_html(
cell_results: list,
seg: SegmentationResult,
num_features: int,
) -> str:
"""Build an HTML result card for single or multi-cell analysis."""
n = len(cell_results)
blasts = [(c, p, g) for c, p, g in cell_results if p > OPTIMAL_THRESHOLD]
normals = [(c, p, g) for c, p, g in cell_results if p <= OPTIMAL_THRESHOLD]
n_blast = len(blasts)
n_normal = len(normals)
blast_pct = n_blast / n * 100 if n > 0 else 0
# Overall assessment
if n_blast == 0:
overall_label = "Normal — No Blasts Detected"
overall_bg = "#071A12"; overall_border = "#0D3B24"; overall_text = "#6EE7B7"
overall_icon = (
''
' '
)
risk_label = "LOW RISK"
risk_emoji = "🟢"
elif blast_pct >= 20:
overall_label = f"AML Suspected — {n_blast}/{n} Blast Cells ({blast_pct:.0f}%)"
overall_bg = "#1C0A0A"; overall_border = "#3B1515"; overall_text = "#FCA5A5"
overall_icon = (
''
' '
)
risk_label = "HIGH RISK"
risk_emoji = "🔴"
else:
overall_label = f"Atypical — {n_blast}/{n} Blast Cells ({blast_pct:.0f}%)"
overall_bg = "#1C1505"; overall_border = "#3B2C0A"; overall_text = "#FCD34D"
overall_icon = (
''
' '
)
risk_label = "MODERATE RISK"
risk_emoji = "🟡"
# Multi-cell info banner
multi_banner = ""
if seg.is_multi_cell:
multi_banner = (
''
f'🔍 Auto-Segmentation: {seg.message}'
'
'
)
# Per-cell table rows (only show for multi-cell)
cell_rows = ""
if n > 1:
rows_html = ""
for cell, prob, gcam in cell_results:
is_blast = prob > OPTIMAL_THRESHOLD
conf = prob if is_blast else 1 - prob
dot = '● ' if is_blast else '● '
cls = "AML Blast" if is_blast else "Normal"
rows_html += (
f''
f'#{cell.index+1} '
f'{dot} {cls} '
f'{prob:.4f} '
f'{conf:.1%} '
f' '
)
cell_rows = (
''
'
Per-Cell Breakdown
'
'
'
''
'Cell '
'Classification '
'P(blast) '
'Confidence '
' {rows_html}
Confidence
{conf:.1%}
P(blast) = {prob:.4f}
threshold = {OPTIMAL_THRESHOLD:.4f}
"""
else:
avg_prob = sum(p for _, p, _ in cell_results) / n
bar_color = "#DC2626" if blast_pct >= 20 else "#D97706" if n_blast > 0 else "#059669"
single_metrics = f"""
Blast Proportion
{n_blast}/{n} ({blast_pct:.0f}%)
Avg P(blast) = {avg_prob:.4f}
threshold = {OPTIMAL_THRESHOLD:.4f}
"""
html = f"""
{multi_banner}
{overall_icon}
{overall_label}
{risk_emoji} {risk_label}
{single_metrics}
{cell_rows}
Cells Analyzed
{n} cell{'s' if n != 1 else ''}
Architecture
ResNet50 + MLP
Features
{num_features} morphological
🧠 ResNet50 Visual Stream
📐 20 Morphological Features
🔥 Grad-CAM Explainability
⚡ <50ms Inference
""")
with gr.Row(equal_height=False):
# ── Left: Upload ─────────────────────────────────
with gr.Column(scale=5, min_width=360):
gr.HTML("""
""")
image_input = gr.Image(
type="pil",
label="Upload microscopic cell image",
elem_classes="upload-zone",
height=300,
sources=["upload", "clipboard"],
)
gr.HTML("""
💡 Accepts both: Upload a single-cell crop
or a whole blood smear field — the system will auto-detect
and segment individual cells for per-cell analysis.
""")
analyze_btn = gr.Button(
"Analyze",
variant="primary",
elem_classes="analyze-btn",
size="lg",
)
# ── Right: Results ───────────────────────────────
with gr.Column(scale=6, min_width=380):
gr.HTML("""
""")
result_output = gr.HTML(
value=_empty_state_html(),
elem_classes="result-area",
)
gr.HTML("""
""")
gradcam_output = gr.Image(
label="Grad-CAM Visualization",
elem_classes="gradcam-img",
height=260,
show_label=False,
)
# ── Footer ───────────────────────────────────────────
gr.HTML("""
""")
# ── Events ───────────────────────────────────────────
analyze_btn.click(
fn=predict,
inputs=[image_input],
outputs=[result_output, gradcam_output],
)
return demo, theme
def _empty_state_html() -> str:
"""Placeholder shown before the user uploads an image."""
return """
No analysis yet
Upload a cell image and click Analyze to begin
"""
# ── Entry Point ──────────────────────────────────────────────
if __name__ == "__main__":
import os
logger.info(f"Current working directory: {os.getcwd()}")
logger.info(f"Files visible in root: {[f.name for f in Path('.').iterdir()]}")
# Try to load model — will prefer final_model.pt if present in Space
load_model()
app, theme = create_gradio_app()
app.launch(
server_name="0.0.0.0",
server_port=7860, # ← important for Hugging Face Spaces
share=False,
show_error=True,
theme=theme,
css=CUSTOM_CSS,
js=FORCE_DARK_JS,
)