A dataset of medication images with instance segmentation masks for preventing adverse drug events

1. Problem Statement

Medication errors and Adverse Drug Events (ADEs) pose significant risks to patient safety, particularly among the elderly and those with complex medication regimens. While AI-driven pill recognition offers a solution, its development is hindered by the lack of comprehensive, real-world datasets.

• Limitations of Existing Datasets: Current public datasets (e.g., NIH Pillbox, CURE) primarily feature single, isolated pills captured under controlled lighting. They lack:
  • Instance Segmentation: Most provide only bounding boxes or incomplete segmentation masks.
  • Real-World Complexity: They fail to capture overlapping pills, partial occlusions, cluttered environments (e.g., dosette boxes), and variable lighting conditions.
  • Generalization: Models trained on these datasets struggle to recognize unseen pill classes in multi-pill, occluded scenarios.

2. Methodology

The authors developed MEDISEG, a new dataset designed to reflect realistic medication handling scenarios.

Data Acquisition & Construction

• Hardware/Setup: Images were captured using an iPhone 12 Pro Max. A standard 4x7 dosette box was used to arrange pills, introducing natural occlusions and overlaps.
• Conditions: Artificial lighting was manipulated (intensity and angle) to create shadows, reflections, and highlights.
• Preprocessing: Images were cropped to isolate pill slots, padded, and resized to 640×640 pixels.
• Annotation:
  • Tool: COCO Annotator.
  • Format: Instance segmentation masks (polygons) were manually drawn for every pill, even in overlapping or occluded states.
  • Quality Control: Strict inter-annotator guidelines and a dual-review process (initial annotator + secondary reviewer) ensured high accuracy.

Dataset Composition

The dataset is divided into two subsets:

1. MEDISEG (3-Pills): Contains 3 distinct pill types designed to challenge recognition algorithms based on subtle shape and color differences (e.g., similar shapes, similar colors). Includes images with 1 to 6 pills.
2. MEDISEG (32-Pills): A comprehensive subset with 32 distinct pill classes featuring diverse shapes, colors, and sizes. It includes fine-grained, visually confusable cases (e.g., identical white tablets) and scenes with up to 13 pills per frame.

• Total Scale: 8,262 images covering diverse conditions (single pills, cluttered boxes, varying orientations).

Experimental Framework

• Models: State-of-the-art object detection models YOLOv8 and YOLOv9 were trained and evaluated.
• Hyperparameter Tuning: A Genetic Algorithm (GA) was used over 70 iterations to optimize learning rates, momentum, and loss weights.
• Few-Shot Learning Protocol: To test transferability, the authors used the FsDet framework.
  • Base Training: Models were pre-trained on MEDISEG (real-world complexity) vs. CURE (controlled, isolated images).
  • Fine-Tuning: Models were adapted to unseen pill classes using limited labels (1-shot, 5-shot, 10-shot).
  • Evaluation: Tested on a held-out set, specifically focusing on a "overlap-only" subset to simulate difficult real-world conditions.

3. Key Contributions

1. MEDISEG Dataset: The first large-scale dataset providing instance segmentation masks for pills in realistic, cluttered environments (dosette boxes) with overlapping and occluded objects.
2. Real-World Simulation: Unlike previous datasets, MEDISEG explicitly captures the complexity of multi-pill scenarios, variable lighting, and partial occlusions common in clinical and home settings.
3. Comprehensive Metadata: Includes regulatory data (Hong Kong drug registration numbers, ingredients, dosage strength) linked to images via metadata.csv.
4. Benchmarking: Established new benchmarks for pill recognition using YOLOv8/v9 and demonstrated the efficacy of few-shot learning in this domain.

4. Results

Supervised Learning Performance

• 3-Pills Subset: Both models achieved high performance, with mAP@50 of 99.5%. YOLOv9 showed superior localization capabilities (higher mAP@50-95) compared to YOLOv8.
• 32-Pills Subset: YOLOv9 outperformed YOLOv8 in handling complex, multi-class scenarios, achieving an mAP@50 of 80.1%. YOLOv9 demonstrated better robustness against occlusions and overlapping objects.
• Error Analysis: Misclassifications primarily occurred between visually similar pills (e.g., Pill B vs. Pill C) when viewed from side profiles, highlighting the difficulty of fine-grained discrimination.

Few-Shot Generalization (Key Finding)

• Superior Transferability: Models base-trained on MEDISEG significantly outperformed those trained on the CURE dataset when adapting to unseen pill classes.
• Overlap-Only Performance: In the challenging "overlap-only" test set:
  • 1-shot: MEDISEG-trained model achieved 0.406 foreground classification accuracy vs. 0.131 for the CURE-trained model.
  • 5-shot: 0.625 (MEDISEG) vs. 0.372 (CURE).
  • 10-shot: 0.740 (MEDISEG) vs. 0.558 (CURE).
• Conclusion: Exposure to realistic multi-object interactions during base training creates more transferable representations, drastically reducing false negatives in occluded scenarios.

5. Significance

• Patient Safety: By enabling AI systems to accurately identify pills in cluttered, real-world environments (like dosette boxes), MEDISEG supports the development of tools that can prevent medication errors, incorrect dosages, and adverse drug reactions.
• Advancement of Computer Vision: The dataset addresses a critical gap in instance segmentation for small, visually similar objects in occluded settings, serving as a robust benchmark for fine-grained object detection.
• Clinical Applicability: The few-shot results suggest that MEDISEG-based pretraining allows healthcare systems to rapidly adapt to new drug formulations with minimal labeled data, a crucial requirement for dynamic clinical environments.
• Open Science: The dataset is publicly available under a CC BY 4.0 license, accompanied by code (Jupyter Notebooks) and detailed metadata, fostering reproducibility and further research in medication safety AI.

Availability:

• Dataset: https://doi.org/10.25383/city.28574786.v1
• Code: https://github.com/williamcwi/MEDISEG