Upload 13 files

.dockerignore

@@ -0,0 +1,9 @@
+__pycache__
+.git
+.venv
+venv
+data
+training
+outputs
+checkpoints
+*.pyc

.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+.env
+model/
+outputs/
+checkpoints/
+data/*/images/*
+data/*/annotations.json

Dockerfile

@@ -0,0 +1,23 @@
+FROM python:3.11-slim
+
+ENV PYTHONUNBUFFERED=1
+ENV PIP_NO_CACHE_DIR=1
+ENV PORT=7860
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libglib2.0-0 \
+    libgl1 \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --upgrade pip && pip install -r requirements.txt
+
+COPY app ./app
+COPY model ./model
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,13 +1,199 @@
----
-title: IcecreamDetect
-emoji: 📚
-colorFrom: indigo
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.25.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Ice Cream Freezer Counter — RT-DETR + Hugging Face
+
+A complete object-detection project for counting individual ice creams in freezer photos and identifying each product type.
+
+**No YOLO is used.** The detector is Hugging Face Transformers' RT-DETR (`PekingU/rtdetr_r50vd`), fine-tuned on your own annotated freezer images.
+
+## What it does
+
+Upload a freezer photo:
+
+- detects every visible ice cream
+- classifies each detected item (for example `cornetto`, `magnum`, `cone`, `cup`)
+- counts items per class
+- returns bounding boxes and confidence scores
+- optionally returns an annotated image
+
+Example response:
+
+```json
+{
+  "total": 17,
+  "counts": {
+    "cornetto": 7,
+    "magnum": 5,
+    "cone": 3,
+    "cup": 2
+  },
+  "detections": [
+    {
+      "class": "cornetto",
+      "confidence": 0.96,
+      "box": [120, 85, 210, 310]
+    }
+  ]
+}
+```
+
+## 1. Define your classes
+
+Edit `training/classes.txt`.
+
+Example:
+
+```text
+cornetto
+magnum
+cone
+cup
+sandwich
+stick
+other
+```
+
+Use the exact product names you want returned by the API. If "correto" is your actual product spelling, use `correto`; otherwise use the correct brand/product name.
+
+## 2. Annotate your images
+
+Use an annotation tool such as CVAT, Label Studio, Roboflow, or another COCO-compatible tool.
+
+Annotate **each individual ice cream**, not just the freezer shelf.
+
+Export in **COCO object-detection format**.
+
+Expected layout:
+
+```text
+data/
+  train/
+    images/
+      0001.jpg
+      0002.jpg
+    annotations.json
+  val/
+    images/
+      1001.jpg
+      1002.jpg
+    annotations.json
+```
+
+The `annotations.json` files must be COCO detection JSON.
+
+Important:
+- draw one bounding box around each visible product
+- assign the correct product class
+- include partially visible products if you want them counted
+- keep train and validation images separate
+- avoid putting near-duplicate photos in both sets
+
+A useful starting target is 500–2,000+ annotated images. More variation is usually more valuable than many almost-identical photos.
+
+## 3. Install
+
+Python 3.10+ is recommended.
+
+```bash
+pip install -r requirements.txt
+```
+
+For NVIDIA GPU training, install the appropriate PyTorch build for your CUDA version first, following the official PyTorch instructions.
+
+## 4. Train
+
+Make sure the dataset is in the structure above, then:
+
+```bash
+python training/train.py \
+  --train-dir data/train \
+  --val-dir data/val \
+  --output-dir model \
+  --epochs 30 \
+  --batch-size 2 \
+  --learning-rate 1e-5
+```
+
+For a small GPU, use `--batch-size 1`.
+
+The script starts from `PekingU/rtdetr_r50vd` and replaces the classification head with your product classes.
+
+At the end, `model/` contains the fine-tuned Hugging Face model and processor.
+
+## 5. Test locally
+
+```bash
+python -m uvicorn app.main:app --host 0.0.0.0 --port 7860
+```
+
+Open:
+
+```text
+http://localhost:7860/docs
+```
+
+Use `POST /predict` to upload an image.
+
+Or from Python:
+
+```python
+import requests
+
+with open("freezer.jpg", "rb") as f:
+    r = requests.post(
+        "http://localhost:7860/predict",
+        files={"file": ("freezer.jpg", f, "image/jpeg")}
+    )
+
+print(r.json())
+```
+
+## 6. Environment variables
+
+The API supports:
+
+```text
+MODEL_ID=./model
+CONFIDENCE_THRESHOLD=0.35
+MAX_IMAGE_MB=15
+```
+
+If the model is stored on the Hugging Face Hub instead of the local filesystem:
+
+```text
+MODEL_ID=your-username/your-icecream-rtdetr
+HF_TOKEN=hf_...
+```
+
+## 7. Deploy to Hugging Face
+
+Create a Hugging Face **Docker Space** and upload this repository.
+
+The included `Dockerfile` listens on port `7860`.
+
+For a private model, add a Space secret named:
+
+```text
+HF_TOKEN
+```
+
+and set:
+
+```text
+MODEL_ID=your-username/your-icecream-rtdetr
+```
+
+If the model is public, `HF_TOKEN` is not required.
+
+## Accuracy notes
+
+Counting accuracy depends heavily on the training data. Freezer images have difficult cases: occlusion, reflections, small objects, similar packaging, tilted products, and products stacked behind each other.
+
+Start with a confidence threshold around `0.35–0.50` and tune it on validation photos.
+
+If two products are touching or one is heavily hidden, the detector can still miss or merge them. Add examples of those exact situations to the training set.
+
+## Why RT-DETR?
+
+RT-DETR is a transformer-based object detector available in Hugging Face Transformers. It produces object boxes and class labels and is suitable for instance counting. This project deliberately does not use YOLO.
+
+Hugging Face reference:
+https://huggingface.co/docs/transformers/model_doc/rt_detr

