Update app.py

app.py

@@ -1,10 +1,8 @@
 import gradio as gr
 import torch
-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
 
@@ -36,8 +34,8 @@ default_images = {
     "Vest": "dress.jpg"
 }
 
-# Convert image paths to a format suitable for Gradio
-default_image_paths = [[img_path] for img_path in default_images.values()]
+# Convert to gallery format (list of (image_path, caption) tuples)
+default_images_gallery = [(path, label) for label, path in default_images.items()]
 
 # **Image Preprocessing Function**
 def preprocess_image(image):
@@ -53,85 +51,63 @@ def preprocess_image(image):
 # **Classification Function**
 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:
         # 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
+            image = Image.fromarray(uploaded_image)
         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.")
+            image = Image.open(image_path)
 
+        image = preprocess_image(image)
+        
         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
-
-            print(f"[DEBUG] Model output shape: {output.shape}")
-            print(f"[DEBUG] Model output values: {output}")
+                output = output[1]
 
-            if output.shape[1] != 14:
-                return f"[ERROR] Model output mismatch! Expected 14 but got {output.shape[1]}."
-
-            # Convert logits to probabilities
             probabilities = F.softmax(output, dim=1)
-            print(f"[DEBUG] Softmax probabilities: {probabilities}")
-
-            # Get predicted class index
             predicted_class = torch.argmax(probabilities, dim=1).item()
-            print(f"[INFO] Predicted class index: {predicted_class} (Class: {class_labels[predicted_class]})")
 
-            # Validate and return the prediction
             if 0 <= predicted_class < len(class_labels):
                 predicted_label = class_labels[predicted_class]
                 confidence = probabilities[0][predicted_class].item() * 100
                 return f"Predicted Class: {predicted_label} (Confidence: {confidence:.2f}%)"
             else:
                 return "[ERROR] Model returned an invalid class index."
-
+    
     except Exception as e:
-        print(f"[ERROR] Exception during classification: {e}")
-        return "Error in classification. Check console for details."
+        return f"Error in classification: {e}"
 
 # **Gradio Interface**
 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 Gallery for Selection**
-    gr.Markdown("### Select a Default Image:")
+    # Gallery to display default images
     gallery = gr.Gallery(
-        value=default_image_paths,
-        label="Available Default Images",
-        show_label=True
+        value=default_images_gallery,  # Provide list of (image, caption) tuples
+        label="Default Images",
+        elem_id="default_gallery"
     )
 
-    # **Default Image Selection Radio**
-    default_selector = gr.Radio(
+    # Default Image Selection
+    default_selector = gr.Dropdown(
         choices=list(default_images.keys()),
         label="Select a Default Image",
         value="Shawl"
     )
 
-    # **File Upload Option**
+    # File Upload Option
     image_upload = gr.Image(type="numpy", label="Or Upload Your Own Image")
 
-    # **Output Text**
+    # Output Text
     output_text = gr.Textbox(label="Classification Result")
 
-    # **Classify Button**
+    # Classify Button
     classify_button = gr.Button("Classify Image")
 
-    # **Define Action**
+    # Define Action
     classify_button.click(
         fn=classify_image,
         inputs=[default_selector, image_upload],