Update app.py

app.py

@@ -7,84 +7,105 @@ from PIL import Image
 import os
 import time
 
-# Define the transform for the input image
+# =========================
+# Image preprocessing
+# =========================
 transform = transforms.Compose([
-    transforms.Resize((150, 150)),
+    transforms.Resize((224, 224)),  # Required for ResNet50
     transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
 ])
 
-# Load the trained ResNet50 model
+# =========================
+# Model definition
+# =========================
 class FineTunedResNet(nn.Module):
     def __init__(self, num_classes=4):
-        super(FineTunedResNet, self).__init__()
-        self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)  # Load pre-trained ResNet50
+        super().__init__()
+        self.resnet = models.resnet50(
+            weights=models.ResNet50_Weights.DEFAULT
+        )
 
-        # Replace the fully connected layer with more layers and batch normalization
         self.resnet.fc = nn.Sequential(
-            nn.Linear(self.resnet.fc.in_features, 1024),  # First additional layer
+            nn.Linear(self.resnet.fc.in_features, 1024),
             nn.BatchNorm1d(1024),
             nn.ReLU(),
             nn.Dropout(0.5),
-            nn.Linear(1024, 512),  # Second additional layer
+
+            nn.Linear(1024, 512),
             nn.BatchNorm1d(512),
             nn.ReLU(),
             nn.Dropout(0.5),
-            nn.Linear(512, 256),  # Third additional layer
+
+            nn.Linear(512, 256),
             nn.BatchNorm1d(256),
             nn.ReLU(),
             nn.Dropout(0.5),
-            nn.Linear(256, num_classes)  # Output layer
+
+            nn.Linear(256, num_classes)
         )
 
     def forward(self, x):
         return self.resnet(x)
 
-model = FineTunedResNet(num_classes=4)
-model_path = 'models/final_fine_tuned_resnet50.pth'
+# =========================
+# Load model
+# =========================
+MODEL_PATH = "models/final_fine_tuned_resnet50.pth"
 
-if not os.path.exists(model_path):
-    raise FileNotFoundError(f"The model file '{model_path}' does not exist. Please check the path.")
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(f"Model not found: {MODEL_PATH}")
 
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+model = FineTunedResNet(num_classes=4)
+model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
 model.eval()
+model.to("cpu")
+
+CLASSES = ["🦠 COVID", "🫁 Normal", "🦠 Pneumonia", "🦠 TB"]
+
+# =========================
+# Prediction function
+# =========================
+def predict(image: Image.Image) -> str:
+    start = time.time()
+
+    image = transform(image).unsqueeze(0)
 
-# Define a function to make predictions
-def predict(image):
-    start_time = time.time()  # Start the timer
-    image = transform(image).unsqueeze(0)  # Transform and add batch dimension
-    
     with torch.no_grad():
         output = model(image)
-        probabilities = F.softmax(output, dim=1)[0]
-        top_prob, top_class = torch.topk(probabilities, 3)
-        classes = ['🦠 COVID', '🫁 Normal', '🦠 Pneumonia', '🦠 TB']  # Adjust based on the classes in your model
-        
-    end_time = time.time()  # End the timer
-    prediction_time = end_time - start_time  # Calculate the prediction time
-
-    # Format the result string
-    result = f"Top Predictions:\n"
-    for i in range(top_prob.size(0)):
-        result += f"{classes[top_class[i]]}: Score {top_prob[i].item()}\n"
-    result += f"Prediction Time: {prediction_time:.2f} seconds"
-    
+        probs = F.softmax(output, dim=1)[0]
+        top_probs, top_idxs = torch.topk(probs, 3)
+
+    elapsed = time.time() - start
+
+    result = "Top Predictions:\n\n"
+    for prob, idx in zip(top_probs, top_idxs):
+        result += f"{CLASSES[idx]} → {prob.item():.4f}\n"
+
+    result += f"\n⏱️ Prediction Time: {elapsed:.2f} seconds"
     return result
 
