"""
Push trained YOLO artifacts from Modal volume to HuggingFace Hub.

Usage:
    modal run finetune/push_yolo_to_hf.py

Reads from Modal volume: kirana-yolo-output  (/output/)
Pushes to: naazimsnh02/yolo26n-indian-fmcg-detection
  - best.pt       (PyTorch weights)
  - best.onnx     (ONNX, opset 12)
  - class_names.json
  - README.md     (model card)
"""

import os
import modal

app = modal.App("kirana-push-yolo")

IMAGE = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("huggingface_hub>=0.30.0")
)

HF_SECRET = modal.Secret.from_name("hf-secret")
HF_REPO   = "naazimsnh02/yolo26n-indian-fmcg-detection"

MODEL_CARD = """\
---
license: apache-2.0
base_model: yolo26n
language:
  - en
tags:
  - object-detection
  - yolo
  - indian-fmcg
  - onnx
  - ultralytics
  - kirana
pipeline_tag: object-detection
datasets:
  - agentsk47/indian-grocery-object-detection-mfsnx
  - iit-patna-qg1jh/grocery_items-7i2em
  - project-c5ho0/indian-market-qieug
---

# YOLO26n — Indian FMCG Product Detection

Fine-tuned [YOLO26n](https://docs.ultralytics.com) on a **merged dataset of three Indian grocery sources**
from Roboflow Universe.  Part of the **Kirana Detective** project — an AI system for small Indian grocery
stores to visually count and reconcile shelf/counter inventory from photos.

## Performance

| Metric | Value |
|---|---|
| mAP50 (all classes) | **0.428** |
| mAP50-95 (all classes) | **0.302** |
| Total classes | 1,831 |
| Validation images | 1,236 |
| Validation instances | 13,443 |

Training ran for **100 epochs** (60 initial + 40 resumed after restart) on an NVIDIA A10G via Modal.

## Training Datasets

| Dataset | Workspace | Version | Images | Classes |
|---|---|---|---|---|
| [Indian Grocery Object Detection](https://universe.roboflow.com/agentsk47/indian-grocery-object-detection-mfsnx) | agentsk47 | v1 | ~400 | 10 |
| [Grocery Items](https://universe.roboflow.com/iit-patna-qg1jh/grocery_items-7i2em) | IIT Patna | v45 | 6,695 | 20 |
| [Indian Market](https://universe.roboflow.com/project-c5ho0/indian-market-qieug) | project-c5ho0 | v2 | 4,694 | 2 |

All three datasets were downloaded in **YOLOv8 format**, class IDs remapped to a unified list, and merged
before training.  The full unified class list (1,831 entries) is available in `class_names.json`.

## Files

| File | Description |
|---|---|
| `best.pt` | PyTorch checkpoint (best mAP50 epoch) |
| `best.onnx` | ONNX export, opset 12 (recommended for inference) |
| `class_names.json` | Full list of 1,831 class names (index = class_id) |

## How to Use

### ONNX Runtime (CPU / any platform)

```python
import json, numpy as np, onnxruntime as ort
from PIL import Image

session    = ort.InferenceSession("best.onnx", providers=["CPUExecutionProvider"])
class_names = json.load(open("class_names.json"))

def preprocess(path, size=640):
    img = Image.open(path).convert("RGB").resize((size, size))
    return (np.array(img, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]

input_name = session.get_inputs()[0].name
outputs    = session.run(None, {input_name: preprocess("shelf.jpg")})
# outputs[0]: (1, 300, 6) — [x1, y1, x2, y2, confidence, class_id]
```

### Ultralytics (PyTorch)

```python
from ultralytics import YOLO

model   = YOLO("best.pt")
results = model.predict("shelf.jpg", imgsz=640, conf=0.25)
results[0].show()
```

## Training Details

| Parameter | Value |
|---|---|
| Base model | YOLO26n |
| Input size | 640 × 640 |
| Epochs | 100 (60 + 40 resumed) |
| Batch size | 16 |
| Early stopping patience | 20 |
| Export format | ONNX opset 12 |
| Hardware | NVIDIA A10G (Modal) |

## Citation

```bibtex
@misc{kirana-detective-yolo-2026,
  title  = {Kirana Detective: YOLO26n Indian FMCG Product Detector},
  author = {Naazim},
  year   = {2026},
  url    = {https://huggingface.co/naazimsnh02/yolo26n-indian-fmcg-detection}
}
```
"""


@app.function(
    image=IMAGE,
    timeout=600,
    secrets=[HF_SECRET],
    volumes={"/output": modal.Volume.from_name("kirana-yolo-output", create_if_missing=False)},
)
def push_to_hub():
    import json
    import shutil
    import tempfile
    from pathlib import Path
    from huggingface_hub import HfApi

    # --- Locate artifacts ---
    output    = Path("/output")
    best_pt   = output / "runs/yolo26n_fmcg/weights/best.pt"
    best_onnx = output / "runs/yolo26n_fmcg/weights/best.onnx"
    cls_json  = output / "class_names.json"

    print("=== Volume contents (/output) ===")
    for p in sorted(output.rglob("*")):
        if p.is_file():
            print(f"  {p.relative_to(output)}  ({p.stat().st_size / 1024:.1f} KB)")

    missing = [p for p in (best_pt, best_onnx, cls_json) if not p.exists()]
    if missing:
        raise FileNotFoundError(f"Missing artifacts: {[str(m) for m in missing]}")

    with open(cls_json) as f:
        classes = json.load(f)
    print(f"\nClass count: {len(classes)}")

    # --- Stage all files into a temp folder, then push as a single commit ---
    api = HfApi(token=os.environ["HF_TOKEN"])
    api.create_repo(HF_REPO, repo_type="model", exist_ok=True, private=False)

    with tempfile.TemporaryDirectory() as staging:
        staging = Path(staging)
        shutil.copy(best_pt,   staging / "best.pt")
        shutil.copy(best_onnx, staging / "best.onnx")
        shutil.copy(cls_json,  staging / "class_names.json")
        (staging / "README.md").write_text(MODEL_CARD, encoding="utf-8")

        print("\nFiles staged for upload:")
        for f in sorted(staging.iterdir()):
            print(f"  {f.name}  ({f.stat().st_size / 1024:.1f} KB)")

        print("\nPushing to HF Hub (single commit)...")
        api.upload_folder(
            folder_path=str(staging),
            repo_id=HF_REPO,
            repo_type="model",
            commit_message="Add best.pt, best.onnx, class_names.json, README (100-epoch FMCG detector)",
        )

    print(f"\nDone — https://huggingface.co/{HF_REPO}")


@app.local_entrypoint()
def main():
    push_to_hub.remote()