Update app.py

app.py

@@ -5,31 +5,31 @@ from torchvision import models, transforms
 from PIL import Image
 from huggingface_hub import hf_hub_download
 import torch.nn as nn
-import torch.optim as optim
 
-# LoadING the VGG16 model
+# Loading the ResNet50 model from your Hugging Face repository
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-vgg16_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_vgg", filename="vgg16_model.pth")
+resnet50_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_resnet", filename="resnet50_model.pth")
 
-vgg16 = models.vgg16(pretrained=True)
-for param in vgg16.parameters():
+# ResNet50 model
+resnet50 = models.resnet50(pretrained=True)
+for param in resnet50.parameters():
     param.requires_grad = False
 
-num_classes = 8 
-vgg16.classifier[6] = nn.Linear(vgg16.classifier[6].in_features, num_classes)
-vgg16 = vgg16.to(device)
+num_classes = 8
+resnet50.fc = nn.Linear(resnet50.fc.in_features, num_classes)
+resnet50 = resnet50.to(device)
 
 # Loading the saved state dict
-checkpoint = torch.load(vgg16_model_path, map_location=device)
-vgg16.load_state_dict(checkpoint['model_state_dict'])
-vgg16.eval()
+checkpoint = torch.load(resnet50_model_path, map_location=device)
+resnet50.load_state_dict(checkpoint['model_state_dict'])
+resnet50.eval()
 
 # Fine-tuned Stable Diffusion model from your Hugging Face repository
 model_id = "sk2003/room-styler"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe.to(device)
 
-# Prediction function for the VGG16 model
+# Prediction function for the ResNet50 model
 def predict_and_show(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -39,10 +39,10 @@ def predict_and_show(image):
     image_tensor = transform(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        outputs = vgg16(image_tensor)
+        outputs = resnet50(image_tensor)
         _, predicted = torch.max(outputs.data, 1)
 
-    class_names = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"]  
+    class_names = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"]
     predicted_label = class_names[predicted.item()]
 
     return predicted_label