init commit

README.md

@@ -1,12 +1,158 @@
----
-title: Computer Vision Classification
-emoji: 🏃
-colorFrom: gray
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.12.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Exercise - Computer Vision Classification & Model Comparison
+
+## 1. Project Overview
+
+This project implements a full image classification pipeline and web application that compares three model types on a custom dataset:
+
+1. Custom transfer learning model (ResNet18 fine-tuned on custom data)
+2. Open-source model (CLIP: `openai/clip-vit-base-patch32`)
+3. Closed-source model (OpenAI Vision API)
+
+The app supports image upload, example images, and side-by-side predictions from all three models.
+
+## 2. Dataset Description
+
+### Dataset
+
+Custom image dataset based on Pokemon classes from the course materials (week 8 style transfer learning setup):
+
+- Classes: `charizard`, `charmander`, `charmeleon`, `ditto`, `eevee`, `ekans`
+- Train split: `data/pokemon/train/<class>`
+- Test split: `data/pokemon/test/<class>`
+
+### Dataset size used in evaluation
+
+- Number of classes: 6
+- Test images: 25
+
+## 3. Preprocessing Steps
+
+For transfer learning training and inference:
+
+- Convert to RGB
+- Train transforms:
+  - `RandomResizedCrop(224)`
+  - `RandomHorizontalFlip()`
+  - `ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2)`
+  - ImageNet normalization
+- Eval transforms:
+  - `Resize(256)`
+  - `CenterCrop(224)`
+  - ImageNet normalization
+
+CLIP uses its own processor and tokenizer.
+
+OpenAI Vision receives the image as base64-encoded JPEG together with a strict label selection prompt.
+
+## 4. Model and Training
+
+### Custom Transfer Learning Model
+
+- Backbone: `ResNet18` with ImageNet pretrained weights
+- Final layer replaced to output 6 classes
+- Loss: CrossEntropyLoss
+- Optimizer: Adam
+- Learning rate: `1e-4`
+- Weight decay: `1e-4`
+- Epochs run: `4`
+- Device used: CPU
+
+### Training Output Artifacts
+
+- Model checkpoint: `models/custom_resnet18.pth`
+- Training metrics: `reports/custom_model_metrics.json`
+
+## 5. Evaluation and Comparison Results
+
+### Summary Accuracy
+
+- Custom transfer learning model: **0.80**
+- CLIP (open-source): **0.72**
+- OpenAI Vision (closed-source): not executed locally in this run because `OPENAI_API_KEY` was not set in the local environment
+
+### Interpretation
+
+- The custom fine-tuned model performs best on this domain-specific dataset.
+- CLIP performs reasonably well in zero-shot mode without task-specific training.
+- OpenAI Vision integration is fully implemented in the app and evaluation pipeline and becomes active once `OPENAI_API_KEY` is configured.
+
+## 6. Application (Hugging Face Space Ready)
+
+The Gradio app provides:
+
+- image upload
+- predictions from custom model, CLIP, and OpenAI Vision
+- built-in example images (`app/examples/*.png`)
+
+### App files
+
+- Main HF entrypoint: `app.py`
+- App implementation: `app/app.py`
+- Space dependencies: `app/requirements.txt`
+
+## 7. Links (to fill after publishing)
+
+- Hugging Face Space: `https://huggingface.co/spaces/<your-username>/<your-space-name>`
+- Hugging Face model repo: `https://huggingface.co/<your-username>/<your-model-repo>`
+
+## 8. OpenAI Key Setup (for Hugging Face Space)
+
+In your Hugging Face Space:
+
+1. Open `Settings`
+2. Go to `Variables and secrets`
+3. Add new secret:
+   - Name: `OPENAI_API_KEY`
+   - Value: use the key provided in the exercise sheet
+
+Do not hardcode the API key in repository files.
+
+## 9. Reproducibility: Local Run Commands
+
+```bash
+# from project root
+python -m venv .venv
+.venv\\Scripts\\activate
+pip install -r requirements.txt
+
+# train custom model
+python -m src.train_custom_model --epochs 4 --batch-size 16
+
+# evaluate and compare models
+python -m src.evaluate_models --openai-max-samples 24
+
+# launch app
+python app.py
+```
+
+## 10. Hugging Face Model Upload
+
+A helper script is included:
+
+```bash
+python -m src.upload_to_hf --repo-id <your-username>/<your-model-repo>
+```
+
+This uploads `models/custom_resnet18.pth` to your Hugging Face model repository.
+
+## 11. Notes
+
+- This implementation follows the mandatory exercise requirements end-to-end (training, comparison, app, example images, documentation).
+- If you want to switch from Pokemon to a car dataset, you can keep the same pipeline and replace `data/pokemon/*` with your car class folders.
+- On Windows with very long folder paths, package installation can fail with path-length errors. In that case, create a short-path venv (example: `C:\cv-modelcmp-venv`) and run the same commands with that Python executable.
+# Computer Vision Classification and Model Comparison
+
+Project implementation for the mandatory exercise.
+
+The submission-specific documentation is in [readme.md](readme.md), following the requested structure.
+
+## Quick start
+
+```bash
+python -m venv .venv
+.venv\\Scripts\\activate
+pip install -r requirements.txt
+python -m src.train_custom_model --epochs 4
+python -m src.evaluate_models
+python app.py
+```