README_HF_SPACE.md

@@ -0,0 +1,37 @@
+# Hugging Face Space configuration
+
+Create a **Docker Space** and upload the project files.
+
+The Space should contain:
+
+```text
+Dockerfile
+requirements.txt
+app/
+model/
+```
+
+For a model stored on the Hugging Face Hub, you can instead omit the local `model/` directory and set the Space variables:
+
+```text
+MODEL_ID=YOUR_USERNAME/YOUR_MODEL_REPO
+HF_TOKEN=hf_your_token_if_private
+```
+
+The API starts on port 7860.
+
+Endpoints:
+
+- `GET /health`
+- `POST /predict`
+- `GET /docs`
+
+Example:
+
+```bash
+curl -X POST \
+  -F "file=@freezer.jpg" \
+  "https://YOUR-SPACE-NAME.hf.space/predict?return_image=true"
+```
+
+The `annotated_image_base64` field contains the annotated result image.

app/main.py

@@ -0,0 +1,140 @@
+import base64
+import io
+import os
+from collections import Counter
+
+import torch
+from fastapi import FastAPI, File, HTTPException, UploadFile
+from fastapi.responses import JSONResponse
+from PIL import Image, ImageDraw, ImageFont
+from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
+
+
+MODEL_ID = os.getenv("MODEL_ID", "./model")
+CONFIDENCE_THRESHOLD = float(os.getenv("CONFIDENCE_THRESHOLD", "0.35"))
+MAX_IMAGE_MB = int(os.getenv("MAX_IMAGE_MB", "15"))
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
+model = RTDetrForObjectDetection.from_pretrained(MODEL_ID)
+model.to(DEVICE)
+model.eval()
+
+app = FastAPI(
+    title="Ice Cream Counter API",
+    version="1.0.0",
+    description="RT-DETR ice cream product detector and counter. No YOLO.",
+)
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "model": MODEL_ID,
+        "device": str(DEVICE),
+        "threshold": CONFIDENCE_THRESHOLD,
+    }
+
+
+def annotate(image, detections):
+    image = image.copy()
+    draw = ImageDraw.Draw(image)
+
+    try:
+        font = ImageFont.truetype("DejaVuSans.ttf", 18)
+    except Exception:
+        font = ImageFont.load_default()
+
+    for d in detections:
+        x1, y1, x2, y2 = d["box"]
+        label = f'{d["class"]} {d["confidence"]:.2f}'
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+
+        bbox = draw.textbbox((x1, y1), label, font=font)
+        draw.rectangle(bbox, fill="red")
+        draw.text((x1, y1), label, fill="white", font=font)
+
+    return image
+
+
+@app.post("/predict")
+async def predict(
+    file: UploadFile = File(...),
+    return_image: bool = True,
+):
+    if not file.content_type or not file.content_type.startswith("image/"):
+        raise HTTPException(400, "Upload a JPG, PNG, WEBP, or other image file.")
+
+    raw = await file.read()
+
+    if len(raw) > MAX_IMAGE_MB * 1024 * 1024:
+        raise HTTPException(
+            413,
+            f"Image is too large. Maximum is {MAX_IMAGE_MB} MB.",
+        )
+
+    try:
+        image = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception as exc:
+        raise HTTPException(400, f"Could not read image: {exc}")
+
+    inputs = processor(images=image, return_tensors="pt")
+    inputs = {
+        k: v.to(DEVICE) if torch.is_tensor(v) else v
+        for k, v in inputs.items()
+    }
+
+    with torch.inference_mode():
+        outputs = model(**inputs)
+
+    target_sizes = torch.tensor(
+        [[image.height, image.width]],
+        device=DEVICE,
+    )
+
+    result = processor.post_process_object_detection(
+        outputs,
+        threshold=CONFIDENCE_THRESHOLD,
+        target_sizes=target_sizes,
+    )[0]
+
+    detections = []
+    counts = Counter()
+
+    for score, label, box in zip(
+        result["scores"],
+        result["labels"],
+        result["boxes"],
+    ):
+        score_value = float(score.item())
+        label_id = int(label.item())
+        class_name = model.config.id2label[label_id]
+
+        coords = [round(float(x), 2) for x in box.tolist()]
+
+        detections.append(
+            {
+                "class": class_name,
+                "confidence": round(score_value, 4),
+                "box": coords,
+            }
+        )
+        counts[class_name] += 1
+
+    response = {
+        "total": len(detections),
+        "counts": dict(sorted(counts.items())),
+        "detections": detections,
+    }
+
+    if return_image:
+        annotated = annotate(image, detections)
+        buf = io.BytesIO()
+        annotated.save(buf, format="JPEG", quality=90)
+        response["annotated_image_base64"] = base64.b64encode(
+            buf.getvalue()
+        ).decode("ascii")
+
+    return JSONResponse(response)

