Update app.py

app.py

@@ -4,6 +4,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.transforms as transforms
 from PIL import Image
+import numpy as np
+import os
 from ResNet_for_CC import CC_model  # Import the model
 
 # Set device (CPU/GPU)
@@ -26,7 +28,15 @@ class_labels = [
     "Vest", "Underwear"
 ]
 
-# ✅ **Updated Image Preprocessing Function**
+# ✅ **Predefined Default Images**
+default_images = {
+    "T-Shirt": "default_images/tshirt.jpg",
+    "Jacket": "default_images/jacket.jpg",
+    "Sweater": "default_images/sweater.jpg",
+    "Dress": "default_images/dress.jpg"
+}
+
+# ✅ **Image Preprocessing Function**
 def preprocess_image(image):
     """Applies necessary transformations to the input image."""
     transform = transforms.Compose([
@@ -38,18 +48,27 @@ def preprocess_image(image):
     return transform(image).unsqueeze(0).to(device)
 
 # ✅ **Classification Function**
-def classify_image(image):
-    """Processes the input image and returns the predicted clothing category."""
-    print("\n[INFO] Received image for classification.")
+def classify_image(selected_default, uploaded_image):
+    """Processes either a default or uploaded image and returns the predicted clothing category."""
 
+    print("\n[INFO] Image selection process started.")
+    
     try:
-        image = Image.fromarray(image)  # Ensure conversion to PIL format
+        # Use the uploaded image if provided; otherwise, use the selected default image
+        if uploaded_image is not None:
+            print("[INFO] Using uploaded image.")
+            image = Image.fromarray(uploaded_image)  # Ensure conversion to PIL format
+        else:
+            print(f"[INFO] Using default image: {selected_default}")
+            image_path = default_images[selected_default]
+            image = Image.open(image_path)  # Load the selected default image
+
         image = preprocess_image(image)  # Apply transformations
         print("[INFO] Image transformed and moved to device.")
 
         with torch.no_grad():
             output = model(image)
-            
+
             # ✅ Ensure output is a tensor (handle tuple case)
             if isinstance(output, tuple):
                 output = output[1]  # Extract the actual output tensor
@@ -81,13 +100,32 @@ def classify_image(image):
         return "Error in classification. Check console for details."
 
 # ✅ **Gradio Interface**
-interface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="numpy"),
-    outputs="text",
-    title="Clothing1M Image Classifier",
-    description="Upload a clothing image, and the model will classify it into one of the 14 categories."
-)
+with gr.Blocks() as interface:
+    gr.Markdown("# Clothing1M Image Classifier")
+    gr.Markdown("Upload a clothing image or select from the predefined images below.")
+
+    # Default Image Selection
+    default_selector = gr.Radio(
+        choices=list(default_images.keys()),
+        label="Select a Default Image",
+        value="T-Shirt"
+    )
+
+    # File Upload Option
+    image_upload = gr.Image(type="numpy", label="Or Upload Your Own Image")
+
+    # Output Text
+    output_text = gr.Textbox(label="Classification Result")
+
+    # Classify Button
+    classify_button = gr.Button("Classify Image")
+
+    # Define Action
+    classify_button.click(
+        fn=classify_image,
+        inputs=[default_selector, image_upload],
+        outputs=output_text
+    )
 
 # ✅ **Run the Interface**
 if __name__ == "__main__":