Upload 2 files

Adjustments to UI

app.py

@@ -8,7 +8,7 @@ import json
 # ---------------------------------------------------------------------------
 # Car brands — must match the classes the ViT model was trained on
 # ---------------------------------------------------------------------------
-CAR_BRANDS = ['BMW', 'Ferrari', 'Ford', 'Jeep', 'Lamborghini', 'Porsche', 'Rolls-Royce', 'Toyota']
+CAR_BRANDS = ['BMW', 'Dodge', 'Ferrari', 'Ford', 'Jeep', 'Lamborghini', 'Porsche', 'Rolls-Royce', 'Toyota']
 
 # ---------------------------------------------------------------------------
 # Load models (loaded once at startup)
@@ -38,47 +38,51 @@ def encode_image_to_base64(image_path: str) -> str:
 
 def classify_with_openai(image_path: str) -> dict:
     """Send image to GPT-4o and ask it to return confidence scores per brand."""
-    api_key = os.environ.get("OPENAI_API_KEY")
-    client = openai.OpenAI(api_key=api_key)
-
-    ext = os.path.splitext(image_path)[1].lower().lstrip(".")
-    mime_type = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"
-    base64_image = encode_image_to_base64(image_path)
-
-    prompt = (
-        f"You are a car classification expert. Classify the car brand shown in this image. "
-        f"The possible classes are: {', '.join(CAR_BRANDS)}. "
-        "Respond ONLY with a valid JSON object where each key is a brand name from the list and "
-        "each value is a confidence score between 0.0 and 1.0. All scores must sum to 1.0. "
-        'Example format: {"BMW": 0.05, "Ferrari": 0.85, "Ford": 0.02, ...}'
-    )
-
-    response = client.chat.completions.create(
-        model="gpt-4o",
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": f"data:{mime_type};base64,{base64_image}"},
-                    },
-                ],
-            }
-        ],
-        max_tokens=300,
-    )
-
-    text = response.choices[0].message.content
     try:
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if not api_key:
+            return {"Error: OPENAI_API_KEY not set": 1.0}
+        client = openai.OpenAI(api_key=api_key)
+
+        ext = os.path.splitext(image_path)[1].lower().lstrip(".")
+        mime_type = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"
+        base64_image = encode_image_to_base64(image_path)
+
+        prompt = (
+            f"You are a car classification expert. Classify the car brand shown in this image. "
+            f"The possible classes are: {', '.join(CAR_BRANDS)}. "
+            "Respond ONLY with a valid JSON object where each key is a brand name from the list and "
+            "each value is a confidence score between 0.0 and 1.0. All scores must sum to 1.0. "
+            'Example format: {"BMW": 0.05, "Ferrari": 0.85, "Ford": 0.02, ...}'
+        )
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:{mime_type};base64,{base64_image}"},
+                        },
+                    ],
+                }
+            ],
+            max_tokens=300,
+        )
+
+        text = response.choices[0].message.content
         start = text.find("{")
         end = text.rfind("}") + 1
         scores = json.loads(text[start:end])
         return {brand: float(scores.get(brand, 0.0)) for brand in CAR_BRANDS}
-    except Exception:
-        uniform = 1.0 / len(CAR_BRANDS)
-        return {brand: uniform for brand in CAR_BRANDS}
+
+    except openai.AuthenticationError:
+        return {"Error: Invalid OpenAI API key": 1.0}
+    except Exception as e:
+        return {f"Error: {str(e)[:60]}": 1.0}
 
 
 # ---------------------------------------------------------------------------
@@ -107,46 +111,64 @@ def classify_car(image):
 # Example images (add representative car images to example_images/)
 # ---------------------------------------------------------------------------
 
