Update app.py

app.py

@@ -6,22 +6,95 @@ import requests
 from PIL import Image
 from torchvision import transforms
 import urllib.request
+import os
+import torch
+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 sklearn.model_selection import train_test_split
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.metrics import precision_recall_fscore_support, accuracy_score
+from torch.cuda.amp import GradScaler, autocast
+import numpy as np
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import matplotlib.pyplot as plt
 
-# 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()
+# --- Define the Model ---
+class FineGrainedClassifier(nn.Module):
+    def __init__(self, num_classes=434):  # Updated to 434 classes
+        super(FineGrainedClassifier, self).__init__()
+        self.image_encoder = models.resnet50(pretrained=True)
+        self.image_encoder.fc = nn.Identity()
+        self.text_encoder = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en')
+        self.classifier = nn.Sequential(
+            nn.Linear(2048 + 768, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)  # Updated to 434 classes
+        )
+    
+    def forward(self, image, input_ids, attention_mask):
+        image_features = self.image_encoder(image)
+        text_output = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)
+        text_features = text_output.last_hidden_state[:, 0, :]
+        combined_features = torch.cat((image_features, text_features), dim=1)
+        output = self.classifier(combined_features)
+        return output
 
-# Load the model and tokenizer from your Hugging Face repository
-model = AutoModel.from_pretrained("Maverick98/EcommerceClassifier")
-tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-base-en")
 
-# Define image preprocessing
+# --- Data Augmentation Setup ---
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(15),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
     transforms.ToTensor(),
     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"
+model.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_checkpoint}/resolve/main/model_checkpoint.pth", map_location=torch.device('cpu')))
+# Load the tokenizer from Jina
+tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-base-en")
+
+# # Define image preprocessing
+# transform = transforms.Compose([
+#     transforms.Resize((224, 224)),
+#     transforms.ToTensor(),
+#     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+# ])
+
 def load_image(image_path_or_url):
     """
     Load an image from a URL or local path and preprocess it.
@@ -45,7 +118,7 @@ def predict(image_path_or_url, title, threshold=0.7):
     image = load_image(image_path_or_url)
     
     # Tokenize the title
-    title_encoding = tokenizer(title, padding='max_length', max_length=32, truncation=True, return_tensors='pt')
+    title_encoding = tokenizer(title, padding='max_length', max_length=200, truncation=True, return_tensors='pt')
     input_ids = title_encoding['input_ids']
     attention_mask = title_encoding['attention_mask']