uploaded the inference endpoint .py file

airad.py

@@ -0,0 +1,86 @@
+import torch
+from torch import nn
+from torchvision import models, transforms
+from PIL import Image
+import cv2
+import numpy as np
+import gdown
+
+class AIRadModel(nn.Module):
+  def __init__(self,num_classes=2):
+    super(AIRadModel,self).__init__()
+    self.model = models.efficientnet_b0(pretrained=False)
+    self.num_features = model.classifier[1].in_features
+    self.model.classifier = nn.Sequential(
+        nn.Dropout(p=0.2),
+        nn.Linear(self.num_features, num_classes)  # Two classes: normal, pneumonia
+    )
+
+  def forward(self, x):
+        return self.model(x)
+
+class AIRadSimModel(nn.Module):
+  def __init__(self, num_classes=2):
+    super(AIRadSimModel,self).__init__()
+    self.sim_model = models.resnet50(pretrained=False)
+    self.sim_model.fc = nn.Linear(self.sim_model.fc.in_features,num_classes)
+
+  def forward(self,x):
+    return self.sim_model(x)
+
+
+def load_model():
+  model = AIRadModel(num_classes=2)
+  file_id = '1CKkdQ5nKWkz3L-ZdgyrJ5SE-oiFwXnSJ'
+  gdrive_url = f"https://drive.google.com/uc?id={file_id}"
+  model_checkpoint = 'model_checkpoint.pth'
+  gdown.download(gdrive_url, model_checkpoint, quiet=False)
+  model.load_state_dict(torch.load(model_checkpoint))
+  model.eval()
+  return model
+
+def load_sim_model():
+  sim_model = AIRadSimModel(num_classes=2)
+  sim_file_id = 'cjdDsW5QAIlOneOPLg0uYqTURSr0oOLq'
+  sim_gdrive_url = f"https://drive.google.com/uc?id={file_id}"
+  sim_model_checkpoint = 'sim_model_checkpoint.pth'
+  gdown.download(sim_gdrive_url, sim_model_checkpoint, quiet=False)
+  sim_model.load_state_dict(torch.load(sim_model_checkpoint))
+  sim_model.eval()
+  return sim_model()
+
+model = load_model()
+sim_model = load_sim_model()
+class_names = {0: 'normal', 1: 'pneumonia'}
+
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+def predict(image_path):
+  image = Image.open(image_path).convert("RGB")
+  image_np = np.array(image)
+  image_np = cv2.bilateralFilter(image_np, 9, 75, 75)
+  image = Image.fromarray(image_np)
+  image_tensor = preprocess(image).unsqueeze(0).to(device)
+
+  # Use ResNet50 to predict if the image is an X-ray
+  with torch.no_grad():
+    sim_output = sim_model(image_tensor)
+    _, predicted_sim = torch.max(sim_output, 1)
+    predicted_class_sim = predicted_sim.item()
+
+  if predicted_class_sim == 1:  
+    with torch.no_grad():
+      output = model(image_tensor)
+      _, predicted = torch.max(output, 1)
+      predicted_class = predicted.item()
+      confidence = torch.nn.functional.softmax(output, dim=1)[0][predicted_class].item()
+      return class_names[predicted_class] ,confidence
+
+  else:
+    return "error"
+