Spaces:
Sleeping
Sleeping
File size: 2,797 Bytes
57bcc48 6bb8cc3 57bcc48 7fea62b f547722 7fea62b 57bcc48 7fea62b f547722 7fea62b b2726e5 7fea62b b2726e5 57bcc48 b2726e5 57bcc48 b2726e5 57bcc48 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.models import vgg16, VGG16_Weights
from PIL import Image
import gradio as gr
import numpy as np
import cv2
from gradcam import GradCAM, apply_heatmap
import base64
from io import BytesIO
# Categories for classification
categories = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
# Load Pretrained VGG16 Model
print("Loading Model...")
device = torch.device("cpu")
vgg_model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
for param in vgg_model.features.parameters():
param.requires_grad = False
vgg_model.classifier[6] = nn.Linear(4096, 4)
vgg_model.load_state_dict(torch.load("vgg_xray_model.pth", map_location=device))
vgg_model = vgg_model.to(device)
vgg_model.eval()
print("Model Loaded Successfully!")
# Initialize Grad-CAM with the last convolutional layer
gradcam = GradCAM(vgg_model, vgg_model.features[-1])
# Image Preprocessing
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# Prediction function with Grad-CAM
def predict(image):
# Preprocess image (add batch dimension)
image_tensor = transform(image).unsqueeze(0).to(device)
# Perform inference first
with torch.no_grad():
output = vgg_model(image_tensor)
_, predicted = torch.max(output, 1)
prediction = categories[predicted.item()]
# Determine status
status = "COVID" if prediction in ["COVID", "Lung_Opacity"] else "Non-COVID"
try:
# Generate Grad-CAM heatmap with gradients
with torch.set_grad_enabled(True):
image_tensor.requires_grad_()
heatmap = gradcam.generate_heatmap(image_tensor, predicted.item())
# Generate overlay with bounding boxes
overlay = apply_heatmap(
image,
heatmap,
min_contour_area=100,
threshold=0.4
)
except Exception as e:
print(f"Heatmap generation failed: {str(e)}")
# Fallback to original image
overlay = image.copy()
# Return both JSON and Image outputs
return (
{"prediction": prediction, "status": status}, # First output (JSON)
overlay # Second output (Image)
)
# Gradio Interface
print("Initializing Gradio Interface...")
interface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[gr.JSON(label="Diagnosis"), gr.Image(type="pil", label="Heatmap")],
title="X-Ray Classifier with Grad-CAM",
allow_flagging="never"
)
print("Launching Gradio App...")
interface.launch(server_name="0.0.0.0", server_port=7860)
print("Gradio App Running!") |