| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| from typing import Dict, List | |
| import gradio as gr | |
| from PIL import Image | |
| ROOT = Path(__file__).resolve().parents[1] | |
| if str(ROOT) not in sys.path: | |
| sys.path.insert(0, str(ROOT)) | |
| from src.inference import ClipPredictor, CustomModelPredictor, OpenAIVisionPredictor | |
| custom = CustomModelPredictor(str(ROOT / "models" / "custom_resnet18.pth")) | |
| if custom.available(): | |
| labels = custom.labels | |
| else: | |
| labels = ["charizard", "charmander", "charmeleon", "ditto", "eevee", "ekans"] | |
| clip_model = ClipPredictor(labels) | |
| openai_model = OpenAIVisionPredictor(labels) | |
| def _format_preds(result: Dict[str, object]) -> str: | |
| if not result.get("available", False): | |
| return f"Unavailable: {result.get('error', 'unknown error')}" | |
| lines: List[str] = [] | |
| top = result.get("top_prediction", {}) | |
| label = top.get("label", "-") | |
| confidence = float(top.get("confidence", 0.0)) | |
| lines.append(f"Top prediction: {label} ({confidence:.2%})") | |
| for pred in result.get("predictions", []): | |
| lines.append(f"- {pred['label']}: {pred['confidence']:.2%}") | |
| raw_response = result.get("raw_response") | |
| if isinstance(raw_response, dict) and raw_response.get("reason"): | |
| lines.append(f"Reason: {raw_response['reason']}") | |
| return "\n".join(lines) | |
| def classify_image(image: Image.Image): | |
| if image is None: | |
| return "No image provided.", "No image provided.", "No image provided." | |
| custom_pred = custom.predict(image) | |
| clip_pred = clip_model.predict(image) | |
| openai_pred = openai_model.predict(image) | |
| return _format_preds(custom_pred), _format_preds(clip_pred), _format_preds(openai_pred) | |
| def get_examples() -> List[List[str]]: | |
| examples_dir = ROOT / "app" / "examples" | |
| if not examples_dir.exists(): | |
| return [] | |
| image_paths = sorted( | |
| [p for p in examples_dir.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}] | |
| ) | |
| return [[str(p)] for p in image_paths] | |
| description = """ | |
| Upload an image and compare predictions from three models: | |
| 1) Custom transfer learning model (ResNet18) | |
| 2) Open-source CLIP model | |
| 3) Closed-source OpenAI vision model | |
| If OPENAI_API_KEY is not set, OpenAI predictions are shown as unavailable. | |
| """ | |
| with gr.Blocks(title="Computer Vision Model Comparison") as demo: | |
| gr.Markdown("# Computer Vision Classification & Model Comparison") | |
| gr.Markdown(description) | |
| with gr.Row(): | |
| image_input = gr.Image(type="pil", label="Upload image") | |
| classify_button = gr.Button("Classify") | |
| with gr.Row(): | |
| custom_output = gr.Textbox(label="Custom Transfer Learning", lines=8) | |
| clip_output = gr.Textbox(label="Open-Source CLIP", lines=8) | |
| openai_output = gr.Textbox(label="Closed-Source OpenAI Vision", lines=8) | |
| classify_button.click(classify_image, inputs=[image_input], outputs=[custom_output, clip_output, openai_output]) | |
| gr.Examples(examples=get_examples(), inputs=image_input, label="Example images") | |
| if __name__ == "__main__": | |
| demo.launch() | |