Upload model_utils.py

model_utils.py

@@ -0,0 +1,42 @@
+
+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