app.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+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__":
+    print("Launching local app on http://127.0.0.1:7860")
+    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)

app/app.py

@@ -0,0 +1,96 @@
+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()

app/requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.44.1
+pillow>=10.0.0
+torch>=2.3.0
+torchvision>=0.18.0
+transformers>=4.40.0
+openai>=1.30.0

models/custom_resnet18.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9805b9724998a575eb468422eeda02f0a368c5b1fd456ba4f340691e4d3e237
+size 44797131

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.3.0
+torchvision>=0.18.0
+transformers>=4.40.0
+gradio>=4.44.0
+pillow>=10.0.0
+numpy>=1.26.0
+scikit-learn>=1.5.0
+openai>=1.30.0
+huggingface_hub>=0.24.0,<1.0.0
+python-dotenv>=1.0.1
+matplotlib>=3.8.0

src/evaluate_models.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Dict, List
+
+from PIL import Image
+from sklearn.metrics import accuracy_score, classification_report
+
+from src.inference import ClipPredictor, CustomModelPredictor, OpenAIVisionPredictor
+
+
+def iter_test_images(data_dir: Path, labels: List[str]) -> List[Dict[str, str]]:
+    items: List[Dict[str, str]] = []
+    test_root = data_dir / "test"
+    for label in labels:
+        class_dir = test_root / label
+        if not class_dir.exists():
+            continue
+        for p in class_dir.iterdir():
+            if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}:
+                items.append({"path": str(p), "label": label})
+    return items
+
+
+def evaluate_custom(custom: CustomModelPredictor, samples: List[Dict[str, str]]) -> Dict[str, object]:
+    y_true = []
+    y_pred = []
+    for item in samples:
+        image = Image.open(item["path"])
+        pred = custom.predict(image)
+        y_true.append(item["label"])
+        y_pred.append(pred["top_prediction"]["label"])
+
+    acc = accuracy_score(y_true, y_pred)
+    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
+    return {"accuracy": float(acc), "classification_report": report}
+
+
+def evaluate_clip(clip: ClipPredictor, samples: List[Dict[str, str]]) -> Dict[str, object]:
+    if not clip.available():
+        return {"available": False, "error": "Could not load CLIP model."}
+
+    y_true = []
+    y_pred = []
+    for item in samples:
+        image = Image.open(item["path"])
+        pred = clip.predict(image)
+        y_true.append(item["label"])
+        y_pred.append(pred["top_prediction"]["label"])
+
+    acc = accuracy_score(y_true, y_pred)
+    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
+    return {"available": True, "accuracy": float(acc), "classification_report": report}
+
+
+def evaluate_openai(openai_model: OpenAIVisionPredictor, samples: List[Dict[str, str]], max_samples: int) -> Dict[str, object]:
+    if not openai_model.available():
+        return {"available": False, "error": "OPENAI_API_KEY missing."}
+
+    subset = samples[:max_samples]
+    y_true = []
+    y_pred = []
+    for item in subset:
+        image = Image.open(item["path"])
+        pred = openai_model.predict(image)
+        y_true.append(item["label"])
+        y_pred.append(pred["top_prediction"]["label"])
+
+    acc = accuracy_score(y_true, y_pred)
+    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
+    return {
+        "available": True,
+        "evaluated_samples": len(subset),
+        "accuracy": float(acc),
+        "classification_report": report,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Compare custom model vs CLIP vs OpenAI Vision.")
+    parser.add_argument("--data-dir", default="data/pokemon")
+    parser.add_argument("--model-path", default="models/custom_resnet18.pth")
+    parser.add_argument("--output", default="reports/model_comparison.json")
+    parser.add_argument("--openai-max-samples", type=int, default=24)
+    args = parser.parse_args()
+
+    data_dir = Path(args.data_dir)
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    custom = CustomModelPredictor(args.model_path)
+    if not custom.available():
+        raise FileNotFoundError(
+            "Custom model checkpoint not found. Train model first with src/train_custom_model.py"
+        )
+
+    labels = custom.labels
+    samples = iter_test_images(data_dir, labels)
+
+    clip = ClipPredictor(labels)
+    openai_model = OpenAIVisionPredictor(labels)
+
+    result = {
+        "dataset": str(data_dir),
+        "num_test_samples": len(samples),
+        "labels": labels,
+        "custom_model": evaluate_custom(custom, samples),
+        "clip_model": evaluate_clip(clip, samples),
+        "openai_model": evaluate_openai(openai_model, samples, args.openai_max_samples),
+    }
+
+    output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(f"Saved comparison report to: {output_path}")
+
+
+if __name__ == "__main__":
+    main()

src/inference.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import base64
+import io
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from openai import OpenAI
+from PIL import Image
+
+TORCH_IMPORT_ERROR: Optional[str] = None
+try:
+    import torch
+    import torch.nn.functional as F
+    from torchvision import models, transforms
+except Exception as exc:  # pragma: no cover - defensive import guard
+    torch = None
+    F = None
+    models = None
+    transforms = None
+    TORCH_IMPORT_ERROR = str(exc)
+
+TRANSFORMERS_IMPORT_ERROR: Optional[str] = None
+try:
+    from transformers import CLIPModel, CLIPProcessor
+except Exception as exc:  # pragma: no cover - defensive import guard
+    CLIPModel = None
+    CLIPProcessor = None
+    TRANSFORMERS_IMPORT_ERROR = str(exc)
+
+
+class CustomModelPredictor:
+    def __init__(self, checkpoint_path: str = "models/custom_resnet18.pth") -> None:
+        self.checkpoint_path = Path(checkpoint_path)
+        self.error: Optional[str] = None
+        self.model = None
+        self.labels: List[str] = []
+        self.image_size = 224
+        self.eval_transform = None
+
+        if torch is None or transforms is None:
+            self.device = None
+            self.error = (
+                "Custom model unavailable because torch/torchvision is missing. "
+                f"Import error: {TORCH_IMPORT_ERROR}"
+            )
+            return
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.eval_transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+
+        if self.checkpoint_path.exists():
+            self._load()
+
+    def _load(self) -> None:
+        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
+        self.labels = checkpoint["labels"]
+        self.image_size = int(checkpoint.get("image_size", 224))
+
+        model = models.resnet18(weights=None)
+        model.fc = torch.nn.Linear(model.fc.in_features, len(self.labels))
+        model.load_state_dict(checkpoint["state_dict"])
+        model.to(self.device)
+        model.eval()
+        self.model = model
+
+    def available(self) -> bool:
+        return self.model is not None and self.error is None
+
+    def predict(self, image: Image.Image, top_k: int = 3) -> Dict[str, object]:
+        if not self.available():
+            return {
+                "model": "custom-transfer-learning",
+                "available": False,
+                "error": self.error or f"Model not found at {self.checkpoint_path}",
+            }
+
+        image = image.convert("RGB")
+        tensor = self.eval_transform(image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(tensor)
+            probs = F.softmax(logits, dim=1).squeeze(0)
+
+        top_probs, top_idx = torch.topk(probs, k=min(top_k, len(self.labels)))
+        predictions = [
+            {"label": self.labels[idx], "confidence": float(prob)}
+            for prob, idx in zip(top_probs.cpu().tolist(), top_idx.cpu().tolist())
+        ]
+
+        return {
+            "model": "custom-transfer-learning",
+            "available": True,
+            "top_prediction": predictions[0],
+            "predictions": predictions,
+        }
+
+
+class ClipPredictor:
+    def __init__(self, labels: List[str], model_name: str = "openai/clip-vit-base-patch32") -> None:
+        self.labels = labels
+        self.model_name = model_name
+        self.error: Optional[str] = None
+        self.device = torch.device("cuda" if torch is not None and torch.cuda.is_available() else "cpu") if torch is not None else None
+        self.available_flag = False
+        self.processor = None
+        self.model = None
+
+        if torch is None or CLIPModel is None or CLIPProcessor is None:
+            self.error = (
+                "CLIP unavailable because required dependencies are missing. "
+                f"torch error: {TORCH_IMPORT_ERROR}; transformers error: {TRANSFORMERS_IMPORT_ERROR}"
+            )
+            return
+
+        if labels:
+            self._load()
+
+    def _load(self) -> None:
+        try:
+            self.processor = CLIPProcessor.from_pretrained(self.model_name)
+            self.model = CLIPModel.from_pretrained(self.model_name).to(self.device)
+            self.model.eval()
+            self.available_flag = True
+        except Exception:
+            self.available_flag = False
+
+    def available(self) -> bool:
+        return self.available_flag
+
+    def predict(self, image: Image.Image, top_k: int = 3) -> Dict[str, object]:
+        if not self.available():
+            return {
+                "model": "clip-open-source",
+                "available": False,
+                "error": self.error or "CLIP model could not be loaded.",
+            }
+
+        prompts = [f"a sprite or photo of a pokemon named {label}" for label in self.labels]
+        inputs = self.processor(text=prompts, images=image.convert("RGB"), return_tensors="pt", padding=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits_per_image = outputs.logits_per_image
+            probs = logits_per_image.softmax(dim=1).squeeze(0)
+
+        top_probs, top_idx = torch.topk(probs, k=min(top_k, len(self.labels)))
+        predictions = [
+            {"label": self.labels[idx], "confidence": float(prob)}
+            for prob, idx in zip(top_probs.cpu().tolist(), top_idx.cpu().tolist())
+        ]
+
+        return {
+            "model": "clip-open-source",
+            "available": True,
+            "top_prediction": predictions[0],
+            "predictions": predictions,
+        }
+
+
+class OpenAIVisionPredictor:
+    def __init__(self, labels: List[str], model_name: str = "gpt-4.1-mini") -> None:
+        self.labels = labels
+        self.model_name = model_name
+        self.api_key = os.getenv("OPENAI_API_KEY", "")
+        self.client: Optional[OpenAI] = None
+        if self.api_key:
+            self.client = OpenAI(api_key=self.api_key)
+
+    def available(self) -> bool:
+        return self.client is not None
+
+    def predict(self, image: Image.Image) -> Dict[str, object]:
+        if not self.available():
+            return {
+                "model": "openai-vision",
+                "available": False,
+                "error": "OPENAI_API_KEY is not set.",
+            }
+
+        buffered = io.BytesIO()
+        image.convert("RGB").save(buffered, format="JPEG")
+        b64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        prompt = (
+            "You are an image classifier. "
+            f"Choose exactly one label from this list: {', '.join(self.labels)}. "
+            "Return strict JSON with keys: label, confidence, reason. "
+            "label must be one of the provided labels. confidence must be in [0,1]."
+        )
+
+        response = self.client.responses.create(
+            model=self.model_name,
+            input=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "input_text", "text": prompt},
+                        {"type": "input_image", "image_url": f"data:image/jpeg;base64,{b64_image}"},
+                    ],
+                }
+            ],
+            temperature=0,
+        )
+
+        text = response.output_text.strip()
+        parsed = self._safe_parse(text)
+        return {
+            "model": "openai-vision",
+            "available": True,
+            "top_prediction": {
+                "label": parsed.get("label", "unknown"),
+                "confidence": float(parsed.get("confidence", 0.0)),
+            },
+            "raw_response": parsed,
+        }
+
+    @staticmethod
+    def _safe_parse(text: str) -> Dict[str, object]:
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            return {"label": "unknown", "confidence": 0.0, "reason": text}

