Update app.py

app.py

@@ -5,9 +5,20 @@ 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
-model = models.resnet18()
-model.fc = torch.nn.Linear(model.fc.in_features, 2) 
+import torch.nn as nn
+
+# 1. Use resnet50 instead of resnet18 to match your weights
+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) 
+
+# 3. Load the weights
+model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
+model.eval()
 model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
 model.eval()