--- license: mit library_name: pytorch tags: - instance-segmentation - mask-rcnn - object-detection - resnet50 - fpn - medical - ocr - pytorch - safetensors pipeline_tag: image-segmentation ---
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/