Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
-from model import FineTunedResNet
 import os
 import time
 
@@ -15,6 +14,31 @@ transform = transforms.Compose([
 ])
 
 # Load the trained ResNet50 model
+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
+
+        # 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.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 512),  # Second additional layer
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, 256),  # Third additional layer
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, num_classes)  # Output layer
+        )
+
+    def forward(self, x):
+        return self.resnet(x)
+
 model = FineTunedResNet(num_classes=4)
 model_path = 'models/final_fine_tuned_resnet50.pth'
 
@@ -89,30 +113,30 @@ custom_css = """
         color: #333333;
         font-weight: bold;
         font-size: 24px;
-        margin-bottom: 10px;
+        margin-bottom: 10px.
     }
     .gradio-description {
         color: #666666;
-        font-size: 16px;
-        margin-bottom: 20px;
+        font-size: 16px.
+        margin-bottom: 20px.
     }
     .gradio-image {
-        border-radius: 10px;
+        border-radius: 10px.
     }
     .gradio-button {
-        background-color: #007bff;
-        color: #ffffff;
-        border: none;
-        padding: 10px 20px;
-        border-radius: 5px;
-        cursor: pointer;
+        background-color: #007bff.
+        color: #ffffff.
+        border: none.
+        padding: 10px 20px.
+        border-radius: 5px.
+        cursor: pointer.
     }
     .gradio-button:hover {
-        background-color: #0056b3;
+        background-color: #0056b3.
     }
     .gradio-label {
-        color: #007bff;
-        font-weight: bold;
+        color: #007bff.
+        font-weight: bold.
     }
 """