models/teeth_segmentation.py

"""
models/teeth_segmentation_torch.py
====================================
Torchvision Mask R-CNN for dental panoramic X-ray tooth segmentation.

COCO pretrained weights loaded automatically.

Dataset:  AKUDENTAL — 333 panoramic X-rays, COCO JSON format
Strategy: Binary first (1 class), then 35 FDI classes

Usage:
    python models/teeth_segmentation_torch.py train
    python models/teeth_segmentation_torch.py evaluate
    python models/teeth_segmentation_torch.py predict --image data/images/1.jpg
"""

import os
import sys
import json
import argparse
import numpy as np
import cv2
import torch
import skimage.io
import skimage.draw
from pathlib import Path
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision.models.detection import maskrcnn_resnet50_fpn, MaskRCNN_ResNet50_FPN_Weights
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
import torchvision.transforms.functional as F

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.append(str(PROJECT_ROOT))

from utils.preprocessing import enhance_contrast, load_image
from configs.model_config import (
    BINARY, NUM_CLASSES, IMAGE_MIN_SIZE, IMAGE_MAX_SIZE,
    EPOCHS, BATCH_SIZE, NUM_WORKERS, LR, MOMENTUM, WEIGHT_DECAY,
    LR_STEP_SIZE, LR_GAMMA, CONF_THRESHOLD, NMS_THRESHOLD, MAX_DETECTIONS,
    ANCHOR_SIZES, ANCHOR_RATIOS, BINARY_CLASSES, FDI_CLASSES,
)

DATA_DIR    = PROJECT_ROOT / 'data'
IMG_DIR     = DATA_DIR / 'processed'
ANN_DIR     = DATA_DIR / 'annotations'
RESULTS_DIR = PROJECT_ROOT / 'outputs' / 'results' / 'maskrcnn_torch'
VIZ_DIR     = PROJECT_ROOT / 'outputs' / 'visualizations'
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
VIZ_DIR.mkdir(parents=True, exist_ok=True)

DEVICE      = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

QUADRANT_COLORS = {
    'UR': (74,  144, 217),
    'UL': (232, 112, 64),
    'LL': (46,  204, 113),
    'LR': (155, 89,  182),
    'XX': (170, 170, 170),
}
FDI_TO_QUADRANT = {
    **{str(fdi): 'UR' for fdi in range(11, 19)},
    **{str(fdi): 'UL' for fdi in range(21, 29)},
    **{str(fdi): 'LL' for fdi in range(31, 39)},
    **{str(fdi): 'LR' for fdi in range(41, 49)},
}


def get_color(cls_name):
    if cls_name in ('tooth', '__background__'):
        return QUADRANT_COLORS['UR']
    return QUADRANT_COLORS.get(FDI_TO_QUADRANT.get(cls_name, 'XX'), (170, 170, 170))


# ── Dataset ───────────────────────────────────────────────────────────────────

