Upload model.py with huggingface_hub

model.py

@@ -1,16 +1,24 @@
 import torch
 import torch.nn as nn
 from torchvision import models
+import huggingface_hub
+from huggingface_hub import hf_hub_download
+import torchvision
 
+model_repo_id = "CSSE416-final-project/faceRecogModel"
+weight_file_id = "modelWeights100.bin"
 
-def load_model(pretrained_weights_path):
-    # Initialize Face-Rego
-    net = models.resnet18(pretrained=False)
-    num_ftrs = net.fc.in_features
-    net.fc = nn.Linear(num_ftrs, 128)  # Match your fine-tuned setup
+def load_model(repo_id):
+    # Download the model weights from the repo
+    weights_path = hf_hub_download(repo_id=model_repo_id, filename=weight_file_id)
 
-    # Load weights
-    state_dict = torch.load(pretrained_weights_path, map_location=torch.device('cpu'))
-    net.load_state_dict(state_dict)
-    net.eval()  # Set to evaluation mode
-    return net
+    # Initialize the ResNet-18 architecture
+    model = torchvision.models.resnet18(pretrained=True)
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your task (e.g., 128 classes)
+
+    # Load the model weights
+    state_dict = torch.load(weights_path, map_location=torch.device("cpu"))
+    model.load_state_dict(state_dict)
+    model.eval()  # Set the model to evaluation mode
+    return model