Update app.py

app.py

@@ -9,6 +9,23 @@ import os
 # Define the number of classes
 num_classes = 2
 
+# Download model from Hugging Face
+def download_model():
+    model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
+    return model_path
+
+# Load the model from Hugging Face
+def load_model(model_path):
+    model = models.resnet50(pretrained=False)
+    model.fc = nn.Linear(model.fc.in_features, num_classes)
+    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+    model.eval()
+    return model
+
+# Download the model and load it
+model_path = download_model()
+model = load_model(model_path)
+
 # Define transformation for image processing
 transform = transforms.Compose([
     transforms.Resize(256),