README.md

metadata

title: Pill Identification System
emoji: 💊
colorFrom: green
colorTo: blue
sdk: docker
app_port: 7860
pinned: false

Pill Identification System (OCR + HSV Color Recognition)

A computer-vision pipeline that identifies medical pills from a single photo by combining YOLO pill detection, multi-angle OCR imprint reading, an HSV-based color-recognition system, and shape analysis, then matches the result against a drug database and returns the most likely medications with their indications, warnings, and side effects.

Built as a Flask web service with a browser front-end for live capture and lookup.

---

Highlights

• YOLOv8 pill detection — locates and crops the pill, with a low-confidence retry pass so faint or low-contrast pills are still found.
• Multi-angle OCR — the imprint is read at 8 rotations (0–315°) through OpenOCR, and the best result is selected by a combined length × confidence score.
• HSV-based color recognition — a robust, contour-aware color extractor that is resilient to lighting, shadows, specular highlights, and imprint text. See below.
• Shape classification — circle / ellipse / other via ellipse-fit ratios with shadow-corrected segmentation.
• Top-N database matching — OCR text is matched front/back against the drug database with an LCS-based scorer, returning ranked candidates (with a low-confidence fallback and a "retake the photo" path for unusable images).

---

System Architecture

        photo (base64 / upload)
                 │
                 ▼
      ┌───────────────────────┐
      │  YOLOv8 pill detection │   conf 0.25 → retry 0.10
      └───────────┬───────────┘
                  │ cropped pill
        ┌─────────┴───────────────────────────┐
        ▼                                       ▼
┌─────────────────────┐              ┌─────────────────────────┐
│  Multi-angle OCR     │              │  HSV color recognition   │
│  (8 rotations,       │              │  + shape classification  │
│   OpenOCR, best pick)│              │  (contour-aware, robust) │
└──────────┬──────────┘              └────────────┬────────────┘
           │ imprint text                          │ color, shape
           └───────────────┬──────────────────────┘
                           ▼
              ┌─────────────────────────┐
              │  Top-N front/back match  │  LCS scorer + color/shape filter
              │  against drug database   │
              └────────────┬────────────┘
                           ▼
              ranked medications + drug info

The pipeline lives in app/utils/pill_detection.py (process_image), with HTTP routing in app/route.py.

---

HSV-based Color Recognition

The color module (app/utils/shape_color_utils.py) uses a robust HSV pipeline rather than a naïve KMeans/RGB average:

1. Contour-aware masking — analyze only pill pixels, not the background.
2. Noise rejection — drop the darkest 15% (shadow), specular highlights, a 5% border band, and dilated imprint pixels.
3. Robust statistics — adaptive hue histogram + IQR-trimmed medians for S and V, so a few stray pixels can't skew the result.
4. Semantic classification — classify_hsv_to_semantic_color() maps HSV to human color labels with lighting-tuned thresholds, plus a color-tolerance table for fuzzy matching (is_color_match_multi, get_color_tolerance).

Key functions: extract_pill_colors_hsv, classify_hsv_to_semantic_color, detect_shape_and_extract_colors, score_candidate.

---

Tech Stack

Python · Flask · OpenCV · NumPy · Ultralytics YOLOv8 · OpenOCR (ONNX Runtime) · PyTorch · scikit-learn · Pandas

---

Project Structure

.
├── app/
│   ├── route.py                 # Flask routes (/upload, /match, /api/*)
│   ├── utils/
│   │   ├── pill_detection.py    # YOLO detect → OCR → color/shape pipeline
│   │   ├── shape_color_utils.py # HSV color recognition + shape detection
│   │   ├── matcher.py           # LCS-based top-N OCR matching
│   │   ├── ocr_utils.py         # OpenOCR wrapper
│   │   ├── data_loader.py       # drug database loading
│   │   └── image_io.py          # safe image reading (HEIC, etc.)
│   ├── templates/ static/       # web front-end
│   └── __init__.py              # create_app()
├── data/                        # drug database (TESTData.xlsx) + pictures/
├── models/                      # best.pt (YOLO) + OCR onnx (auto-downloaded)
├── main.py                      # app entry point
├── setup_models.py              # downloads OCR models
└── requirements.txt

---

Getting Started

# 1. Install dependencies
pip install -r requirements.txt

# 2. Download the OCR models (best.pt YOLO weights are already included)
python setup_models.py

# 3. Run
python main.py
# server starts on http://localhost:10000

API

Endpoint	Method	Purpose
/	GET	Web UI
/upload	POST	Detect + OCR + color/shape from a base64 image
/match	POST	Match {texts, colors, shape} against the drug DB
/api/status	GET	Health / data-loaded status

---

Notes on Data & Models

• models/best.pt (YOLO detector) ships with the repo.
• OCR models are excluded from git and fetched by setup_models.py.
• data/pictures/ (the ~534 MB drug image database) is git-ignored to keep the repo lean — remove that line from .gitignore if you want to publish it, or drop your own images there.