update 2

app.py

@@ -6,15 +6,8 @@ import requests
 from PIL import Image
 from torchvision import transforms
 import urllib.request
-import os
-import torch
+from torchvision import models
 import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import DataLoader, Dataset, DistributedSampler
-from transformers import AutoModel, AutoTokenizer
-from torchvision import models, transforms
-from torch.cuda.amp import GradScaler, autocast
-import numpy as np
 
 # --- Define the Model ---
 class FineGrainedClassifier(nn.Module):
@@ -43,7 +36,6 @@ class FineGrainedClassifier(nn.Module):
         output = self.classifier(combined_features)
         return output
 
-
 # --- Data Augmentation Setup ---
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -54,60 +46,43 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# def load_model_checkpoint(model, checkpoint_path, device):
-#     checkpoint = torch.load(checkpoint_path, map_location=device)
-    
-#     # Strip the "module." prefix from the keys in the state_dict if they exist
-#     state_dict = checkpoint['model_state_dict']
-#     new_state_dict = {}
-    
-#     for k, v in state_dict.items():
-#         if k.startswith("module."):
-#             new_state_dict[k[7:]] = v  # Remove "module." prefix
-#         else:
-#             new_state_dict[k] = v
-    
-#     model.load_state_dict(new_state_dict)
-#     return model
-
 # Load the label-to-class mapping from your Hugging Face repository
 label_map_url = "https://huggingface.co/Maverick98/EcommerceClassifier/resolve/main/label_to_class.json"
 label_to_class = requests.get(label_map_url).json()
 
-
 # Load your custom model from Hugging Face
 model = FineGrainedClassifier(num_classes=len(label_to_class))
-model_checkpoint = "Maverick98/EcommerceClassifier"
 checkpoint_url = f"https://huggingface.co/Maverick98/EcommerceClassifier/resolve/main/model_checkpoint.pth"
 checkpoint = torch.hub.load_state_dict_from_url(checkpoint_url, map_location=torch.device('cpu'))
-# Extract and load the model state_dict
-model.load_state_dict(checkpoint['model_state_dict'])
+
+# Strip the "module." prefix from the keys in the state_dict if they exist
+new_state_dict = {}
+for k, v in checkpoint.items():
+    if k.startswith("module."):
+        new_state_dict[k[7:]] = v  # Remove "module." prefix
+    else:
+        new_state_dict[k] = v
+
+model.load_state_dict(new_state_dict)
 
 # Load the tokenizer from Jina
 tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-base-en")
 
-
-def load_image(image_path_or_url):
+def load_image(image):
     """
-    Load an image from a URL or local path and preprocess it.
+    Preprocess the uploaded image.
     """
-    if image_path_or_url.startswith("http"):
-        with urllib.request.urlopen(image_path_or_url) as url:
-            image = Image.open(url).convert('RGB')
-    else:
-        image = Image.open(image_path_or_url).convert('RGB')
-    
     image = transform(image)
     image = image.unsqueeze(0)  # Add batch dimension
     return image
 
-def predict(image_path_or_url, title, threshold=0.7):
+def predict(image, title, threshold=0.7):
     """
     Predict the top 3 categories for the given image and title.
     Includes "Others" if the confidence of the top prediction is below the threshold.
     """
     # Preprocess the image
-    image = load_image(image_path_or_url)
+    image = load_image(image)
     
     # Tokenize the title
     title_encoding = tokenizer(title, padding='max_length', max_length=200, truncation=True, return_tensors='pt')
@@ -138,7 +113,7 @@ def predict(image_path_or_url, title, threshold=0.7):
 
 # Define the Gradio interface
 title_input = gr.inputs.Textbox(label="Product Title", placeholder="Enter the product title here...")
-image_input = gr.inputs.Textbox(label="Image URL or Path", placeholder="Enter image URL or local path here...")
+image_input = gr.inputs.Image(type="pil", label="Upload Image")
 output = gr.outputs.JSON(label="Top 3 Predictions with Probabilities")
 
 gr.Interface(