class TeethDataset(Dataset):

    def __init__(self, annotation_file, img_dir, binary=True, transforms=None):
        with open(annotation_file) as f:
            self.coco = json.load(f)

        self.img_dir    = Path(img_dir)
        self.binary     = binary
        self.transforms = transforms

        # Build category map
        if binary:
            self.cat_map    = {cat['id']: 1 for cat in self.coco['categories']}
            self.class_names = BINARY_CLASSES
        else:
            sorted_cats     = sorted(self.coco['categories'], key=lambda x: x['id'])
            self.cat_map    = {cat['id']: i+1 for i, cat in enumerate(sorted_cats)}
            self.class_names = FDI_CLASSES

        # Index annotations by image_id
        self.anns_by_image = {}
        for ann in self.coco['annotations']:
            self.anns_by_image.setdefault(ann['image_id'], []).append(ann)

        # Filter to images with annotations that exist on disk
        self.images = []
        for img in self.coco['images']:
            path = self.img_dir / img['file_name']
            if path.exists() and self.anns_by_image.get(img['id']):
                self.images.append(img)

        print(f'  Loaded {len(self.images)} images')

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_info = self.images[idx]
        img_path = self.img_dir / img_info['file_name']

        # Load image
        image = load_image(str(img_path))

        H, W = image.shape[:2]
        anns = self.anns_by_image.get(img_info['id'], [])

        masks    = []
        boxes    = []
        labels   = []

        for ann in anns:
            seg = ann.get('segmentation', [])
            if not seg or not isinstance(seg[0], list):
                continue
            flat = seg[0]
            if len(flat) < 6:
                continue

            xs = np.array(flat[0::2])
            ys = np.array(flat[1::2])

            # Draw polygon mask
            mask = np.zeros((H, W), dtype=np.uint8)
            rr, cc = skimage.draw.polygon(ys, xs)
            rr = np.clip(rr, 0, H-1)
            cc = np.clip(cc, 0, W-1)
            mask[rr, cc] = 1

            if mask.sum() == 0:
                continue

            # Bounding box from mask
            rows = np.any(mask, axis=1)
            cols = np.any(mask, axis=0)
            y1, y2 = np.where(rows)[0][[0, -1]]
            x1, x2 = np.where(cols)[0][[0, -1]]

            if x2 <= x1 or y2 <= y1:
                continue

            masks.append(mask)
            boxes.append([x1, y1, x2, y2])
            labels.append(self.cat_map.get(ann['category_id'], 1))

        if not masks:
            # Return empty target if no valid annotations
            target = {
                'boxes':   torch.zeros((0, 4), dtype=torch.float32),
                'labels':  torch.zeros(0, dtype=torch.int64),
                'masks':   torch.zeros((0, H, W), dtype=torch.uint8),
                'image_id': torch.tensor([img_info['id']]),
                'area':    torch.zeros(0, dtype=torch.float32),
                'iscrowd': torch.zeros(0, dtype=torch.int64),
            }
        else:
            boxes_t  = torch.as_tensor(boxes,  dtype=torch.float32)
            labels_t = torch.as_tensor(labels, dtype=torch.int64)
            masks_t  = torch.as_tensor(np.stack(masks), dtype=torch.uint8)
            area     = (boxes_t[:, 3] - boxes_t[:, 1]) * (boxes_t[:, 2] - boxes_t[:, 0])

            target = {
                'boxes':    boxes_t,
                'labels':   labels_t,
                'masks':    masks_t,
                'image_id': torch.tensor([img_info['id']]),
                'area':     area,
                'iscrowd':  torch.zeros(len(masks), dtype=torch.int64),
            }

        # Convert image to tensor
        image_t = F.to_tensor(image)

        if self.transforms:
            image_t, target = self.transforms(image_t, target)

        return image_t, target


def collate_fn(batch):
    return tuple(zip(*batch))


# ── Model ─────────────────────────────────────────────────────────────────────

def build_model(num_classes):
    """
    Build Mask R-CNN with ResNet50+FPN backbone.
    Loads COCO pretrained weights and replaces heads for num_classes.
    """
    # Load pretrained model
    model = maskrcnn_resnet50_fpn(
        weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT,
        box_nms_thresh = NMS_THRESHOLD,
        box_score_thresh = CONF_THRESHOLD,
        box_detections_per_img = MAX_DETECTIONS,
    )

    # Replace box predictor head
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # Replace mask predictor head
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer     = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask, hidden_layer, num_classes
    )

    return model


# ── Training ──────────────────────────────────────────────────────────────────