src/train_custom_model.py

@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from sklearn.metrics import classification_report
+from torch.utils.data import DataLoader, random_split
+from torchvision import datasets, models, transforms
+
+
+@dataclass
+class TrainConfig:
+    data_dir: str = "data/pokemon"
+    output_model_path: str = "models/custom_resnet18.pth"
+    output_metrics_path: str = "reports/custom_model_metrics.json"
+    batch_size: int = 16
+    num_epochs: int = 8
+    learning_rate: float = 1e-4
+    weight_decay: float = 1e-4
+    val_split: float = 0.2
+    image_size: int = 224
+    seed: int = 42
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def get_device() -> torch.device:
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def build_transforms(image_size: int) -> Tuple[transforms.Compose, transforms.Compose]:
+    train_tfms = transforms.Compose(
+        [
+            transforms.RandomResizedCrop(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+    eval_tfms = transforms.Compose(
+        [
+            transforms.Resize(int(image_size * 1.14)),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+    return train_tfms, eval_tfms
+
+
+def build_dataloaders(config: TrainConfig) -> Tuple[DataLoader, DataLoader, DataLoader, List[str]]:
+    train_tfms, eval_tfms = build_transforms(config.image_size)
+
+    train_root = Path(config.data_dir) / "train"
+    test_root = Path(config.data_dir) / "test"
+
+    full_train_dataset = datasets.ImageFolder(train_root, transform=train_tfms)
+    eval_train_dataset = datasets.ImageFolder(train_root, transform=eval_tfms)
+    test_dataset = datasets.ImageFolder(test_root, transform=eval_tfms)
+
+    num_train = len(full_train_dataset)
+    num_val = int(num_train * config.val_split)
+    num_train_final = num_train - num_val
+
+    train_subset, val_subset_indices = random_split(
+        full_train_dataset,
+        [num_train_final, num_val],
+        generator=torch.Generator().manual_seed(config.seed),
+    )
+
+    # Recreate val subset with eval transforms by reusing indices from split.
+    val_indices = val_subset_indices.indices
+    val_subset = torch.utils.data.Subset(eval_train_dataset, val_indices)
+
+    train_loader = DataLoader(train_subset, batch_size=config.batch_size, shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_subset, batch_size=config.batch_size, shuffle=False, num_workers=0)
+    test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=0)
+
+    classes = full_train_dataset.classes
+    return train_loader, val_loader, test_loader, classes
+
+
+def create_model(num_classes: int) -> nn.Module:
+    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+    in_features = model.fc.in_features
+    model.fc = nn.Linear(in_features, num_classes)
+    return model
+
+
+def evaluate(model: nn.Module, loader: DataLoader, device: torch.device) -> Tuple[float, List[int], List[int]]:
+    model.eval()
+    correct = 0
+    total = 0
+    preds_all: List[int] = []
+    labels_all: List[int] = []
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.to(device)
+            outputs = model(images)
+            _, preds = torch.max(outputs, dim=1)
+
+            total += labels.size(0)
+            correct += (preds == labels).sum().item()
+            preds_all.extend(preds.cpu().numpy().tolist())
+            labels_all.extend(labels.cpu().numpy().tolist())
+
+    acc = correct / total if total > 0 else 0.0
+    return acc, preds_all, labels_all
+
+
+def train(config: TrainConfig) -> Dict[str, object]:
+    set_seed(config.seed)
+    device = get_device()
+
+    train_loader, val_loader, test_loader, classes = build_dataloaders(config)
+
+    model = create_model(num_classes=len(classes)).to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5)
+
+    best_val_acc = 0.0
+    best_state = None
+    history = []
+
+    for epoch in range(config.num_epochs):
+        model.train()
+        running_loss = 0.0
+        running_correct = 0
+        running_total = 0
+
+        for images, labels in train_loader:
+            images = images.to(device)
+            labels = labels.to(device)
+
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * labels.size(0)
+            _, preds = torch.max(outputs, dim=1)
+            running_correct += (preds == labels).sum().item()
+            running_total += labels.size(0)
+
+        scheduler.step()
+
+        train_loss = running_loss / max(1, running_total)
+        train_acc = running_correct / max(1, running_total)
+        val_acc, _, _ = evaluate(model, val_loader, device)
+
+        history.append(
+            {
+                "epoch": epoch + 1,
+                "train_loss": round(train_loss, 5),
+                "train_acc": round(train_acc, 5),
+                "val_acc": round(val_acc, 5),
+            }
+        )
+
+        print(
+            f"Epoch {epoch + 1}/{config.num_epochs} "
+            f"- train_loss: {train_loss:.4f} "
+            f"- train_acc: {train_acc:.4f} "
+            f"- val_acc: {val_acc:.4f}"
+        )
+
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+
+    if best_state is not None:
+        model.load_state_dict(best_state)
+
+    test_acc, test_preds, test_labels = evaluate(model, test_loader, device)
+
+    target_names = classes
+    cls_report = classification_report(
+        test_labels,
+        test_preds,
+        target_names=target_names,
+        output_dict=True,
+        zero_division=0,
+    )
+
+    output_model = Path(config.output_model_path)
+    output_model.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "state_dict": model.state_dict(),
+            "labels": classes,
+            "image_size": config.image_size,
+            "architecture": "resnet18",
+        },
+        output_model,
+    )
+
+    result = {
+        "config": asdict(config),
+        "device": str(device),
+        "num_classes": len(classes),
+        "labels": classes,
+        "best_val_acc": round(best_val_acc, 5),
+        "test_acc": round(test_acc, 5),
+        "history": history,
+        "classification_report": cls_report,
+        "model_path": str(output_model),
+    }
+
+    output_metrics = Path(config.output_metrics_path)
+    output_metrics.parent.mkdir(parents=True, exist_ok=True)
+    output_metrics.write_text(json.dumps(result, indent=2), encoding="utf-8")
+
+    print(f"Saved model to: {output_model}")
+    print(f"Saved metrics to: {output_metrics}")
+    print(f"Final test accuracy: {test_acc:.4f}")
+
+    return result
+
+
+def parse_args() -> TrainConfig:
+    parser = argparse.ArgumentParser(description="Train transfer learning classifier on custom image data.")
+    parser.add_argument("--data-dir", default="data/pokemon")
+    parser.add_argument("--output-model", default="models/custom_resnet18.pth")
+    parser.add_argument("--output-metrics", default="reports/custom_model_metrics.json")
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--epochs", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--val-split", type=float, default=0.2)
+    parser.add_argument("--image-size", type=int, default=224)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    return TrainConfig(
+        data_dir=args.data_dir,
+        output_model_path=args.output_model,
+        output_metrics_path=args.output_metrics,
+        batch_size=args.batch_size,
+        num_epochs=args.epochs,
+        learning_rate=args.lr,
+        weight_decay=args.weight_decay,
+        val_split=args.val_split,
+        image_size=args.image_size,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    cfg = parse_args()
+    train(cfg)

src/upload_to_hf.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import HfApi, upload_file
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Upload custom model checkpoint to Hugging Face model repo.")
+    parser.add_argument("--repo-id", required=True, help="e.g., username/pokemon-transfer-resnet18")
+    parser.add_argument("--model-path", default="models/custom_resnet18.pth")
+    parser.add_argument("--private", action="store_true")
+    args = parser.parse_args()
+
+    model_path = Path(args.model_path)
+    if not model_path.exists():
+        raise FileNotFoundError(f"Model file not found: {model_path}")
+
+    api = HfApi()
+    api.create_repo(repo_id=args.repo_id, repo_type="model", private=args.private, exist_ok=True)
+
+    upload_file(
+        path_or_fileobj=str(model_path),
+        path_in_repo=model_path.name,
+        repo_id=args.repo_id,
+        repo_type="model",
+    )
+
+    print(f"Uploaded {model_path} to https://huggingface.co/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()