Update model card with full training details and usage

993fb4c verified 5 months ago

5.75 kB

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

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