Create app.py

app.py

@@ -0,0 +1,114 @@
+import gradio as gr
+import cv2
+import torch
+import torchvision
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+import os
+import io
+
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+
+# Class names and colors
+CLASS_NAMES = {1: 'Nipple', 2: 'Lump'}
+CLASS_COLORS = {1: 'white', 2: 'white'}
+
+def preprocess_image(image):
+    """Load and preprocess image for Faster R-CNN."""
+    # Convert PIL Image to numpy array
+    image = np.array(image)
+    
+    # Convert RGB to RGB (already in correct format from Gradio)
+    image = image.astype(np.float32) / 255.0  # Normalize to [0,1]
+    
+    # Normalize using ImageNet mean and std
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    image = (image - mean) / std
+    
+    return torch.tensor(image.transpose(2, 0, 1), dtype=torch.float32)
+
+def load_model(checkpoint_path, device):
+    """Load Faster R-CNN model with fine-tuned weights."""
+    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None)
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(CLASS_NAMES) + 1)
+    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
+    model.to(device).eval()
+    return model
+
+def predict(image, score_thresh=0.5):
+    """Run inference and return image with bounding boxes."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Load model
+    checkpoint_path = "lumps.pth"  # This will be downloaded from the model repo
+    model = load_model(checkpoint_path, device)
+    
+    # Preprocess image
+    image_tensor = preprocess_image(image)
+    
+    # Run inference
+    model.eval()
+    with torch.no_grad():
+        preds = model([image_tensor.to(device)])[0]
+    
+    boxes, labels, scores = preds['boxes'].cpu().numpy(), preds['labels'].cpu().numpy(), preds['scores'].cpu().numpy()
+    
+    # Filter based on confidence threshold
+    keep = scores >= score_thresh
+    boxes, labels, scores = boxes[keep], labels[keep], scores[keep]
+    
+    # Convert tensor back to image
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    image_np = image_tensor.cpu().permute(1, 2, 0).numpy() * std + mean
+    image_np = np.clip(image_np, 0, 1)
+    
+    # Create figure
+    fig, ax = plt.subplots(1, figsize=(12, 12))
+    ax.imshow(image_np)
+    
+    # Draw bounding boxes
+    for box, label, score in zip(boxes, labels, scores):
+        xmin, ymin, xmax, ymax = box
+        rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, 
+                                linewidth=3, edgecolor=CLASS_COLORS.get(label, 'blue'), 
+                                facecolor='none')
+        ax.add_patch(rect)
+        ax.text(xmin, ymin - 10, f"{CLASS_NAMES.get(label, f'class_{label}')} ({score:.2f})", 
+                fontsize=14, color='white', backgroundcolor='black', weight='bold')
+    
+    plt.axis('off')
+    
+    # Convert plot to image
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    buf.seek(0)
+    result_image = Image.open(buf)
+    plt.close()
+    
+    return result_image
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Upload Breast Image"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence Threshold")
+    ],
+    outputs=gr.Image(type="pil", label="Detection Results"),
+    title="Breast Lumps Detection",
+    description="""Upload a breast image to detect lumps and nipples using a Faster R-CNN model.\n\n
+    ⚠️ **Important Medical Disclaimer**: This is a screening tool for research and assistive purposes only. 
+    It should NOT be used as the sole basis for medical diagnosis. All detections must be reviewed and confirmed 
+    by qualified medical professionals. This model is not FDA approved or certified for clinical diagnosis.""",
+    examples=None,
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    demo.launch()