---
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](#hsv-based-color-recognition).
- **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`](app/utils/pill_detection.py)
(`process_image`), with HTTP routing in [`app/route.py`](app/route.py).

---

## HSV-based Color Recognition

The color module ([`app/utils/shape_color_utils.py`](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

```bash
# 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.