chest_x_ray.py

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
import warnings

warnings.filterwarnings("ignore")

# Define model
class ResNet_Classifier(nn.Module):
    def __init__(self, num_classes=4):
        super(ResNet_Classifier, self).__init__()
        self.cnn = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=False)
        self.cnn.fc = nn.Linear(self.cnn.fc.in_features, num_classes)

    def forward(self, x):
        return self.cnn(x)

# Load model
model_path = "chest_x_ray.bin"  # Adjust path as needed
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = ResNet_Classifier(num_classes=4).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

# Image preprocessing
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

def predict(image):
    image = transform(image).unsqueeze(0).to(device)  # Add batch dimension
    with torch.inference_mode():
        output = model(image)
    
    # Class mapping
    labels = {0: "COVID19", 1: "NORMAL", 2: "PNEUMONIA", 3: "TUBERCULOSIS"}
    return labels[int(output.argmax(dim=1)[0])]

# Launch Gradio app
app = gr.Interface(fn=predict, inputs="image", outputs="label")
app.launch()