Update app.py

app.py

@@ -1,35 +1,67 @@
-import gradio as gr
+import torch
+import torchvision.transforms as transforms
 from PIL import Image
-from inference import load_model, run_inference
+import gradio as gr
 
+from model import NetFeat, NetClassifier  # Ensure this imports correctly based on your setup
 
-net_feat, net_cls = load_model()
+# Define class labels for clothing classification
+CLOTHING_CLASSES = [
+    "T-shirt", "Shirt", "Shawl", "Dress", "Vest", "Underwear", "Cardigan", "Jacket",
+    "Sweater", "Hoodie", "Knitwear", "Chiffon", "Downcoat", "Suit"
+]
 
-def classify_image(image):
-    """Classify the image using the loaded model and return the predicted class."""
-    try:
-        prediction = run_inference(image, net_feat, net_cls)
-        return prediction
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
+# Load the model
+def load_model():
+    model_filename = 'netBest.pth'  # Adjust the path as necessary
+    net_feat = NetFeat(arch='resnet18', pretrained=False, dataset='Clothing1M')
+    net_cls = NetClassifier(feat_dim=512, nb_cls=14)
+    
+    state_dict = torch.load(model_filename, map_location=torch.device('cpu'))
+    if "feat" in state_dict:
+        net_feat.load_state_dict(state_dict['feat'], strict=False)
+    if "cls" in state_dict:
+        net_cls.load_state_dict(state_dict['cls'], strict=False)
 
+    net_feat.eval()
+    net_cls.eval()
+    return net_feat, net_cls
 
-example_images = [
-    "sample1.jpg", "sample2.jpg", "sample3.jpg",
-    "sample4.jpg", "sample5.jpg", "sample6.jpg"
-]
+# Preprocess image for model input
+def preprocess_image(image):
+    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])
+    ])
+    image = Image.open(image).convert("RGB")
+    return transform(image).unsqueeze(0)  # Add batch dimension
 
+# Run inference on the image
+def run_inference(image, net_feat, net_cls):
+    image_tensor = preprocess_image(image)
+    with torch.no_grad():
+        feature_vector = net_feat(image_tensor)
+        output = net_cls(feature_vector)
+    predicted_index = output.argmax(dim=1).item()
+    return CLOTHING_CLASSES[predicted_index]
 
+# Load models
+net_feat, net_cls = load_model()
+
+# Define the Gradio interface function
+def classify_image(image):
+    return run_inference(image, net_feat, net_cls)
+
+# Set up Gradio interface
 interface = gr.Interface(
     fn=classify_image,
-    inputs=gr.inputs.Image(type="pil", label="Upload an Image"),  
-    outputs=gr.outputs.Textbox(label="Predicted Clothing1M Class"),
+    inputs=gr.Image(shape=(224, 224), label="Upload an Image"),
+    outputs=gr.Textbox(label="Predicted Clothing1M Class"),
     title="Clothing1M Classifier",
-    description="Upload an image of clothing, and the classifier will identify its category among 14 different types, such as T-shirts, Dresses, Jackets, etc.",
-    examples=example_images,
-    theme="default"  
+    description="Upload an image of clothing to classify it into one of 14 categories."
 )
 
-
+# Launch the application
 if __name__ == "__main__":
     interface.launch()