Update README.md

README.md

@@ -8,7 +8,9 @@ tags:
   - resnet50
   - fpn
   - medical
+  - ocr
   - pytorch
+  - safetensors
 pipeline_tag: image-segmentation
 ---
 
@@ -18,41 +20,62 @@ pipeline_tag: image-segmentation
 
 # MediSeG — Mask R-CNN for Medicine Box Instance Segmentation
 
-Custom **Mask R-CNN** (ResNet-50 + FPN + RPN + ROI Align, implemented from scratch in PyTorch) that detects and segments medicine boxes (`Med-Box`) in images.
+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.
 
-## Model description
+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>
 
-- **Backbone**: ResNet-50 (ImageNet-pretrained), feature maps C2–C5.
-- **Neck**: Feature Pyramid Network (FPN), levels P2–P5, 256 channels.
-- **Region proposal**: RPN with 3 anchor ratios per level (0.5, 1.0, 2.0), one fixed anchor scale per level (32/64/128/256).
-- **ROI heads**: shared two-layer MLP trunk → classification head + box regression head (7×7 ROI Align) and a 4-conv + deconv mask head (14×14 → 28×28 ROI Align).
-- **Classes**: 2 (`background`, `Med-Box`).
-- **Input size**: 640×640 RGB.
+## Model description
 
-Full architecture diagram: see `ARCHITECTURE.md` in the source repository.
+- **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
 
-Detect and segment medicine boxes in photos (e.g. pharmacy / assistive applications). Not validated for clinical or safety-critical 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
 
-- 389 training images, 84 validation images, 84 test images.
-- Annotations: YOLO-seg polygons (4-point quadrilaterals per instance), converted to bounding boxes + rasterized binary masks.
-- Single object class: `Med-Box`.
+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
 
-- Optimizer: SGD, `lr=0.005`, `momentum=0.9`, `weight_decay=0.0005`, `StepLR` (step=8, gamma=0.1).
-- Losses: RPN objectness (BCE) + RPN box regression (smooth-L1), ROI classification (cross-entropy) + ROI box regression (smooth-L1, class-specific), mask (BCE on 28×28 predictions cropped/resized from ground-truth masks).
-- Positive/negative sampling: IoU ≥ 0.7 / < 0.3 for RPN anchors, IoU ≥ 0.5 for ROI proposals (see `app/config.py` for exact sampling ratios and batch sizes).
+- 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
 
-Evaluated with **Dice** and **IoU** on the held-out test split (combined instance masks vs. combined ground-truth masks per image). Fill in scores after running:
+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 app.evaluate
+python -m model_rcnn_scratch.evaluate
 ```
 
 | Metric | Value |
@@ -62,34 +85,100 @@ python -m app.evaluate
 
 ## How to use
 
+### Weights only
+
 ```python
 import torch
-from app.model_arch.mask_rcnn_full import MaskRCNN
+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)
-checkpoint = torch.load("best_model.pth", map_location="cpu")
-model.load_state_dict(checkpoint["model_state_dict"])
+model.load_state_dict(load_file(weights, device="cpu"))
 model.eval()
 
-image = torch.rand(1, 3, 640, 640)  # replace with a real preprocessed image
+image, scale = letterbox(Image.open("photo.jpg").convert("RGB"))  # 640x640, aspect ratio kept
 with torch.no_grad():
-    detections = model(image)  # list of {boxes, labels, scores, masks} per image
+    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 (`Med-Box`); not a general-purpose detector.
-- Fixed 640×640 input resize (no letterboxing) — non-square source images are stretched.
-- Small dataset (389 training images) — expect limited generalization outside the training distribution.
+- 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
 
-If you use this model, please cite the repository:
-
 ```
 @misc{mediseg2026,
   title  = {MediSeG: Mask R-CNN for Medicine Box Instance Segmentation},
   year   = {2026},
-  url    = {<repo-url>}
+  url    = {https://huggingface.co/ApyHTML19/MediSeg-Mask-RCNN}
 }
 ```