ApyHTML19/MediSeg-Mask-RCNN

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.

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); 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):

python -m model_rcnn_scratch.evaluate

Metric	Value
Dice	TBD
IoU	TBD

How to use

Weights only

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

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/.

Several boxes, French labels

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

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

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

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.

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