import torch
import torchvision
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
import torchvision.transforms as T

def get_model(num_classes=3, weights_path=r"C:\Ishan_works\Pycharm_project\skin_detectron\skin_pallor_segment\Saved_model\mask_rcnn_conjunctiva.pth", device="cpu"):
    device = torch.device(device)

    model = maskrcnn_resnet50_fpn(pretrained=False)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, 256, num_classes)

    model.load_state_dict(torch.load(weights_path, map_location=device))
    model.to(device)
    model.eval()
    return model

def predict(model, image, device="cpu", class_names=None, threshold=0.76):
    device = torch.device(device)
    transform = T.ToTensor()
    image_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(image_tensor)[0]

    results = []
    for i in range(len(outputs["scores"])):
        if outputs["scores"][i] >= threshold:
            label = outputs["labels"][i].item()
            mask = outputs["masks"][i, 0].cpu().numpy()
            results.append({
                "label": class_names[label] if class_names else str(label),
                "mask": mask
            })

    return results