def train():
    print(f'Device: {DEVICE}')
    print(f'Binary: {BINARY}')

    num_classes = 2 if BINARY else 36  # bg + classes

    # Datasets
    dataset_train = TeethDataset(ANN_DIR / 'train.json', IMG_DIR, binary=BINARY)
    dataset_val   = TeethDataset(ANN_DIR / 'val.json',   IMG_DIR, binary=BINARY)

    loader_train = DataLoader(
        dataset_train, batch_size=BATCH_SIZE,
        shuffle=True, num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
    )
    loader_val = DataLoader(
        dataset_val, batch_size=1,
        shuffle=False, num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
    )

    # Model
    model = build_model(num_classes)
    model.to(DEVICE)

    # Optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(
        params, lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY
    )
    lr_scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=LR_STEP_SIZE, gamma=LR_GAMMA # step_size reduces LR every LR_STEP_SIZE
    )

    # CSV logger
    import csv
    csv_path = RESULTS_DIR / 'training_history.csv'
    csv_file = open(csv_path, 'w', newline='')
    csv_writer = csv.writer(csv_file)
    csv_writer.writerow(['epoch', 'train_loss', 'val_loss', 'loss_classifier',
                         'loss_box_reg', 'loss_mask', 'loss_objectness', 'loss_rpn_box_reg'])

    best_loss = float('inf')

    print(f'\nTraining for {EPOCHS} epochs...')
    print(f'Train: {len(dataset_train)} images')
    print(f'Val:   {len(dataset_val)} images\n')

    for epoch in range(1, EPOCHS + 1):
        model.train()
        epoch_losses = []

        for i, (images, targets) in enumerate(loader_train):
            images  = [img.to(DEVICE) for img in images]
            targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]

            loss_dict = model(images, targets)
            losses    = sum(loss for loss in loss_dict.values())

            optimizer.zero_grad()
            losses.backward()
            optimizer.step()

            epoch_losses.append(losses.item())

            if (i + 1) % 20 == 0:
                print(f'  Epoch {epoch}/{EPOCHS} step {i+1}/{len(loader_train)} '
                      f'loss={losses.item():.4f}')
                
        #Validationn phase
        model.train()
        val_losses = []

        with torch.no_grad():
            for images, targets in loader_val:
                images = [img.to(DEVICE) for img in images]
                targets = [{k: v.to(DEVICE) for k,v in t.items()} for t in targets]
                loss_dict_val = model(images,targets)
                val_loss = sum(loss for loss in loss_dict_val.values())    
                val_losses.append(val_loss.item())

        lr_scheduler.step()

        avg_train = np.mean(epoch_losses)
        avg_val = np.mean(val_losses)
        ld       = {k: v.item() for k, v in loss_dict.items()}

        print(f'Epoch {epoch}/{EPOCHS}  '
              f'train={avg_train:.4f}  '
              f'val={avg_val:.4f}  '
              f'lr={optimizer.param_groups[0]["lr"]:.6f}')

        csv_writer.writerow([
            epoch, avg_train, avg_val,
            ld.get('loss_classifier', 0),
            ld.get('loss_box_reg', 0),
            ld.get('loss_mask', 0),
            ld.get('loss_objectness', 0),
            ld.get('loss_rpn_box_reg', 0),
        ])
        csv_file.flush()

        # Save best model
        if avg_val < best_loss:
            best_loss = avg_val
            torch.save(model.state_dict(), RESULTS_DIR / 'best.pth')
            print(f'  → Saved best model (val_loss={best_loss:.4f})')

        # Save latest
        torch.save(model.state_dict(), RESULTS_DIR / 'last.pth')

    csv_file.close()
    print(f'\nTraining complete!')
    print(f'Best weights: {RESULTS_DIR}/best.pth')


# ── Evaluation ────────────────────────────────────────────────────────────────

def evaluate():
    """Run COCO-style evaluation using pycocotools"""
    #pycocotools reads coco objects and calculates IoU,precision,recall,mAP
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval
    from pycocotools import mask as maskUtils
    import copy, tempfile, os

    num_classes = 2 if BINARY else 36
    weights_path = RESULTS_DIR / 'best.pth'
    assert weights_path.exists(), f'Weights not found: {weights_path}'

    model = build_model(num_classes)
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE, weights_only=True))
    model.to(DEVICE)
    model.eval()

    class_names = BINARY_CLASSES if BINARY else FDI_CLASSES

    dataset_val = TeethDataset(ANN_DIR / 'val.json', IMG_DIR, binary=BINARY)
    loader_val  = DataLoader(dataset_val, batch_size=1,
                             shuffle=False, collate_fn=collate_fn)

    coco_gt      = COCO(str(ANN_DIR / 'val.json'))
    coco_results = []

    print(f'Evaluating {len(dataset_val)} val images...')
    with torch.no_grad():
        for images, targets in loader_val:
            images  = [img.to(DEVICE) for img in images]
            outputs = model(images)

            for target, output in zip(targets, outputs):
                image_id = target['image_id'].item()
                boxes    = output['boxes'].cpu().numpy()
                scores   = output['scores'].cpu().numpy()
                labels   = output['labels'].cpu().numpy()
                masks    = output['masks'].cpu().numpy()  # (N, 1, H, W)

                for i in range(len(boxes)):
                    if scores[i] < CONF_THRESHOLD:
                        continue
                    #i-th tooth, 0-drops channel dimension
                    #as model outputs probabilites, >0.5 converts them into yes/no decisions
                    # turn into np.uint8 as pycocotools expects ints not bools
                    mask = (masks[i, 0] > 0.5).astype(np.uint8)
                    rle  = maskUtils.encode(np.asfortranarray(mask))
                    rle['counts'] = rle['counts'].decode('utf-8')

                    x1, y1, x2, y2 = boxes[i]
                    coco_results.append({
                        'image_id':    image_id,
                        'category_id': int(labels[i]),
                        'segmentation': rle,
                        'bbox':        [float(x1), float(y1),
                                        float(x2-x1), float(y2-y1)],
                        'score':       float(scores[i]),
                    })

    if not coco_results:
        print('No detections above threshold.')
        return

 
    # Binary mode: remap all GT categories to 1
    if BINARY:
        val_gt = copy.deepcopy(coco_gt.dataset)
        for ann in val_gt['annotations']:
            ann['category_id'] = 1
        val_gt['categories'] = [{'id': 1, 'name': 'tooth'}]
        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
            json.dump(val_gt, f)
            tmp_path = f.name
        coco_gt = COCO(tmp_path)
        os.unlink(tmp_path)


    # converrts lists of prediction dicts into COCO obejct that
    # pycocotools can compare against ground truth
    coco_dt  = coco_gt.loadRes(coco_results) 
    # creates evaluator: coco_gt - gt annotations, coco_dt -model prediction
    # 'segm' evaluate mask IoU
    coco_eval = COCOeval(coco_gt, coco_dt, 'segm')
    #for every detection computes of matches a gt annotation
    coco_eval.evaluate()
    #aggregates per image results into PR curves across all images and IoU thresholds
    coco_eval.accumulate()
    coco_eval.summarize()

    mAP50 = coco_eval.stats[1]
    mAP   = coco_eval.stats[0]

    print('RESULTS SUMMARY')
    print('='*55)
    print(f'Model:     Mask R-CNN ResNet50+FPN (torchvision)')
    print(f'Dataset:   AKUDENTAL (333 panoramic X-rays)')
    print(f'Classes:   {"1 (binary tooth)" if BINARY else "35 FDI"}')
    print(f'mAP@50:    {mAP50*100:.1f}%')
    print(f'mAP@50-95: {mAP*100:.1f}%')


