""" inference.py ------------ Step 8: Report Assembly Purpose: End-to-end inference pipeline for a single uploaded image. Runs all steps in sequence and assembles a complete biological observation report. Added features: - Confidence scores per cell (Random Forest probability) - GradCAM heatmap on U-Net encoder (visual explanation) - PDF report generation (downloadable summary) Pipeline: 1. Load + preprocess uploaded image 2. U-Net predicts binary mask + GradCAM heatmap 3. Connected components → cell instances 4. skimage.regionprops → 7 morphological features per cell 5. Random Forest → health label + confidence score per cell 6. Benchmark comparison → population status 7. Assemble structured report + PDF Input: Any brightfield image (.tif / .png / .jpg / .npy) models/unet_best.pth -> trained U-Net models/health_classifier.pkl -> trained RF classifier Output: InferenceReport dataclass outputs/report.pdf -> PDF report outputs/gradcam.png -> GradCAM heatmap """ from pathlib import Path from dataclasses import dataclass, field from typing import List, Dict, Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import cv2 import joblib import matplotlib.pyplot as plt import matplotlib.cm as cm from skimage import measure from reportlab.lib.pagesizes import A4 from reportlab.lib import colors from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle from reportlab.lib.units import cm as rcm from reportlab.platypus import (SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image as RLImage) from reportlab.lib.enums import TA_CENTER import os # ── Paths (HF compatible) ───────────────────────────────────────────────────── BASE_DIR = Path(__file__).resolve().parent.parent UNET_PATH = BASE_DIR / "models" / "unet_best.pth" CLASSIFIER_PATH = BASE_DIR / "models" / "health_classifier.pkl" OUTPUT_DIR = BASE_DIR / "outputs" OUTPUT_DIR.mkdir(exist_ok=True) # ── Report dataclass ────────────────────────────────────────────────────────── @dataclass class InferenceReport: filename : str n_cells : int confluency_pct : float mean_area : float mean_circularity : float mean_solidity : float mean_intensity : float healthy_pct : float stressed_pct : float apoptotic_pct : float mean_confidence : float overall_status : str observations : List[str] = field(default_factory=list) recommendations : List[str] = field(default_factory=list) caveats : List[str] = field(default_factory=list) overlay_image : Optional[np.ndarray] = field(default=None) gradcam_image : Optional[np.ndarray] = field(default=None) cell_details : List[Dict] = field(default_factory=list) def to_dict(self): return { "filename" : self.filename, "n_cells" : self.n_cells, "confluency_pct" : self.confluency_pct, "mean_area" : self.mean_area, "mean_circularity": self.mean_circularity, "mean_solidity" : self.mean_solidity, "mean_intensity" : self.mean_intensity, "healthy_pct" : self.healthy_pct, "stressed_pct" : self.stressed_pct, "apoptotic_pct" : self.apoptotic_pct, "mean_confidence" : self.mean_confidence, "overall_status" : self.overall_status, "observations" : self.observations, "recommendations" : self.recommendations, "caveats" : self.caveats, } # ── U-Net ───────────────────────────────────────────────────────────────────── class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Dropout2d(0.1), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch)) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 = self.up(x1) diffY = x2.size(2) - x1.size(2) diffX = x2.size(3) - x1.size(3) x1 = F.pad(x1, [diffX//2, diffX-diffX//2, diffY//2, diffY-diffY//2]) return self.conv(torch.cat([x2, x1], dim=1)) class UNet(nn.Module): def __init__(self, base=32): super().__init__() self.inc = DoubleConv(1, base) self.down1 = Down(base, base*2) self.down2 = Down(base*2, base*4) self.down3 = Down(base*4, base*8) self.down4 = Down(base*8, base*16) self.up1 = Up(base*16 + base*8, base*8) self.up2 = Up(base*8 + base*4, base*4) self.up3 = Up(base*4 + base*2, base*2) self.up4 = Up(base*2 + base, base) self.out = nn.Conv2d(base, 1, 1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return torch.sigmoid(self.out(x)) # ── Model loaders ───────────────────────────────────────────────────────────── _unet = None _classifier = None def load_unet(device): global _unet if _unet is None: model = UNet(base=32).to(device) model.load_state_dict(torch.load(UNET_PATH, map_location=device, weights_only=True)) model.eval() _unet = model return _unet def load_classifier(): global _classifier if _classifier is None: _classifier = joblib.load(CLASSIFIER_PATH) return _classifier # ── Preprocessing ───────────────────────────────────────────────────────────── def preprocess(image_input, size=256): if isinstance(image_input, (str, Path)): p = Path(image_input) if p.suffix == ".npy": img = np.load(p).astype(np.float32) filename = p.name if img.max() <= 1.0: return img, filename else: img = cv2.imread(str(p), cv2.IMREAD_ANYDEPTH | cv2.IMREAD_GRAYSCALE) if img is None: raise FileNotFoundError(f"Cannot read: {p}") filename = p.name else: img = image_input filename = "uploaded_image" if img.dtype != np.uint8: img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) if img.ndim == 3: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img = cv2.resize(img, (size, size), interpolation=cv2.INTER_LANCZOS4) clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) img = clahe.apply(img) img = cv2.GaussianBlur(img, (3, 3), 0) arr = img.astype(np.float32) / 255.0 return arr, filename # ── GradCAM ─────────────────────────────────────────────────────────────────── def compute_gradcam(img, model, device): model.eval() tensor = torch.tensor(img[None, None], requires_grad=False).to(device) activations = {} gradients = {} def fwd_hook(module, input, output): activations["down4"] = output.detach() def bwd_hook(module, grad_in, grad_out): gradients["down4"] = grad_out[0].detach() fwd_handle = model.down4.register_forward_hook(fwd_hook) bwd_handle = model.down4.register_full_backward_hook(bwd_hook) model.zero_grad() output = model(tensor) loss = output.mean() loss.backward() fwd_handle.remove() bwd_handle.remove() acts = activations["down4"].squeeze() grads = gradients["down4"].squeeze() weights = grads.mean(dim=(1, 2)) cam = (weights[:, None, None] * acts).sum(dim=0) cam = F.relu(cam) cam = cam.cpu().numpy() if cam.max() > 0: cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) cam_resized = cv2.resize(cam, (img.shape[1], img.shape[0])) heatmap = (cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8) base = (np.stack([img] * 3, axis=-1) * 255).astype(np.uint8) overlay = cv2.addWeighted(base, 0.5, heatmap, 0.5, 0) return overlay # ── Segmentation ────────────────────────────────────────────────────────────── def segment(img, model, device, threshold=0.5): tensor = torch.tensor(img[None, None]).to(device) with torch.no_grad(): prob = model(tensor).squeeze().cpu().numpy() return (prob > threshold).astype(np.uint8) # ── Instance detection ──────────────────────────────────────────────────────── def get_instances(mask, min_area=50, max_area=5000): kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) clean = cv2.morphologyEx( (mask * 255).astype(np.uint8), cv2.MORPH_OPEN, kernel) label_map = measure.label(clean, connectivity=2) props = measure.regionprops(label_map) filtered = np.zeros_like(label_map) new_id = 1 for p in props: if min_area <= p.area <= max_area: filtered[label_map == p.label] = new_id new_id += 1 return filtered # ── Feature extraction ──────────────────────────────────────────────────────── def extract_features(img, label_map): props = measure.regionprops(label_map, intensity_image=img) cells = [] for p in props: area = float(p.area) perimeter = float(p.perimeter) if p.perimeter > 0 else 1.0 circularity = 4 * np.pi * area / (perimeter ** 2 + 1e-6) region_px = img[label_map == p.label] cells.append({ "cell_id" : int(p.label), "area" : area, "perimeter" : perimeter, "circularity" : circularity, "eccentricity" : float(p.eccentricity), "solidity" : float(p.solidity), "mean_intensity": float(region_px.mean()), "std_intensity" : float(region_px.std()), }) return cells # ── Health classification + confidence ──────────────────────────────────────── FEATURE_COLS = ["area", "perimeter", "circularity", "eccentricity", "solidity", "mean_intensity", "std_intensity"] def classify_health(cells, clf_bundle): if not cells: return [], {}, [] clf = clf_bundle["clf"] le = clf_bundle["le"] X = np.array([[c[f] for f in FEATURE_COLS] for c in cells]) preds = le.inverse_transform(clf.predict(X)) proba = clf.predict_proba(X) confidences = proba.max(axis=1).tolist() health_summary = {"healthy": 0, "stressed": 0, "apoptotic": 0} for label in preds: health_summary[label] = health_summary.get(label, 0) + 1 return list(preds), health_summary, confidences # ── Overlay ─────────────────────────────────────────────────────────────────── def build_overlay(img, label_map, health_preds): COLOR = { "healthy" : (0, 200, 0), "stressed" : (0, 200, 200), "apoptotic": (0, 0, 220), } rgb = (np.stack([img] * 3, axis=-1) * 255).astype(np.uint8) for cell_id, label in enumerate(health_preds, start=1): color = COLOR.get(label, (128, 128, 128)) rgb[label_map == cell_id] = color return rgb # ── Benchmark ───────────────────────────────────────────────────────────────── def benchmark_status(value, metric): if metric == "confluency_pct": if value < 2: return "concerning" elif value < 5: return "below_normal" elif value <= 20: return "within_normal" else: return "above_normal" elif metric == "mean_circularity": if value >= 0.65: return "within_normal" elif value >= 0.40: return "below_normal" else: return "concerning" elif metric == "mean_solidity": if value >= 0.85: return "within_normal" else: return "below_normal" elif metric == "apoptotic_pct": if value <= 20: return "within_normal" elif value <= 30: return "above_normal" else: return "concerning" elif metric == "healthy_pct": if value >= 60: return "within_normal" elif value >= 40: return "below_normal" else: return "concerning" return "unknown" # ── PDF generation ──────────────────────────────────────────────────────────── def generate_pdf(report, overlay_path, gradcam_path, out_path): doc = SimpleDocTemplate(out_path, pagesize=A4, rightMargin=2*rcm, leftMargin=2*rcm, topMargin=2*rcm, bottomMargin=2*rcm) styles = getSampleStyleSheet() story = [] title_style = ParagraphStyle("title", fontSize=16, fontName="Helvetica-Bold", alignment=TA_CENTER, spaceAfter=12) story.append(Paragraph("Brightfield Cell Population Analysis Report", title_style)) story.append(Paragraph(f"File: {report.filename}", styles["Normal"])) story.append(Spacer(1, 0.4*rcm)) status_color = { "healthy_population" : colors.green, "mildly_suboptimal" : colors.orange, "suboptimal" : colors.darkorange, "stressed_or_abnormal": colors.red, }.get(report.overall_status, colors.grey) status_style = ParagraphStyle("status", fontSize=13, fontName="Helvetica-Bold", textColor=status_color, spaceAfter=8) story.append(Paragraph( f"Overall Status: {report.overall_status.replace('_', ' ').upper()}", status_style)) story.append(Spacer(1, 0.3*rcm)) story.append(Paragraph("Population Metrics", styles["Heading2"])) table_data = [ ["Metric", "Value", "Reference Range"], ["Total cells", str(report.n_cells), "50–300 cells/FOV"], ["Confluency", f"{report.confluency_pct:.1f}%", "5–20% (BBBC006)"], ["Mean cell area", f"{report.mean_area:.0f} px²", "100–3000 px²"], ["Mean circularity", f"{report.mean_circularity:.3f}", "≥ 0.65 (healthy)"], ["Mean solidity", f"{report.mean_solidity:.3f}", "≥ 0.85 (healthy)"], ["Mean intensity", f"{report.mean_intensity:.3f}", "0.25–0.65 a.u."], ["Healthy cells", f"{report.healthy_pct:.1f}%", "≥ 60%"], ["Stressed cells", f"{report.stressed_pct:.1f}%", "< 25%"], ["Apoptotic cells", f"{report.apoptotic_pct:.1f}%", "< 20%"], ["Mean confidence", f"{report.mean_confidence:.3f}", "0–1 (higher=better)"], ] t = Table(table_data, colWidths=[6*rcm, 4*rcm, 6*rcm]) t.setStyle(TableStyle([ ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#2d3748")), ("TEXTCOLOR", (0, 0), (-1, 0), colors.white), ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"), ("FONTSIZE", (0, 0), (-1, 0), 10), ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.HexColor("#f7fafc"), colors.white]), ("FONTSIZE", (0, 1), (-1, -1), 9), ("GRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#e2e8f0")), ("ALIGN", (1, 0), (-1, -1), "CENTER"), ("PADDING", (0, 0), (-1, -1), 6), ])) story.append(t) story.append(Spacer(1, 0.4*rcm)) story.append(Paragraph("Segmentation Overlay", styles["Heading2"])) if os.path.exists(overlay_path): story.append(RLImage(overlay_path, width=8*rcm, height=8*rcm)) story.append(Spacer(1, 0.3*rcm)) story.append(Paragraph("GradCAM Heatmap", styles["Heading2"])) story.append(Paragraph( "Regions highlighted in red/yellow indicate areas the U-Net " "focused on when predicting cell locations.", styles["Normal"])) if os.path.exists(gradcam_path): story.append(RLImage(gradcam_path, width=8*rcm, height=8*rcm)) story.append(Spacer(1, 0.3*rcm)) story.append(Paragraph("Biological Observations", styles["Heading2"])) for obs in report.observations: story.append(Paragraph(f"• {obs}", styles["Normal"])) story.append(Spacer(1, 0.3*rcm)) story.append(Paragraph("Recommendations", styles["Heading2"])) for rec in report.recommendations: story.append(Paragraph(f"• {rec}", styles["Normal"])) story.append(Spacer(1, 0.3*rcm)) story.append(Paragraph("Caveats", styles["Heading2"])) for cav in report.caveats: story.append(Paragraph(f"⚑ {cav}", styles["Normal"])) doc.build(story) # ── Report assembly ─────────────────────────────────────────────────────────── def build_report(filename, n_cells, confluency, cells, health_summary, confidences, img, label_map, health_preds): total = max(n_cells, 1) healthy_pct = round(100 * health_summary.get("healthy", 0) / total, 1) stressed_pct = round(100 * health_summary.get("stressed", 0) / total, 1) apoptotic_pct = round(100 * health_summary.get("apoptotic", 0) / total, 1) mean_conf = round(float(np.mean(confidences)), 4) if confidences else 0.0 mean_area = round(np.mean([c["area"] for c in cells]), 2) if cells else 0 mean_circ = round(np.mean([c["circularity"] for c in cells]), 4) if cells else 0 mean_sol = round(np.mean([c["solidity"] for c in cells]), 4) if cells else 0 mean_inten = round(np.mean([c["mean_intensity"] for c in cells]), 4) if cells else 0 metrics = { "confluency_pct" : confluency, "mean_circularity" : mean_circ, "mean_solidity" : mean_sol, "apoptotic_pct" : apoptotic_pct, "healthy_pct" : healthy_pct, } statuses = {m: benchmark_status(v, m) for m, v in metrics.items()} n_issues = sum(1 for s in statuses.values() if s in ["below_normal", "above_normal", "concerning"]) n_concern = sum(1 for s in statuses.values() if s == "concerning") if n_concern >= 2: overall = "stressed_or_abnormal" elif n_issues >= 3: overall = "suboptimal" elif n_issues >= 1: overall = "mildly_suboptimal" else: overall = "healthy_population" obs = [] if statuses["confluency_pct"] == "within_normal": obs.append(f"Confluency ({confluency:.1f}%) is within the normal range " f"for BBBC006 sparse plate format (5–20%).") elif statuses["confluency_pct"] == "below_normal": obs.append(f"Confluency ({confluency:.1f}%) is below normal (5–20%). " f"Possible low seeding density or impaired attachment.") else: obs.append(f"Confluency ({confluency:.1f}%) is very low — near background.") if statuses["mean_circularity"] == "within_normal": obs.append(f"Mean circularity ({mean_circ:.3f}) indicates well-rounded " f"healthy cell morphology (threshold ≥ 0.65).") else: obs.append(f"Mean circularity ({mean_circ:.3f}) is below healthy threshold " f"(0.65) — elongated or irregular cell shapes detected.") if statuses["apoptotic_pct"] == "within_normal": obs.append(f"Apoptotic fraction ({apoptotic_pct:.1f}%) is within " f"acceptable range (< 20%).") else: obs.append(f"Apoptotic fraction ({apoptotic_pct:.1f}%) is elevated. " f"Warrants further investigation.") if healthy_pct >= 60: obs.append(f"Majority of cells ({healthy_pct:.1f}%) show healthy morphology.") else: obs.append(f"Only {healthy_pct:.1f}% of cells show healthy morphology " f"— below the expected threshold of 60%.") obs.append(f"Mean classifier confidence: {mean_conf:.3f} " f"({'high' if mean_conf > 0.85 else 'moderate' if mean_conf > 0.70 else 'low'}).") recs = [] if statuses["confluency_pct"] in ["below_normal", "concerning"]: recs.append("Verify seeding density and allow additional time " "for cell attachment before imaging.") if statuses["apoptotic_pct"] != "within_normal": recs.append("Confirm apoptosis with Annexin V / PI staining.") recs.append("Check media freshness, CO₂ stability, " "and incubator temperature.") if statuses["mean_circularity"] != "within_normal": recs.append("Irregular morphology detected — check for cytoskeletal " "stress (osmolarity, pH, mechanical disruption).") if not recs: recs.append("Population metrics are within reference ranges. " "Standard monitoring schedule is appropriate.") recs.append("Single time-point analysis cannot distinguish growth " "inhibition from cytotoxicity. " "Parallel control imaging recommended.") caveats = [ "Analysis is based on morphological features only. " "Functional assays required for definitive conclusions.", "U-Net trained for 10 epochs — segmentation may under-detect " "cells. Retrain with 50 epochs for improved accuracy.", "Reference ranges adjusted for BBBC006 sparse plate format.", ] cell_details = [] for cell, label, conf in zip(cells, health_preds, confidences): cell_details.append({ "cell_id" : cell["cell_id"], "health" : label, "confidence" : round(conf, 4), "area" : round(cell["area"], 1), "circularity": round(cell["circularity"], 4), "solidity" : round(cell["solidity"], 4), }) overlay = build_overlay(img, label_map, health_preds) return InferenceReport( filename=filename, n_cells=n_cells, confluency_pct=round(confluency, 2), mean_area=mean_area, mean_circularity=mean_circ, mean_solidity=mean_sol, mean_intensity=mean_inten, healthy_pct=healthy_pct, stressed_pct=stressed_pct, apoptotic_pct=apoptotic_pct, mean_confidence=mean_conf, overall_status=overall, observations=obs, recommendations=recs, caveats=caveats, overlay_image=overlay, cell_details=cell_details, ) # ── Full pipeline ───────────────────────────────────────────────────────────── def run_inference(image_input, size=256, threshold=0.5, save_pdf=True, save_gradcam=True) -> InferenceReport: device = "cuda" if torch.cuda.is_available() else "cpu" img, filename = preprocess(image_input, size=size) unet = load_unet(device) mask = segment(img, unet, device, threshold=threshold) confluency = round(float(mask.mean() * 100), 2) gradcam_img = compute_gradcam(img, unet, device) gradcam_path = str(OUTPUT_DIR / "gradcam.png") if save_gradcam: cv2.imwrite(gradcam_path, cv2.cvtColor(gradcam_img, cv2.COLOR_RGB2BGR)) label_map = get_instances(mask) n_cells = int(label_map.max()) cells = extract_features(img, label_map) clf_bundle = load_classifier() health_preds, health_summary, confidences = classify_health( cells, clf_bundle) report = build_report(filename, n_cells, confluency, cells, health_summary, confidences, img, label_map, health_preds) report.gradcam_image = gradcam_img overlay_path = str(OUTPUT_DIR / "overlay.png") if report.overlay_image is not None: cv2.imwrite(overlay_path, cv2.cvtColor(report.overlay_image, cv2.COLOR_RGB2BGR)) if save_pdf: pdf_path = str(OUTPUT_DIR / "report.pdf") generate_pdf(report, overlay_path, gradcam_path, pdf_path) return report # ── Test run ────────────────────────────────────────────────────────────────── if __name__ == "__main__": test_img = sorted( (BASE_DIR / "data" / "processed").glob("*.npy"))[0] print(f"Testing on: {test_img.name}\n") report = run_inference(test_img) print(f"Cells : {report.n_cells}") print(f"Confluency : {report.confluency_pct}%") print(f"Healthy : {report.healthy_pct}%") print(f"Overall status : {report.overall_status}") print(f"Confidence : {report.mean_confidence}") print("\n✅ Inference complete.")