data/train/PUT_IMAGES_HERE.txt

@@ -0,0 +1 @@
+Put training images in this directory's images/ subdirectory and COCO annotations in annotations.json.

data/val/PUT_IMAGES_HERE.txt

@@ -0,0 +1 @@
+Put validation images in this directory's images/ subdirectory and COCO annotations in annotations.json.

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.3
+torchvision>=0.18
+transformers>=4.50
+huggingface_hub>=0.25
+Pillow>=10.0
+fastapi>=0.115
+uvicorn[standard]>=0.30
+python-multipart>=0.0.9
+numpy>=1.26
+tqdm>=4.66
+pycocotools>=2.0.8

training/classes.txt

@@ -0,0 +1,10 @@
+Carnavalita
+Kimo-COno
+Squizz
+Oreo
+Moro
+Dulce
+KitKat
+Cadbury
+Mega
+other

training/convert_coco_categories.py

@@ -0,0 +1,19 @@
+"""
+Optional helper.
+
+Your COCO exporter should already contain the same categories in train and val.
+This script prints the category IDs/names so you can compare them with classes.txt.
+
+Usage:
+    python training/convert_coco_categories.py data/train/annotations.json
+"""
+
+import json
+import sys
+
+if len(sys.argv) != 2:
+    raise SystemExit("Usage: python training/convert_coco_categories.py annotations.json")
+
+data = json.load(open(sys.argv[1], "r", encoding="utf-8"))
+for cat in sorted(data["categories"], key=lambda x: x["id"]):
+    print(f'{cat["id"]}: {cat["name"]}')

