| from __future__ import annotations
|
|
|
| import os
|
| import sys
|
| from functools import lru_cache
|
| from pathlib import Path
|
| from typing import Dict, List
|
|
|
| import gradio as gr
|
| from PIL import Image
|
|
|
| ROOT = Path(__file__).resolve().parent
|
| if str(ROOT) not in sys.path:
|
| sys.path.insert(0, str(ROOT))
|
|
|
| from src.inference import ClipPredictor, CustomModelPredictor, OpenAIVisionPredictor
|
|
|
|
|
| labels = ["charizard", "charmander", "charmeleon", "ditto", "eevee", "ekans"]
|
|
|
|
|
| @lru_cache(maxsize=1)
|
| def get_predictors() -> tuple[CustomModelPredictor, ClipPredictor, OpenAIVisionPredictor]:
|
| custom = CustomModelPredictor(str(ROOT / "models" / "custom_resnet18.pth"))
|
| predictor_labels = custom.labels if custom.available() else labels
|
| clip_model = ClipPredictor(predictor_labels)
|
| openai_model = OpenAIVisionPredictor(predictor_labels)
|
| return custom, clip_model, openai_model
|
|
|
|
|
| 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, clip_model, openai_model = get_predictors()
|
| 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__":
|
| in_space = bool(os.getenv("SPACE_ID"))
|
| if in_space:
|
| demo.launch(ssr_mode=False)
|
| else:
|
| print("Launching local app on http://127.0.0.1:7860")
|
| demo.launch(server_name="127.0.0.1", server_port=7860, share=False, ssr_mode=False)
|
|
|
