Update model/predict.py

model/predict.py

@@ -1,12 +1,16 @@
 import torch
-from torchvision import transforms
+from torchvision import models, transforms
 from PIL import Image
 
-MODEL_PATH = "models/deeptrust.pt"
+MODEL_PATH = "models/deeptrust_weights.pt"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the local model
-model = torch.load(MODEL_PATH, map_location=device)
+# Recreate the same model architecture
+model = models.efficientnet_b0(pretrained=False)
+model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)
+
+# Load the weights only (state_dict)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
 model.eval()
 
 # Transform for input images
@@ -40,3 +44,4 @@ def predict(image: Image.Image):
     trust_score = int(conf.item()*100 if pred.item() == 0 else (1 - conf.item())*100)
     
     return pred.item(), conf.item(), trust_score, tensor, model
+