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

MediSeG banner

# 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**.

Mask R-CNN detections on three medicine boxes Full pipeline output: each box identified with name, dosage, form and active ingredient

Left: Mask R-CNN output (boxes + masks + scores). Right: full pipeline output (OCR + LLM + web search).

## 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/_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/`: `_detections.jpg`, `_box.png` (masked crop), `_analysis.json`, `_annotated.jpg`. ## Examples All images below are real outputs of the model and pipeline, stored in [`assets/examples/`](assets/examples/). ### Several boxes, French labels

Masked crop of the Doliprane box

Masked crop (`_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

Aureomycine tube and box, French and Arabic labels

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

Vermox box with an Arabic label

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

False positives on a glass and a text block

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.

Two overlapping Colocalm boxes, correctly segmented LLM misidentification of one box as Lokelma

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} } ```