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