Update app.py

app.py

@@ -6,31 +6,31 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import torch.nn as nn
 
-# Loading the ResNet50 model from your Hugging Face repository
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-resnet50_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_resnet", filename="resnet50_model.pth")
 
-# ResNet50 model
-resnet50 = models.resnet50(pretrained=True)
-for param in resnet50.parameters():
-    param.requires_grad = False
+# Finetuned Resnet-50 model is downloaded
+vgg16_model_path = hf_hub_download(repo_id="sk2003/style_recognizer_vgg", filename="vgg16_model.pth")
+
+vgg16 = models.vgg16(pretrained=True)
+for param in vgg16.parameters():
+    param.requires_grad = False # freezing parameters
 
 num_classes = 8
-resnet50.fc = nn.Linear(resnet50.fc.in_features, num_classes)
-resnet50 = resnet50.to(device)
+vgg16.fc = nn.Linear(vgg16.fc.in_features, num_classes)
+vgg16 = vgg16.to(device)
 
-# Loading the saved state dict
-checkpoint = torch.load(resnet50_model_path, map_location=device)
-resnet50.load_state_dict(checkpoint['model_state_dict'])
-resnet50.eval()
+# Loading the model
+checkpoint = torch.load(vgg16_model_path, map_location=device)
+vgg16.load_state_dict(checkpoint['model_state_dict'])
+vgg16.eval() # setting to evaluation mode to disable batch-norm and dropout layers
 
-# Fine-tuned Stable Diffusion model from your Hugging Face repository
+# Fine-tuned Stable Diffusion model
 model_id = "sk2003/room-styler"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe.to(device)
 
 # Prediction function for the ResNet50 model
-def predict_and_show(image):
+def predict(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -42,10 +42,10 @@ def predict_and_show(image):
         outputs = resnet50(image_tensor)
         _, predicted = torch.max(outputs.data, 1)
 
-    class_names = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"]
-    predicted_label = class_names[predicted.item()]
+    classes = ["Classic", "Modern", "Vintage", "Glamour", "Scandinavian", "Rustic", "ArtDeco", "Industrial"]
+    pred = classes[predicted.item()]
 
-    return predicted_label
+    return pred
 
 # Generation function for the Stable Diffusion model
 def generate_image(prompt):
@@ -54,14 +54,16 @@ def generate_image(prompt):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Room Style Recognition and Generation")
+    gr.Markdown("## Room Style Recognition and Generation") # title
 
+   # 1st tab
     with gr.Tab("Recognize Room Style"):
         image_input = gr.Image(type="pil")
         label_output = gr.Textbox()
         btn_predict = gr.Button("Predict Style")
-        btn_predict.click(predict_and_show, inputs=image_input, outputs=label_output)
+        btn_predict.click(predict, inputs=image_input, outputs=label_output)
 
+    # 2nd tab
     with gr.Tab("Generate Room Style"):
         text_input = gr.Textbox(placeholder="Enter a prompt for room style...")
         image_output = gr.Image()
@@ -69,3 +71,4 @@ with gr.Blocks() as demo:
         btn_generate.click(generate_image, inputs=text_input, outputs=image_output)
 
 demo.launch()
+