Upload chest_x_ray.py

chest_x_ray.py

@@ -0,0 +1,44 @@
+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()