Spaces:
Sleeping
Sleeping
File size: 1,620 Bytes
b01878e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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
|