Upload app.py with huggingface_hub

app.py

@@ -1,71 +1,52 @@
 import torch
-import torchvision.models as models
+import torchvision.transforms as transforms
 import gradio as gr
-from huggingface_hub import hf_hub_download
+from torchvision import models
 from PIL import Image
-from torchvision import transforms
 
-# ✅ Download model checkpoint from Hugging Face Hub
-model_path = hf_hub_download(
-    repo_id="PrachiY/image-classification-model",
-    filename="clothing1m.pth.tar"
-)
+# Define Clothing1M class labels
+clothing1m_classes = [
+    "T-shirt", "Shirt", "Knitwear", "Chiffon", "Sweater", "Hoodie", "Windbreaker",
+    "Jacket", "Down Coat", "Suits", "Shawl", "Dress", "Vest", "Underwear", "Shorts",
+    "Trousers", "Jeans", "Leather Shoes", "Casual Shoes", "Sport Shoes", "Sandals"
+]
 
-# ✅ Load the Model
+# ✅ Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"✅ Using device: {device}")
-
-model = models.resnet50(pretrained=False)
-checkpoint = torch.load(model_path, map_location=device)
-
-if "model" in checkpoint:
-    model.load_state_dict(checkpoint["model"], strict=False)
-elif "state_dict" in checkpoint:
-    model.load_state_dict(checkpoint["state_dict"], strict=False)
-else:
-    model.load_state_dict(checkpoint, strict=False)
 
-# Ensure correct output layer (21 classes for Clothing1M)
-model.fc = torch.nn.Linear(2048, 21)
+# ✅ Load model
+model = models.resnet50(weights=None)  # Ensure correct architecture
+num_ftrs = model.fc.in_features
+model.fc = torch.nn.Linear(num_ftrs, 21)  # Match Clothing1M class count
+model.load_state_dict(torch.load("model.pth", map_location=device))  # Load weights
 model.to(device)
 model.eval()
 
-# ✅ Define Clothing1M Class Labels
-class_labels = [
-    "T-shirt", "Shirt", "Knitwear", "Chiffon", "Sweater", "Hoodie", "Windbreaker",
-    "Jacket", "Downcoat", "Suits", "Shawl", "Dress", "Vest", "Underwear",
-    "Hat", "Sock", "Jeans", "Sweatpants", "Trousers", "Shorts", "Skirt"
-]
-
-# ✅ Image Preprocessing
+# ✅ Define image preprocessing
 def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
-    image = transform(image).unsqueeze(0).to(device)  # Ensure tensor is on GPU if available
-    return image
+    return transform(image).unsqueeze(0).to(device)
 
-# ✅ Prediction Function
-def predict(image):
+# ✅ Define inference function
+def classify_image(image):
     image_tensor = preprocess_image(image)
     with torch.no_grad():
         output = model(image_tensor)
         predicted_class_idx = output.argmax(dim=1).item()
+    predicted_class_name = clothing1m_classes[predicted_class_idx] if predicted_class_idx < len(clothing1m_classes) else "Unknown"
+    return f"Predicted Class: {predicted_class_name}"
 
-        if predicted_class_idx >= len(class_labels):
-            return f"Predicted Class: Unknown (Index {predicted_class_idx} out of range)"
-
-        return f"Predicted Class: {class_labels[predicted_class_idx]}"
-
-# ✅ Gradio Interface
+# ✅ Create Gradio Interface
 interface = gr.Interface(
-    fn=predict,
+    fn=classify_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Clothing1M Image Classifier",
-    description="Upload an image to classify it into one of 21 clothing categories."
+    title="Clothing1M Classifier",
+    description="Upload an image of clothing and get the predicted category."
 )
 
 if __name__ == "__main__":