-# Example images with labels
+# =========================
+# Example images
+# =========================
 examples = [
-    ['examples/Pneumonia/02009view1_frontal.jpg', '🦠 Pneumonia'],
-    ['examples/Pneumonia/02055view1_frontal.jpg', '🦠 Pneumonia'],
-    ['examples/Pneumonia/03152view1_frontal.jpg', '🦠 Pneumonia'],
-    ['examples/COVID/11547_2020_1200_Fig3_HTML-a.png', '🦠 COVID'],
-    ['examples/COVID/11547_2020_1200_Fig3_HTML-b.png', '🦠 COVID'],
-    ['examples/COVID/11547_2020_1203_Fig1_HTML-b.png', '🦠 COVID'],
-    ['examples/Normal/06bc1cfe-23a0-43a4-a01b-dfa10314bbb0.jpg', '🫁 Normal'],
-    ['examples/Normal/08ae6c0b-d044-4de2-a410-b3cf8dc65868.jpg', '🫁 Normal'],
-    ['examples/Normal/IM-0178-0001.jpeg', '🫁 Normal']
+    ["examples/Pneumonia/02009view1_frontal.jpg"],
+    ["examples/Pneumonia/02055view1_frontal.jpg"],
+    ["examples/Pneumonia/03152view1_frontal.jpg"],
+    ["examples/COVID/11547_2020_1200_Fig3_HTML-a.png"],
+    ["examples/COVID/11547_2020_1200_Fig3_HTML-b.png"],
+    ["examples/COVID/11547_2020_1203_Fig1_HTML-b.png"],
+    ["examples/Normal/06bc1cfe-23a0-43a4-a01b-dfa10314bbb0.jpg"],
+    ["examples/Normal/08ae6c0b-d044-4de2-a410-b3cf8dc65868.jpg"],
+    ["examples/Normal/IM-0178-0001.jpeg"]
 ]
 
-# Load visualization images
+# =========================
+# Visualization images
+# =========================
 visualization_images = [
     "pictures/1.png",
     "pictures/2.png",
@@ -93,86 +114,41 @@ visualization_images = [
     "pictures/5.png"
 ]
 
-# Function to display visualization images
 def display_visualizations():
-    return [Image.open(image) for image in visualization_images]
-
-# Custom CSS to enhance appearance (injected via HTML)
-custom_css = """
-    <style>
-    body {
-        font-family: 'Arial', sans-serif;
-        background-color: #f5f5f5;
-    }
-    .gradio-container {
-        background-color: #ffffff;
-        border: 1px solid #e6e6e6;
-        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-        border-radius: 10px;
-        padding: 20px;
-    }
-    .gradio-title {
-        color: #333333;
-        font-weight: bold;
-        font-size: 24px;
-        margin-bottom: 10px;
-    }
-    .gradio-description {
-        color: #666666;
-        font-size: 16px;
-        margin-bottom: 20px;
-    }
-    .gradio-image {
-        border-radius: 10px;
-    }
-    .gradio-button {
-        background-color: #007bff;
-        color: #ffffff;
-        border: none;
-        padding: 10px 20px;
-        border-radius: 5px;
-        cursor: pointer;
-    }
-    .gradio-button:hover {
-        background-color: #0056b3;
-    }
-    .gradio-label {
-        color: #007bff;
-        font-weight: bold;
-    }
-    </style>
-"""
+    return [Image.open(path) for path in visualization_images]
 
-# Create Gradio interfaces
+# =========================
+# Gradio interfaces
+# =========================
 prediction_interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Chest X-ray Image"),
-    outputs=gr.Label(label="Predicted Disease"),
+    inputs=gr.Image(type="pil", label="Upload Chest X-ray"),
+    outputs=gr.Textbox(label="Prediction Result"),
     examples=examples,
+    cache_examples=False,  # IMPORTANT for HF Spaces
     title="Lung Disease Detection XVI",
-    description=f"""
-        Upload a chest X-ray image to detect lung diseases such as 🦠 COVID-19, 🦠 Pneumonia, 🫁 Normal, or 🦠 TB. 
-        Use the example images to see how the model works.
-        {custom_css}
-    """
+    description="""
+Upload a chest X-ray image to detect:
+🦠 COVID-19 • 🦠 Pneumonia • 🫁 Normal • 🦠 Tuberculosis
+"""
 )
 
 visualization_interface = gr.Interface(
     fn=display_visualizations,
     inputs=None,
-    outputs=[gr.Image(type="pil", label=f"Visualization {i+1}") for i in range(len(visualization_images))],
-    title="Model Performance Visualizations",
-    description=f"""
-        Here are some visualizations that depict the performance of the model during training and testing.
-        {custom_css}
-    """
+    outputs=[
+        gr.Image(type="pil", label=f"Visualization {i+1}")
+        for i in range(len(visualization_images))
+    ],
+    title="Model Performance Visualizations"
 )
 
-# Combine interfaces into a tabbed interface
 app = gr.TabbedInterface(
     interface_list=[prediction_interface, visualization_interface],
     tab_names=["Predict", "Model Performance"]
 )
 
-# Launch the interface
-app.launch(share=True)
+# =========================
+# Launch (HF Spaces safe)
+# =========================
+app.launch()