Spaces:
Sleeping
Sleeping
| 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. | |