README.md

---
library_name: transformers
license: apache-2.0
base_model: PekingU/rtdetr_r18vd_coco_o365
tags:
  - object-detection
  - pill-detection
  - medical
  - rt-detr
  - computer-vision
  - counting
  - healthcare
datasets:
  - Francesco/pills-sxdht
pipeline_tag: object-detection
model-index:
  - name: rtdetr-pill-detector
    results:
      - task:
          type: object-detection
          name: Object Detection
        dataset:
          name: pills-sxdht (Roboflow 100)
          type: Francesco/pills-sxdht
        metrics:
          - type: loss
            value: 3.528
            name: Best Validation Loss
---

# 💊 RT-DETR Pill Detector & Counter

A real-time medicine pill detection model that can **detect and count pills, capsules, and specific medications** in images.

**🎮 Try it live:** [**Pill Detector Demo**](https://huggingface.co/spaces/SARANGx/pill-detector-demo)

## Model Description

This model is a fine-tuned [RT-DETR R18](https://huggingface.co/PekingU/rtdetr_r18vd_coco_o365) (Real-Time DEtection TRansformer with ResNet-18 backbone) for detecting medicine pills in images.

### Why RT-DETR?
- **NMS-free architecture** — Each pill is detected exactly once, making it ideal for **accurate counting**
- **End-to-end detection** — No post-processing like Non-Maximum Suppression needed
- **Real-time capable** — 20M parameters, fast inference even on CPU
- **Strong baseline** — Pre-trained on Objects365 (365 categories, 2M images)

### What it detects (9 classes)

| Class | Description |
|-------|-------------|
| `pills` | Generic pill detection |
| `Cipro 500` | Ciprofloxacin 500mg |
| `Ibuphil 600 mg` | Ibuprofen 600mg |
| `Ibuphil Cold 400-60` | Ibuprofen/Pseudoephedrine combination |
| `Xyzall 5mg` | Levocetirizine 5mg |
| `blue` | Blue-colored pills |
| `pink` | Pink-colored pills |
| `red` | Red-colored pills |
| `white` | White-colored pills |

## Usage

### Quick Start with Pipeline

```python
from transformers import pipeline
from PIL import Image

detector = pipeline("object-detection", model="SARANGx/rtdetr-pill-detector")
image = Image.open("pills.jpg")

results = detector(image, threshold=0.5)
for r in results:
    print(f"{r['label']}: {r['score']:.2%} at {r['box']}")

print(f"Total pills: {len(results)}")
```

### Manual Inference (more control)

```python
import torch
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
from PIL import Image
from collections import Counter

model_id = "SARANGx/rtdetr-pill-detector"
device = "cuda" if torch.cuda.is_available() else "cpu"

image_processor = RTDetrImageProcessor.from_pretrained(model_id)
model = RTDetrForObjectDetection.from_pretrained(model_id).to(device).eval()

image = Image.open("pills.jpg").convert("RGB")
inputs = image_processor(images=image, return_tensors="pt").to(device)

with torch.no_grad():
    outputs = model(**inputs)

# Post-process — boxes in original image coordinates
target_sizes = torch.tensor([(image.height, image.width)], device=device)
results = image_processor.post_process_object_detection(
    outputs, target_sizes=target_sizes, threshold=0.5
)[0]

# Count pills by class
counts = Counter()
for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
    label = model.config.id2label[label_id.item()]
    counts[label] += 1
    x1, y1, x2, y2 = box.tolist()
    print(f"  {label}: {score:.2%} at [{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]")

print(f"\nTotal pills detected: {sum(counts.values())}")
for label, count in sorted(counts.items(), key=lambda x: -x[1]):
    print(f"  {label}: {count}")
```

## Training Details

### Dataset
- **[Francesco/pills-sxdht](https://huggingface.co/datasets/Francesco/pills-sxdht)** from Roboflow 100
- 316 training / 45 validation / 90 test images
- All images 640×640 with COCO-format bounding box annotations
- 9 object classes (pills + specific medications + color categories)

### Training Configuration
| Parameter | Value |
|-----------|-------|
| Base model | PekingU/rtdetr_r18vd_coco_o365 (Objects365 pretrained) |
| Image size | 480×480 |
| Epochs | 20 |
| Batch size | 8 |
| Learning rate | 5e-5 |
| LR scheduler | Cosine with 50 warmup steps |
| Optimizer | AdamW (fused) |
| Max grad norm | 0.1 |
| Augmentations | HorizontalFlip, ColorJitter, RandomBrightnessContrast, GaussNoise, Blur |

### Training Loss Curve

| Epoch | Train Loss | Eval Loss |
|-------|-----------|-----------|
| 1 | 36.61 | 25.70 |
| 5 | 8.57 | 5.21 |
| 10 | 6.32 | 3.79 |
| 15 | 5.55 | 3.53 |
| **19** | **5.31** | **3.53 (best)** |
| 20 | 5.78 | 3.58 |

Best validation loss: **3.528** at epoch 19 (loaded as final checkpoint).

### Technical Notes
- Classification head re-initialized from 80 COCO classes → 9 pill classes
- Auxiliary loss enabled for training stability
- `freeze_backbone_batch_norms=True` to preserve pretrained backbone statistics
- Focal loss (α=0.75, γ=2.0) for handling class imbalance

## Limitations
- Trained on a small dataset (316 images) — may not generalize well to all pill types
- Best on images similar to training data (top-down views, clean backgrounds)
- Color-based classes (blue, pink, red, white) may overlap with medication-specific classes
- Not intended for medical decision-making — for counting/inventory purposes only

## Framework Versions
- Transformers 5.6.1
- PyTorch 2.11.0
- Datasets 4.8.4

## Citation

If you use this model, please cite the underlying RT-DETR architecture:

```bibtex
@article{zhao2024detrs,
  title={DETRs Beat YOLOs on Real-time Object Detection},
  author={Zhao, Yian and Lv, Wenyu and Xu, Shangliang and Wei, Jinman and Wang, Guanzhong and Dang, Qingqing and Liu, Yi and Chen, Jie},
  journal={CVPR},
  year={2024}
}
```