training/train.py

@@ -0,0 +1,280 @@
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from transformers import RTDetrImageProcessor, RTDetrForObjectDetection
+
+
+BASE_MODEL = "PekingU/rtdetr_r50vd"
+
+
+def load_classes(path):
+    classes = [x.strip() for x in Path(path).read_text().splitlines() if x.strip()]
+    if not classes:
+        raise ValueError("classes.txt is empty")
+    return classes
+
+
+class COCODetectionDataset(Dataset):
+    def __init__(self, image_dir, annotation_file, processor):
+        self.image_dir = Path(image_dir)
+        self.processor = processor
+
+        coco = json.loads(Path(annotation_file).read_text())
+        self.images = {x["id"]: x for x in coco["images"]}
+
+        # The training model needs contiguous labels 0..N-1.
+        # categories.json is converted here from the COCO category IDs.
+        categories = sorted(coco["categories"], key=lambda x: x["id"])
+        self.category_id_to_label = {
+            cat["id"]: i for i, cat in enumerate(categories)
+        }
+
+        anns_by_image = {}
+        for ann in coco["annotations"]:
+            if ann.get("iscrowd", 0):
+                continue
+            anns_by_image.setdefault(ann["image_id"], []).append(ann)
+
+        self.records = []
+        for image_id, info in self.images.items():
+            self.records.append(
+                {
+                    "image_id": image_id,
+                    "file_name": info["file_name"],
+                    "width": info["width"],
+                    "height": info["height"],
+                    "annotations": anns_by_image.get(image_id, []),
+                }
+            )
+
+    def __len__(self):
+        return len(self.records)
+
+    def __getitem__(self, idx):
+        record = self.records[idx]
+        image_path = self.image_dir / record["file_name"]
+        image = Image.open(image_path).convert("RGB")
+
+        annotations = []
+        for ann in record["annotations"]:
+            x, y, w, h = ann["bbox"]
+            if w <= 0 or h <= 0:
+                continue
+
+            label = self.category_id_to_label[ann["category_id"]]
+            annotations.append(
+                {
+                    "id": ann["id"],
+                    "image_id": record["image_id"],
+                    "category_id": label,
+                    "bbox": [x, y, w, h],
+                    "area": float(ann.get("area", w * h)),
+                    "iscrowd": 0,
+                }
+            )
+
+        target = {
+            "image_id": record["image_id"],
+            "annotations": annotations,
+        }
+
+        encoded = self.processor(
+            images=image,
+            annotations=target,
+            return_tensors="pt",
+        )
+
+        # Remove batch dimension. DataLoader will create it.
+        encoded["pixel_values"] = encoded["pixel_values"].squeeze(0)
+        if "pixel_mask" in encoded:
+            encoded["pixel_mask"] = encoded["pixel_mask"].squeeze(0)
+        encoded["labels"] = encoded["labels"][0]
+        return encoded
+
+
+def collate_fn(batch):
+    pixel_values = torch.stack([x["pixel_values"] for x in batch])
+    pixel_mask = None
+    if "pixel_mask" in batch[0]:
+        pixel_mask = torch.stack([x["pixel_mask"] for x in batch])
+
+    labels = [x["labels"] for x in batch]
+    result = {"pixel_values": pixel_values, "labels": labels}
+    if pixel_mask is not None:
+        result["pixel_mask"] = pixel_mask
+    return result
+
+
+def evaluate(model, loader, device):
+    model.eval()
+    total = 0.0
+    count = 0
+    with torch.no_grad():
+        for batch in loader:
+            batch = {
+                k: (v.to(device) if torch.is_tensor(v) else v)
+                for k, v in batch.items()
+            }
+            out = model(**batch)
+            total += float(out.loss.item())
+            count += 1
+    model.train()
+    return total / max(count, 1)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train-dir", default="data/train")
+    parser.add_argument("--val-dir", default="data/val")
+    parser.add_argument("--classes", default="training/classes.txt")
+    parser.add_argument("--output-dir", default="model")
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--learning-rate", type=float, default=1e-5)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--num-workers", type=int, default=2)
+    parser.add_argument("--grad-accumulation", type=int, default=1)
+    parser.add_argument("--resume", default=None)
+    args = parser.parse_args()
+
+    classes = load_classes(args.classes)
+    label2id = {name: i for i, name in enumerate(classes)}
+    id2label = {i: name for i, name in enumerate(classes)}
+
+    train_ann = Path(args.train_dir) / "annotations.json"
+    val_ann = Path(args.val_dir) / "annotations.json"
+
+    processor = RTDetrImageProcessor.from_pretrained(BASE_MODEL)
+
+    train_ds = COCODetectionDataset(
+        Path(args.train_dir) / "images",
+        train_ann,
+        processor,
+    )
+    val_ds = COCODetectionDataset(
+        Path(args.val_dir) / "images",
+        val_ann,
+        processor,
+    )
+
+    # Make sure both datasets use exactly the class list supplied by the user.
+    train_categories = len(train_ds.category_id_to_label)
+    val_categories = len(val_ds.category_id_to_label)
+    if train_categories != len(classes) or val_categories != len(classes):
+        raise ValueError(
+            "The number of COCO categories does not match classes.txt. "
+            f"classes.txt={len(classes)}, train={train_categories}, val={val_categories}"
+        )
+
+    model = RTDetrForObjectDetection.from_pretrained(
+        BASE_MODEL,
+        num_labels=len(classes),
+        id2label=id2label,
+        label2id=label2id,
+        ignore_mismatched_sizes=True,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=torch.cuda.is_available(),
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        pin_memory=torch.cuda.is_available(),
+    )
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.learning_rate,
+        weight_decay=args.weight_decay,
+    )
+
+    scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
+    start_epoch = 0
+
+    if args.resume:
+        checkpoint = torch.load(args.resume, map_location=device)
+        model.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        start_epoch = checkpoint["epoch"] + 1
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    best_val = float("inf")
+
+    for epoch in range(start_epoch, args.epochs):
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        running = 0.0
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{args.epochs}")
+        for step, batch in enumerate(pbar):
+            batch = {
+                k: (v.to(device) if torch.is_tensor(v) else v)
+                for k, v in batch.items()
+            }
+
+            with torch.amp.autocast(
+                device_type="cuda",
+                enabled=torch.cuda.is_available(),
+            ):
+                out = model(**batch)
+                loss = out.loss / args.grad_accumulation
+
+            scaler.scale(loss).backward()
+
+            if (step + 1) % args.grad_accumulation == 0:
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+
+            running += float(loss.item()) * args.grad_accumulation
+            pbar.set_postfix(loss=f"{running / (step + 1):.4f}")
+
+        val_loss = evaluate(model, val_loader, device)
+        print(f"epoch={epoch + 1} validation_loss={val_loss:.4f}")
+
+        checkpoint = {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+        }
+        torch.save(checkpoint, output_dir / "last_checkpoint.pt")
+
+        if val_loss < best_val:
+            best_val = val_loss
+            model.save_pretrained(output_dir)
+            processor.save_pretrained(output_dir)
+            (output_dir / "classes.json").write_text(
+                json.dumps(
+                    {"id2label": id2label, "label2id": label2id},
+                    indent=2,
+                )
+            )
+            print(f"Saved best model to {output_dir}")
+
+    # Always save the final model as well.
+    model.save_pretrained(output_dir)
+    processor.save_pretrained(output_dir)
+    print(f"Final model saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()