# ── Prediction ────────────────────────────────────────────────────────────────

def predict(image_path):
    num_classes  = 2 if BINARY else 36
    weights_path = RESULTS_DIR / 'best.pth'
    assert weights_path.exists(), f'Weights not found: {weights_path}'

    class_names = BINARY_CLASSES if BINARY else FDI_CLASSES

    model = build_model(num_classes)
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE, weights_only=True))
    model.to(DEVICE)
    model.eval()

    image    = load_image(str(image_path))
    enhanced = enhance_contrast(image, method='clahe')
    image_t  = F.to_tensor(enhanced).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        outputs = model(image_t)

    output = outputs[0]
    boxes  = output['boxes'].cpu().numpy()
    scores = output['scores'].cpu().numpy()
    labels = output['labels'].cpu().numpy()
    masks  = output['masks'].cpu().numpy()  # (N, 1, H, W)

    keep = scores >= CONF_THRESHOLD
    boxes, scores, labels, masks = boxes[keep], scores[keep], labels[keep], masks[keep]
    n = len(boxes)
    print(f'Detected: {n} instances')
    if n == 0:
        return enhanced

    output_img = enhanced.copy()
    for i in range(n):
        cls_name  = class_names[labels[i]] if labels[i] < len(class_names) else 'tooth'
        color     = get_color(cls_name)
        mask      = (masks[i, 0] > 0.5)

        overlay   = output_img.copy()
        overlay[mask] = color
        output_img = cv2.addWeighted(output_img, 0.55, overlay, 0.45, 0)

        x1, y1, x2, y2 = map(int, boxes[i])
        cv2.rectangle(output_img, (x1, y1), (x2, y2), color, 2)

        # FDI label at centroid
        mask_u8 = mask.astype(np.uint8)
        M = cv2.moments(mask_u8)
        if M['m00'] > 0:
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])
        else:
            cx, cy = (x1+x2)//2, (y1+y2)//2
        cv2.putText(output_img, cls_name, (cx-10, cy+5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)

    out_path = VIZ_DIR / f'torch_pred_{Path(image_path).name}'
    cv2.imwrite(str(out_path), cv2.cvtColor(output_img, cv2.COLOR_RGB2BGR))
    print(f'Saved: {out_path}')
    return output_img

# used in Gradio, loads once at startup then reuse for every request
def load_inference_model():
    """Load model for use in Gradio/FastAPI."""
    weights_path = RESULTS_DIR / 'best.pth'

    # Downloads model weights stored on huggingface
    if not weights_path.exists():
        from huggingface_hub import hf_hub_download
        print("Downloading weights from Hugging Face Hub ...")
        hf_hub_download(
            repo_id = "chocodo/dental-segmentation-maskrcnn-torch",
            filename = "best.pth",
            local_dir = str(RESULTS_DIR),
        )
        print("Weights downloaded.")

    num_classes  = 2 if BINARY else 36
    model = build_model(num_classes)
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE, weights_only=True))
    model.to(DEVICE)
    model.eval()
    return model


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Mask R-CNN dental tooth segmentation')
    parser.add_argument('command', choices=['train', 'evaluate', 'predict'])
    parser.add_argument('--image', help='Image path for predict command')
    args = parser.parse_args()

    if args.command == 'train':
        train()
    elif args.command == 'evaluate':
        evaluate()
    elif args.command == 'predict':
        assert args.image, '--image required'
        predict(args.image)