import os os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" os.environ["OMP_NUM_THREADS"] = "1" import sys import torch torch.set_num_threads(1) import torch.nn.functional as F import numpy as np try: import gradio as gr HAS_GRADIO = True except ImportError: gr = None HAS_GRADIO = False from PIL import Image, ImageDraw, ImageFont from pathlib import Path os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Add repository root and model subdirectories to path WORKSPACE_ROOT = Path(__file__).resolve().parent.parent if str(WORKSPACE_ROOT) not in sys.path: sys.path.insert(0, str(WORKSPACE_ROOT)) m2_path = WORKSPACE_ROOT / "models_suite/model2_choroidalyzer" if str(m2_path) not in sys.path: sys.path.insert(0, str(m2_path)) # Import models from models_suite from models_suite.model1_oct5k_layers.unet_layers import RetinalLayersUNet from models_suite.model2_choroidalyzer.choroidalyze.model import UNet as ChoroidalyzerUNet from models_suite.model3_hrf_dme.hrf_aunet import HRFAttentionUNet as HRF_AttentionUNet from models_suite.model4_oimhs_hole_cysts.oimhs_unet import OIMHSUNet from models_suite.model5_oct5k_detection.detector import OCTPathologyDetector, OCT5K_DETECTION_CLASSES try: import spaces IS_HF_SPACE = True except ImportError: spaces = None IS_HF_SPACE = False # Device configuration: # Local runs: use Mac GPU (mps/cuda) if available # Deployed HF Space: ConvNeXtV2 uses @spaces.GPU for ZeroGPU, Segmentation models use CPU only if IS_HF_SPACE: SEGMENT_DEVICE = torch.device("cpu") elif torch.cuda.is_available(): SEGMENT_DEVICE = torch.device("cuda") elif torch.backends.mps.is_available(): SEGMENT_DEVICE = torch.device("mps") else: SEGMENT_DEVICE = torch.device("cpu") DEVICE = SEGMENT_DEVICE print(f"Initializing OCT Analyzer Suite | Deployed HF Space: {IS_HF_SPACE} | Segmentation Device: {SEGMENT_DEVICE}") try: import numpy._core.multiarray as np_multiarray torch.serialization.add_safe_globals([np_multiarray.scalar]) except Exception: pass def safe_load_ckpt(cp_path, device): try: return torch.load(cp_path, map_location=device, weights_only=True) except Exception: return torch.load(cp_path, map_location=device, weights_only=False) def get_checkpoint_file(local_rel_path: str, bucket_key: str) -> Path: target_path = WORKSPACE_ROOT / local_rel_path if target_path.exists(): return target_path target_path.parent.mkdir(parents=True, exist_ok=True) token = os.getenv("HF_TOKEN") try: from huggingface_hub import hf_hub_download print(f"Downloading weight from HF Bucket: segmentation/{bucket_key}...", flush=True) cached_path = hf_hub_download( repo_id="NMundhra/OCT-Image-Classifier-Model-storage", filename=f"segmentation/{bucket_key}", repo_type="dataset", token=token ) import shutil shutil.copy(cached_path, target_path) print(f"Successfully cached {bucket_key} to {target_path}", flush=True) except Exception as e: print(f"Warning: Could not download segmentation/{bucket_key} from bucket: {e}", flush=True) return target_path # Initialize model instances and load checkpoints def load_suite(): print("Loading M1...", flush=True) m1 = RetinalLayersUNet(in_channels=1, num_classes=6) cp1 = get_checkpoint_file("models_suite/model1_oct5k_layers/checkpoints/best_model.pth", "model1_oct5k_layers.pth") if cp1.exists(): ckpt = safe_load_ckpt(cp1, DEVICE) m1.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt) m1.to(DEVICE).eval() print("M1 Done.", flush=True) print("Loading M2...", flush=True) m2 = ChoroidalyzerUNet(in_channels=1, out_channels=3, depth=7, channels='8_doublemax-64', up_type='conv_then_interpolate', extra_out_conv=True) cp2 = get_checkpoint_file("models_suite/model2_choroidalyzer/checkpoints/best_model.pth", "model2_choroidalyzer.pth") if cp2.exists(): ckpt = safe_load_ckpt(cp2, DEVICE) m2.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt) m2.to(DEVICE).eval() print("M2 Done.", flush=True) print("Loading M3...", flush=True) m3 = HRF_AttentionUNet(n_channels=3, n_classes=1) cp3 = get_checkpoint_file("models_suite/model3_hrf_dme/checkpoints/best_model.pth", "model3_hrf_dme.pth") if cp3.exists(): ckpt = safe_load_ckpt(cp3, DEVICE) m3.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt) m3.to(DEVICE).eval() print("M3 Done.", flush=True) print("Loading M4...", flush=True) m4 = OIMHSUNet(in_channels=1, num_classes=5) cp4 = get_checkpoint_file("models_suite/model4_oimhs_hole_cysts/checkpoints/best_model.pth", "model4_oimhs_hole_cysts.pth") if cp4.exists(): ckpt = safe_load_ckpt(cp4, DEVICE) m4.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt) m4.to(DEVICE).eval() print("M4 Done.", flush=True) print("Loading M5...", flush=True) m5 = OCTPathologyDetector(num_classes=10) cp5 = get_checkpoint_file("models_suite/model5_oct5k_detection/checkpoints/best_model.pth", "model5_oct5k_detection.pth") if cp5.exists(): ckpt = safe_load_ckpt(cp5, DEVICE) m5.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt) m5.to(DEVICE).eval() print("M5 Done.", flush=True) return m1, m2, m3, m4, m5 model1, model2, model3, model4, model5 = load_suite() print("OCT Analyser 5-Model Suite Loaded Successfully!", flush=True) # Preprocessing helpers def preprocess_image(image: Image.Image, target_size=(256, 256)): gray = image.convert("L").resize(target_size) arr = np.array(gray, dtype=np.float32) / 255.0 tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0) # (1, 1, H, W) return gray, tensor def create_pure_mask(orig_img: Image.Image, mask: np.ndarray, num_classes: int, cmap: list) -> Image.Image: h, w = mask.shape mask_rgba = np.zeros((h, w, 4), dtype=np.uint8) for c in range(1, num_classes): color = cmap[c % len(cmap)] mask_rgba[mask == c] = [color[0], color[1], color[2], 255] return Image.fromarray(mask_rgba, mode="RGBA").resize(orig_img.size, Image.NEAREST) def create_segmentation_overlay(orig_img: Image.Image, mask: np.ndarray, num_classes: int, cmap: list): orig_rgb = orig_img.convert("RGB") h, w = mask.shape overlay = np.zeros((h, w, 3), dtype=np.uint8) for c in range(1, num_classes): color = cmap[c % len(cmap)] overlay[mask == c] = color overlay_img = Image.fromarray(overlay).resize(orig_rgb.size) blended = Image.blend(orig_rgb, overlay_img, alpha=0.45) return blended, Image.fromarray(overlay).resize(orig_rgb.size) def preprocess_rgb_image(image: Image.Image, target_size=(256, 256)): rgb = image.convert("RGB").resize(target_size) arr = np.array(rgb, dtype=np.float32).transpose(2, 0, 1) / 255.0 # (3, H, W) tensor = torch.from_numpy(arr).unsqueeze(0) # (1, 3, H, W) return rgb, tensor COLOR_MAP = [ [0, 0, 0], # 0: BG [255, 50, 50], # 1: Red [50, 255, 50], # 2: Green [50, 50, 255], # 3: Blue [255, 255, 50], # 4: Yellow [255, 50, 255], # 5: Magenta [50, 255, 255] # 6: Cyan ] # Inference handlers for each model def predict_model1(image): if image is None: return None, "Please upload an OCT scan image." gray, tensor = preprocess_image(image, (256, 256)) with torch.no_grad(): logits = model1(tensor.to(DEVICE)) preds = torch.argmax(logits, dim=1).cpu().numpy()[0] overlay = create_segmentation_overlay(image, preds, 6, COLOR_MAP) layer_names = ["Background", "ILM -> OPL", "OPL -> IS-OS", "IS-OS -> IBRPE", "IBRPE -> OBRPE", "Choroid & Below"] counts = {layer_names[c]: int((preds == c).sum()) for c in range(6)} info = f"Retinal Layer Segmentation Complete.\nPixel Breakdown:\n" + "\n".join([f"• {k}: {v} px" for k, v in counts.items()]) return overlay, info def predict_model2(image): if image is None: return None, "Please upload an OCT scan image." gray, tensor = preprocess_image(image, (256, 256)) with torch.no_grad(): logits = model2(tensor.to(DEVICE)) probs = torch.sigmoid(logits).cpu().numpy()[0, 0] mask = (probs > 0.5).astype(np.uint8) overlay = create_segmentation_overlay(image, mask, 2, [[0,0,0], [0, 220, 255]]) choroid_area = int(mask.sum()) mean_thickness = float(mask.sum(axis=0).mean()) info = f"Choroid Region Analysis:\n• Choroid Area: {choroid_area} px\n• Est. Mean Thickness: {mean_thickness:.2f} px" return overlay, info def predict_model3(image): if image is None: return None, "Please upload an OCT scan image." rgb, tensor = preprocess_rgb_image(image, (256, 256)) with torch.no_grad(): logits = model3(tensor.to(DEVICE)) probs = torch.sigmoid(logits).cpu().numpy()[0, 0] mask = (probs > 0.5).astype(np.uint8) overlay = create_segmentation_overlay(image, mask, 2, [[0,0,0], [255, 0, 120]]) fluid_px = int(mask.sum()) info = f"High-Resolution HRF DME Fluid & Lesion Attention Analysis:\n• Pathological Fluid / Lesion Region: {fluid_px} px" return overlay, info def predict_model4(image): if image is None: return None, "Please upload an OCT scan image." gray, tensor = preprocess_image(image, (256, 256)) with torch.no_grad(): logits = model4(tensor.to(DEVICE)) preds = torch.argmax(logits, dim=1).cpu().numpy()[0] overlay = create_segmentation_overlay(image, preds, 5, COLOR_MAP) classes = ["Background", "Macular Hole", "Choroid", "Retina", "Intraretinal Cysts (IRC)"] counts = {classes[c]: int((preds == c).sum()) for c in range(5)} info = "OIMHS Pathology Analysis:\n" + "\n".join([f"• {k}: {v} px" for k, v in counts.items()]) return overlay, info def predict_model5(image, score_threshold=0.5): if image is None: return None, "Please upload an OCT scan image." orig_rgb = image.convert("RGB") gray, tensor = preprocess_image(image, (256, 256)) # Faster R-CNN expects 3-channel input input_tensor = tensor.repeat(1, 3, 1, 1).to(DEVICE) with torch.no_grad(): outputs = model5([input_tensor[0]])[0] boxes = outputs["boxes"].cpu().numpy() labels = outputs["labels"].cpu().numpy() scores = outputs["scores"].cpu().numpy() # Scale boxes back to original image size orig_w, orig_h = orig_rgb.size scale_x = orig_w / 256.0 scale_y = orig_h / 256.0 draw_img = orig_rgb.copy() draw = ImageDraw.Draw(draw_img) detected_items = [] for box, label, score in zip(boxes, labels, scores): if score >= score_threshold: x1, y1, x2, y2 = box[0] * scale_x, box[1] * scale_y, box[2] * scale_x, box[3] * scale_y cls_name = OCT5K_DETECTION_CLASSES[label] if label < len(OCT5K_DETECTION_CLASSES) else f"Class {label}" draw.rectangle([x1, y1, x2, y2], outline="red", width=3) draw.text((x1 + 4, max(0, y1 - 12)), f"{cls_name} {score:.2f}", fill="yellow") detected_items.append(f"• {cls_name}: Conf {score:.2%} at [{int(x1)}, {int(y1)}, {int(x2)}, {int(y2)}]") info = f"OCT Pathology Object Detector ({len(detected_items)} objects detected above threshold {score_threshold:.2f}):\n" if detected_items: info += "\n".join(detected_items) else: info += "No objects detected above threshold." return draw_img, info # Build Gradio UI for HF Space if HAS_GRADIO: with gr.Blocks(title="OCT Analyser 5-Model Microservice Suite") as demo: gr.Markdown("# 👁️ OCT Analyser 5-Model Suite (Hugging Face API Services Deployment)") gr.Markdown("Comprehensive Optical Coherence Tomography (OCT) Deep Learning Suite providing 5 API services for segmentation, choroid analysis, fluid/lesion quantification, macular hole/cyst detection, and 9-class biomarker object detection.") with gr.Tabs(): with gr.TabItem("Model 1: Retinal Layers U-Net"): gr.Markdown("### 6-Class Retinal Layer Segmentation U-Net (OCT5K Benchmark)") with gr.Row(): with gr.Column(): img1 = gr.Image(type="pil", label="Input OCT Scan") btn1 = gr.Button("Segment Retinal Layers", variant="primary") with gr.Column(): out_img1 = gr.Image(type="pil", label="6-Layer Segmentation Overlay") txt1 = gr.Textbox(label="Layer Metrics", lines=7) btn1.click(predict_model1, inputs=img1, outputs=[out_img1, txt1], api_name="predict_model1") with gr.TabItem("Model 2: Choroidalyzer U-Net"): gr.Markdown("### Choroid Region & Thickness Quantification U-Net") with gr.Row(): with gr.Column(): img2 = gr.Image(type="pil", label="Input OCT Scan") btn2 = gr.Button("Analyze Choroid Region", variant="primary") with gr.Column(): out_img2 = gr.Image(type="pil", label="Choroid Mask Overlay") txt2 = gr.Textbox(label="Choroid Metrics", lines=5) btn2.click(predict_model2, inputs=img2, outputs=[out_img2, txt2], api_name="predict_model2") with gr.TabItem("Model 3: HRF Attention U-Net"): gr.Markdown("### High-Resolution Fluid & Lesion Attention U-Net (HRF DME/AMD)") with gr.Row(): with gr.Column(): img3 = gr.Image(type="pil", label="Input OCT Scan") btn3 = gr.Button("Segment Fluid & Lesions", variant="primary") with gr.Column(): out_img3 = gr.Image(type="pil", label="Fluid & Lesion Mask Overlay") txt3 = gr.Textbox(label="Fluid & Lesion Metrics", lines=5) btn3.click(predict_model3, inputs=img3, outputs=[out_img3, txt3], api_name="predict_model3") with gr.TabItem("Model 4: OIMHS Hole & Cyst U-Net"): gr.Markdown("### Macular Hole & Intraretinal Cyst (IRC) U-Net (OIMHS)") with gr.Row(): with gr.Column(): img4 = gr.Image(type="pil", label="Input OCT Scan") btn4 = gr.Button("Detect Hole & Cysts", variant="primary") with gr.Column(): out_img4 = gr.Image(type="pil", label="Hole & Cyst Segmentation Overlay") txt4 = gr.Textbox(label="Pathology Metrics", lines=6) btn4.click(predict_model4, inputs=img4, outputs=[out_img4, txt4], api_name="predict_model4") with gr.TabItem("Model 5: OCT Pathology Detector"): gr.Markdown("### Faster R-CNN 9-Class Biomarker Object Detector") with gr.Row(): with gr.Column(): img5 = gr.Image(type="pil", label="Input OCT Scan") thresh = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.05, label="Confidence Threshold") btn5 = gr.Button("Detect Biomarker Bounding Boxes", variant="primary") with gr.Column(): out_img5 = gr.Image(type="pil", label="Biomarker Bounding Box Detections") txt5 = gr.Textbox(label="Detection Results", lines=8) btn5.click(predict_model5, inputs=[img5, thresh], outputs=[out_img5, txt5], api_name="predict_model5") if __name__ == "__main__": if HAS_GRADIO: demo.launch(server_name="0.0.0.0", server_port=7860) else: print("Gradio is not installed locally. Run pip install gradio to launch the server web interface.")