README.md

---
license: mit
library_name: pytorch
tags:
  - instance-segmentation
  - mask-rcnn
  - object-detection
  - resnet50
  - fpn
  - medical
  - ocr
  - pytorch
  - safetensors
pipeline_tag: image-segmentation
---

<p align="center">
  <img src="assets/banner.png" alt="MediSeG banner" width="100%">
</p>

# MediSeG — Mask R-CNN for Medicine Box Instance Segmentation

Custom **Mask R-CNN** (ResNet-50 + FPN + RPN + ROI Align, implemented from scratch in PyTorch — no `torchvision.models.detection`) that detects and segments medicine packages (`medicine_box`) in photos.

The model is the first stage of the MediSeG pipeline: **detection + mask → masked crop → OCR (French + Arabic) → LLM identification → web search → annotated summary**.

<p align="center">
  <img src="assets/examples/01_multi_box_detections.jpg" alt="Mask R-CNN detections on three medicine boxes" width="49%">
  <img src="assets/examples/01_multi_box_annotated.jpg" alt="Full pipeline output: each box identified with name, dosage, form and active ingredient" width="49%">
</p>
<p align="center"><em>Left: Mask R-CNN output (boxes + masks + scores). Right: full pipeline output (OCR + LLM + web search).</em></p>

## Model description

- **Backbone**: ResNet-50 (ImageNet-pretrained), feature maps C2–C5; the top 3 ResNet stages are fine-tuned, BatchNorm stays frozen.
- **Neck**: Feature Pyramid Network (FPN), 5 levels, 256 channels.
- **Region proposal**: RPN, one anchor size per level (32 / 64 / 128 / 256 / 512) × 3 aspect ratios (0.5, 1.0, 2.0); NMS 0.7, 300 proposals kept at test time.
- **ROI heads**: two-layer MLP trunk → classification head + class-specific box regression head (7×7 ROI Align); 4-conv + deconv mask head (14×14 ROI Align → 28×28 masks).
- **Classes**: 2 (`0 = background`, `1 = medicine_box`).
- **Input**: RGB, letterboxed to 640×640 (aspect ratio preserved, black padding, top-left aligned).
- **Weights**: `checkpoints/best_model.safetensors` in this repository.

## Intended use

Locate and segment medicine packages (boxes, tubes, sachets) in photos, typically to crop them for OCR — e.g. pharmacy tooling or assistive apps that read a medicine's label.

