This model is a fine-tuned
RT-DETR R18 (Real-Time DEtection TRansformer with ResNet-18 backbone) for detecting medicine pills in images.
1from transformers import pipeline
2from PIL import Image
3
4detector = pipeline("object-detection", model="SARANGx/rtdetr-pill-detector")
5image = Image.open("pills.jpg")
6
7results = detector(image, threshold=0.5)
8for r in results:
9 print(f"{r['label']}: {r['score']:.2%} at {r['box']}")
10
11print(f"Total pills: {len(results)}")
1import torch
2from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
3from PIL import Image
4from collections import Counter
5
6model_id = "SARANGx/rtdetr-pill-detector"
7device = "cuda" if torch.cuda.is_available() else "cpu"
8
9image_processor = RTDetrImageProcessor.from_pretrained(model_id)
10model = RTDetrForObjectDetection.from_pretrained(model_id).to(device).eval()
11
12image = Image.open("pills.jpg").convert("RGB")
13inputs = image_processor(images=image, return_tensors="pt").to(device)
14
15with torch.no_grad():
16 outputs = model(**inputs)
17
18# Post-process — boxes in original image coordinates
19target_sizes = torch.tensor([(image.height, image.width)], device=device)
20results = image_processor.post_process_object_detection(
21 outputs, target_sizes=target_sizes, threshold=0.5
22)[0]
23
24# Count pills by class
25counts = Counter()
26for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
27 label = model.config.id2label[label_id.item()]
28 counts[label] += 1
29 x1, y1, x2, y2 = box.tolist()
30 print(f" {label}: {score:.2%} at [{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]")
31
32print(f"\nTotal pills detected: {sum(counts.values())}")
33for label, count in sorted(counts.items(), key=lambda x: -x[1]):
34 print(f" {label}: {count}")
1@article{zhao2024detrs,
2 title={DETRs Beat YOLOs on Real-time Object Detection},
3 author={Zhao, Yian and Lv, Wenyu and Xu, Shangliang and Wei, Jinman and Wang, Guanzhong and Dang, Qingqing and Liu, Yi and Chen, Jie},
4 journal={CVPR},
5 year={2024}
6}