SARANGx/rtdetr-pill-detector

💊 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

Model Description

This model is a fine-tuned RT-DETR R18 (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

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)

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 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:

@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}
}