"""
Upload locally-downloaded YOLO artifacts to HuggingFace Hub.

Prerequisites:
    pip install huggingface_hub
    Set HF_TOKEN env var  (or run `huggingface-cli login`)

Usage:
    python finetune/upload_yolo_to_hf.py

Uploads from:  model_artifacts/yolo26n_fmcg/
    best.pt, best.onnx, class_names.json  →  naazimsnh02/yolo26n-indian-fmcg-detection
"""

import os
from pathlib import Path

HF_REPO    = "naazimsnh02/yolo26n-indian-fmcg-detection"
ARTIFACTS  = Path(__file__).parent.parent / "model_artifacts" / "yolo26n_fmcg"

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 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 initial + 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}
}
```
"""


def main():
    from huggingface_hub import HfApi

    token = os.environ.get("HF_TOKEN")
    if not token:
        raise EnvironmentError("HF_TOKEN env var not set. Run: set HF_TOKEN=hf_...")

    api = HfApi(token=token)

    files = {
        "best.pt":          ARTIFACTS / "best.pt",
        "best.onnx":        ARTIFACTS / "best.onnx",
        "class_names.json": ARTIFACTS / "class_names.json",
    }
    for name, path in files.items():
        if not path.exists():
            raise FileNotFoundError(f"Missing: {path}")

    print(f"Repo: {HF_REPO}")
    api.create_repo(HF_REPO, repo_type="model", exist_ok=True, private=False)

    # Upload all files in a single commit to avoid the "no changes" skip bug
    from huggingface_hub import CommitOperationAdd

    operations = []
    for repo_path, local_path in files.items():
        size_mb = local_path.stat().st_size / 1024 / 1024
        print(f"  Staging {repo_path}  ({size_mb:.1f} MB)")
        operations.append(CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=str(local_path)))

    operations.append(
        CommitOperationAdd(
            path_in_repo="README.md",
            path_or_fileobj=MODEL_CARD.encode("utf-8"),
        )
    )
    print("  Staging README.md")

    print("\nCommitting...")
    commit = api.create_commit(
        repo_id=HF_REPO,
        repo_type="model",
        operations=operations,
        commit_message="Add best.pt, best.onnx, class_names.json, README (100-epoch FMCG detector)",
    )
    print(f"Done — {commit.commit_url}")


if __name__ == "__main__":
    main()