Update app.py

app.py

@@ -6,19 +6,22 @@ from PIL import Image
 # 1. Load your model (Ensure this matches your training architecture)
 # Change 'models.resnet18' if you used a different one
 # --- UPDATED MODEL ARCHITECTURE ---
-from torchvision import models
 import torch.nn as nn
+from torchvision import models
 
-# 1. Use resnet50 instead of resnet18 to match your weights
-model = models.resnet50() 
+# 1. Initialize ResNet-50 (matches the 2048 feature size in your error)
+model = models.resnet50()
 
-# 2. Update the final layer to match your 2 classes (Defect, Normal)
-# Note: ResNet50 uses 2048 features in the final layer
-model.fc = nn.Linear(2048, 2) 
+# 2. Recreate the EXACT Sequential head used during your training
+# This fixes the "Missing key: fc.0.weight" and "fc.3.weight" errors
+model.fc = nn.Sequential(
+    nn.Linear(2048, 256),   # fc.0
+    nn.ReLU(),              # fc.1
+    nn.Dropout(0.4),        # fc.2
+    nn.Linear(256, 2)       # fc.3
+)
 
-# 3. Load the weights
-model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
-model.eval()
+# 3. Load your weights
 model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
 model.eval()