**Not** validated for clinical or safety-critical use. The downstream LLM identification can be wrong (see [Limitations](#limitations)); never rely on it to decide what medicine to take.

## Training data

Built from two public sources merged into one COCO instance-segmentation dataset (full report: `data/report.md` in the source repository):

| Split | Images | Instances |
|---|---:|---:|
| train | 540 | 806 |
| val | 68 | 121 |
| test | 68 | 100 |

- Sources: `main_ar_fr` (French/Arabic packages, 4-point polygons) and `medicine_packv2` (polygons). A third source with bounding boxes only was excluded.
- Split 80/10/10 with seed 42, **grouped by source photo** (near-duplicates merged by perceptual hash) to avoid train/test leakage — the original splits leaked 128 images.
- Single class: `medicine_box`.

## Training procedure

- SGD, `lr=0.01`, `momentum=0.9`, `weight_decay=1e-4`, batch size 8, 30 epochs.
- Linear warmup over 200 iterations, LR ×0.1 at epochs 20 and 26, gradient clipping at 10, mixed precision on GPU.
- Losses (equal weights): RPN objectness (BCE) + RPN box (smooth-L1), ROI classification (cross-entropy) + ROI box (smooth-L1), mask (BCE on 28×28).
- Sampling: RPN anchors positive at IoU ≥ 0.7 / negative < 0.3 (256 per image, 50 % positive); ROI proposals positive at IoU ≥ 0.5 (512 per image, 25 % positive).
- Augmentation: random horizontal flip (p = 0.5).

All hyperparameters live in `model_rcnn_scratch/config.py`.

## Evaluation

Dice and IoU on the held-out test split (union of predicted masks with score ≥ 0.5 vs. union of ground-truth masks, per image):

```bash
python -m model_rcnn_scratch.evaluate
```

| Metric | Value |
|---|---|
| Dice | _TBD_ |
| IoU  | _TBD_ |

## How to use

### Weights only

```python
import torch
import torchvision.transforms.functional as TF
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file

from model_rcnn_scratch.dataset import letterbox
from model_rcnn_scratch.model_architecture.mask_rcnn_full import MaskRCNN

weights = hf_hub_download("ApyHTML19/MediSeg-Mask-RCNN", "checkpoints/best_model.safetensors")
model = MaskRCNN(num_classes=2, pretrained_backbone=False)
model.load_state_dict(load_file(weights, device="cpu"))
model.eval()

image, scale = letterbox(Image.open("photo.jpg").convert("RGB"))  # 640x640, aspect ratio kept
with torch.no_grad():
    detections = model(TF.to_tensor(image).unsqueeze(0))[0]  # {boxes, labels, scores, masks}

keep = detections["scores"] >= 0.5
boxes = detections["boxes"][keep] / scale  # back to original image coordinates
```

The model code lives in the source repository (`model_rcnn_scratch/`); `model_rcnn_scratch.hub.load_model()` does the download + loading above in one call.

### Command line

```bash
python -m model_rcnn_scratch.predict path/to/photo.jpg        # detection only -> outputs_results/<name>_pred.png
python -m harness.pipeline path/to/photo.jpg                  # detection + masked crop + OCR, no API key needed
python -m operation.run path/to/photo.jpg                     # full pipeline (needs OPENAI_API_KEY + TAVILY_API_KEY)
```

Full pipeline outputs, in `outputs_results/pipeline/`: `<name>_detections.jpg`, `<name>_box<i>.png` (masked crop), `<name>_analysis.json`, `<name>_annotated.jpg`.

## Examples

All images below are real outputs of the model and pipeline, stored in [`assets/examples/`](assets/examples/).

### Several boxes, French labels

<p align="center">
  <img src="assets/examples/01_multi_box_crop0.png" alt="Masked crop of the Doliprane box" width="35%">
</p>

Masked crop (`<name>_box0.png`) sent to OCR. The three boxes are detected with score ≈ 1.00; the pipeline identifies **Doliprane 1000 mg**, **Ibuprofène Mylan 200 mg** and **Efferalgan 500 mg**, each with its active ingredient (paracetamol / ibuprofen).

### Bilingual French / Arabic tube and box

<p align="center">
  <img src="assets/examples/02_bilingual_tube_annotated.jpg" alt="Aureomycine tube and box, French and Arabic labels" width="80%">
</p>

The French tube (#0) and the Arabic box (#1) are both identified as **Auréomycine 1 %** (chlortetracycline, ophthalmic ointment). Box #2 is a duplicate partial detection of the tube's label.

### Arabic-only label

<p align="center">
  <img src="assets/examples/03_arabic_box_annotated.jpg" alt="Vermox box with an Arabic label" width="60%">
</p>

Arabic OCR reads **فيرموكسمجم** ("Vermox mg", words merged), Latin OCR reads **500** and **janssen**, and the LLM resolves it to **Vermox 500 mg** (mebendazole).

### Failure cases

<p align="center">
  <img src="assets/examples/04_failure_false_positives_annotated.jpg" alt="False positives on a glass and a text block" width="60%">
</p>

Advertising image: the Propalgina Plus box (#0) is correct, but the model also fires on a **block of text** (#1, score 0.78) and a **glass** (#2, score 0.50). The LLM then "identifies" the text block as paracetamol 500 mg from the words it contains.

<p align="center">
  <img src="assets/examples/05_failure_llm_misid_detections.jpg" alt="Two overlapping Colocalm boxes, correctly segmented" width="49%">
  <img src="assets/examples/05_failure_llm_misid_annotated.jpg" alt="LLM misidentification of one box as Lokelma" width="49%">
</p>

Segmentation of the two overlapping boxes is correct and the OCR reads the Arabic brand name correctly (كولوكالم, Colocalm — mebeverine), but the LLM maps it to **Lokelma**, a different medicine. Box #1 gets the right active ingredient without a name.

## Limitations

- Single class (`medicine_box`); not a general-purpose detector. Any rectangular, label-like region (text panels, product packaging, glasses) can trigger a detection — see the failure cases above. Raise the score threshold (default 0.5) to trade recall for precision.
- Small dataset (540 training images), mostly French/Arabic packages photographed on plain backgrounds — expect weaker results on cluttered scenes, blisters, bottles or other scripts.
- Overlapping boxes may produce duplicate or partial detections.
- The identification stage depends on OCR quality and on an external LLM + web search; it can return a wrong medicine with high confidence. Treat it as a hint, not a diagnosis.
- `Backbone_RES` downloads ImageNet ResNet-50 weights when training from scratch (not needed for inference with `pretrained_backbone=False`).

## Citation

```
@misc{mediseg2026,
  title  = {MediSeG: Mask R-CNN for Medicine Box Instance Segmentation},
  year   = {2026},
  url    = {https://huggingface.co/ApyHTML19/MediSeg-Mask-RCNN}
}
```