Update app.py

app.py

@@ -19,10 +19,11 @@ from PIL import Image
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load a vanilla VGG16 model pretrained on ImageNet
+# Load vanilla VGG16 pretrained on ImageNet
 model = models.vgg16(weights="IMAGENET1K_V1").to(device)
 model.eval()
 
+# Download ImageNet labels
 LABELS_URL = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
 try:
     LABELS_CACHE = requests.get(LABELS_URL, timeout=5).json()
@@ -30,13 +31,13 @@ except Exception as e:
     print(f"Could not fetch ImageNet labels: {e}")
     LABELS_CACHE = [f"Class {i}" for i in range(1000)]
 
-# Transformation pipeline for input images
+# Transform pipeline
 transform_pipeline = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],  # ImageNet means
-        std=[0.229, 0.224, 0.225]    # ImageNet std
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
     )
 ])
 
@@ -45,117 +46,86 @@ transform_pipeline = transforms.Compose([
 # -----------------------------
 def classify_image(image, confidence_threshold=0.0):
     """
-    Classify an image using pretrained VGG16 on ImageNet.
-    Returns the top-3 predictions above confidence_threshold.
+    Classify an image using the pretrained VGG16 on ImageNet.
+    Returns top-3 predictions above the given confidence_threshold.
     """
     try:
-        # Convert Gradio's image (numpy) to PIL
+        # Convert Gradio's numpy image to PIL
         if isinstance(image, np.ndarray):
             image_pil = Image.fromarray(image.astype('uint8'), 'RGB')
         else:
             image_pil = Image.open(image).convert('RGB')
 
-        # Apply preprocessing
+        # Preprocess
         input_tensor = transform_pipeline(image_pil).unsqueeze(0).to(device)
 
         # Inference
         with torch.no_grad():
             output = model(input_tensor)
-            probabilities = torch.nn.functional.softmax(output, dim=1)
+            probs = torch.nn.functional.softmax(output, dim=1)
 
         # Top-3 predictions
-        top_probs, top_classes = torch.topk(probabilities, 3)
+        top_probs, top_cls_idxs = torch.topk(probs, 3)
         top_probs = top_probs[0].cpu().numpy()
-        top_classes = top_classes[0].cpu().numpy()
+        top_cls_idxs = top_cls_idxs[0].cpu().numpy()
 
-        # Build dictionary with label -> probability
         results = {}
-        for prob, cls_idx in zip(top_probs, top_classes):
-            if prob >= confidence_threshold:
-                label = LABELS_CACHE[cls_idx] if LABELS_CACHE else f"Class {cls_idx}"
-                results[label] = float(prob)
+        for p, cidx in zip(top_probs, top_cls_idxs):
+            if p >= confidence_threshold:
+                label = LABELS_CACHE[cidx] if LABELS_CACHE else f"Class {cidx}"
+                results[label] = float(p)
 
         if not results:
             return "No predictions above the confidence threshold."
         return results
+
     except Exception as e:
         return f"Error during classification: {str(e)}"
 
 # -----------------------------
-#  CUSTOM CSS FOR BACKGROUND
+#  (OPTIONAL) CUSTOM CSS
 # -----------------------------
-# Replace the background color/image/gradient with whatever you prefer.
-# You can also style text, buttons, etc.
 custom_css = """
 body {
     margin: 0;
     padding: 0;
-    background: linear-gradient(135deg, #f2f2f2, #dceeff);
-    font-family: 'Arial', sans-serif;
-}
-
-#title {
-    font-size: 2.5rem;
-    text-align: center;
-    margin-top: 20px;
-    font-weight: bold;
-    color: #333;
+    background: linear-gradient(135deg, #f6f9fc, #ddeefc);
+    font-family: "Helvetica", sans-serif;
 }
-
-#subtext {
+h1, p {
     text-align: center;
-    font-size: 1rem;
-    color: #555;
-    margin-bottom: 20px;
+    margin-bottom: 1rem;
 }
 """
 
 # -----------------------------
-#  BUILD GRADIO INTERFACE
+#  BUILD THE GRADIO APP
 # -----------------------------
-def build_interface():
-    # Inputs
-    image_input = gr.Image(type="numpy", label="Upload an Image")
-    confidence_slider = gr.Slider(
-        minimum=0.0,
-        maximum=1.0,
-        value=0.0,
-        step=0.01,
-        label="Confidence Threshold"
-    )
+def build_app():
+    with gr.Blocks(css=custom_css) as demo:
+        gr.HTML("<h1>VGG16 ImageNet Classifier</h1>")
+        gr.HTML("<p>Upload an image to see the top 3 predicted ImageNet classes.</p>")
 
-    # Outputs
-    label_output = gr.Label(num_top_classes=3)
+        with gr.Box():
+            # Place widgets in a vertical layout
+            image_input = gr.Image(type="numpy", label="Upload an Image")
+            confidence_slider = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Confidence Threshold")
+            classify_button = gr.Button("Classify")
 
-    # An optional HTML block (header text, etc.)
-    with gr.Blocks(css=custom_css) as demo:
-        gr.HTML("<h1 id='title'>VGG16 ImageNet Classifier</h1>")
-        gr.HTML("<p id='subtext'>Upload an image to see top ImageNet predictions from a pretrained VGG16 model.</p>")
-
-        # Main interface
-        with gr.Row():
-            with gr.Column():
-                image_in = image_input
-                conf_slider = confidence_slider
-            with gr.Column():
-                label_out = label_output
-
-        # Create the main Interface
-        btn = gr.Button("Classify")
-        btn.click(
+            label_output = gr.Label(num_top_classes=3, label="Prediction Results")
+
+        # Connect button click to classification
+        classify_button.click(
             fn=classify_image,
-            inputs=[image_in, conf_slider],
-            outputs=label_out
+            inputs=[image_input, confidence_slider],
+            outputs=label_output
         )
 
     return demo
 
 # -----------------------------
-#  LAUNCH APP
+#  LAUNCH
 # -----------------------------
 if __name__ == "__main__":
-    interface = build_interface()
-    # You can set a Gradio theme if you want (e.g., 'Soft', 'Monochrome', 'Glass')
-    interface.launch(share=True)  # share=True if you want a shareable public link locally
-
-
+    demo = build_app()
+    demo.launch()