+_img_dir = os.path.join(os.path.dirname(__file__), "example_images")
 example_images = [
-    ["example_images/ferrari.jpg"],
-    ["example_images/lamborghini.jpg"],
-    ["example_images/bmw.jpg"],
-    ["example_images/jeep.jpg"],
-    ["example_images/ford.jpg"],
-    ["example_images/toyota.jpg"],
-    ["example_images/porsche.jpg"],
-    ["example_images/rolls_royce.jpg"],
+    [os.path.join(_img_dir, f)]
+    for f in sorted(os.listdir(_img_dir))
+    if f.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))
 ]
 
 # ---------------------------------------------------------------------------
 # Gradio UI
 # ---------------------------------------------------------------------------
 
-with gr.Blocks(title="Car Brand Classification Comparison") as demo:
-    gr.Markdown("# Car Brand Classification — Model Comparison")
-    gr.Markdown(
-        "Upload a car image to compare predictions from three different models:\n\n"
-        "- **Custom ViT** — fine-tuned on Stanford Cars (8 brand classes)\n"
-        "- **CLIP** — zero-shot classification with `openai/clip-vit-large-patch14`\n"
-        "- **OpenAI GPT-4o** — vision LLM classification\n\n"
-        f"**Classes:** {', '.join(CAR_BRANDS)}"
-    )
+css = """
+.title { text-align: center; margin-bottom: 0.25rem; }
+.subtitle { text-align: center; color: #6b7280; margin-bottom: 1.5rem; font-size: 0.95rem; }
+.model-header { font-weight: 600; font-size: 1rem; margin-bottom: 0.25rem; padding: 0.4rem 0.75rem;
+                border-radius: 6px; background: #f3f4f6; }
+.classify-btn { max-width: 200px; margin: 0 auto; }
+footer { display: none !important; }
+"""
 
-    with gr.Row():
-        input_image = gr.Image(type="filepath", label="Upload Car Image")
+with gr.Blocks(title="Car Brand Classification", css=css, theme=gr.themes.Soft()) as demo:
 
-    classify_btn = gr.Button("Classify", variant="primary")
+    # Header
+    gr.Markdown("# 🚗 Car Brand Classification", elem_classes="title")
+    gr.Markdown(
+        "Upload a car image and compare predictions from three models side by side.",
+        elem_classes="subtitle"
+    )
 
+    # Input + button
     with gr.Row():
-        with gr.Column():
-            gr.Markdown("### Custom ViT Model")
-            vit_output = gr.Label(num_top_classes=8, label="ViT Predictions")
-        with gr.Column():
-            gr.Markdown("### CLIP Zero-Shot")
-            clip_output = gr.Label(num_top_classes=8, label="CLIP Predictions")
-        with gr.Column():
-            gr.Markdown("### OpenAI GPT-4o Vision")
-            openai_output = gr.Label(num_top_classes=8, label="OpenAI Predictions")
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                type="filepath",
+                label="Car Image",
+                height=280,
+            )
+            classify_btn = gr.Button("Classify", variant="primary", elem_classes="classify-btn")
+
+        # Results
+        with gr.Column(scale=2):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("**Custom ViT** — fine-tuned", elem_classes="model-header")
+                    vit_output = gr.Label(num_top_classes=5, label="")
+                with gr.Column():
+                    gr.Markdown("**CLIP** — zero-shot", elem_classes="model-header")
+                    clip_output = gr.Label(num_top_classes=5, label="")
+                with gr.Column():
+                    gr.Markdown("**GPT-4o** — vision LLM", elem_classes="model-header")
+                    openai_output = gr.Label(num_top_classes=5, label="")
+
+    # Examples
+    gr.Examples(
+        examples=example_images,
+        inputs=input_image,
+        label="Example Images",
+    )
 
     classify_btn.click(
         fn=classify_car,
@@ -154,10 +176,4 @@ with gr.Blocks(title="Car Brand Classification Comparison") as demo:
         outputs=[vit_output, clip_output, openai_output],
     )
 
-    gr.Examples(
-        examples=example_images,
-        inputs=input_image,
-        label="Example Images",
-    )
-
 demo.launch()