我已經訓練了一個掩碼RCNN網絡，例如蘋果的分割。我能夠加載權重并為我的測試圖像生成預測。正在生成的掩碼似乎位于正確的位置，但掩模本身沒有真正的形式。它看起來就像一堆像素訓練是根據本文中的數據集完成的，以下是用于訓練和生成權重的代碼的github鏈接預測代碼如下。（我省略了我創建路徑變量并分配路徑的部分）import osimport globimport numpy as npimport pandas as pdimport cv2 as cvimport fileinputimport torchimport torch.utils.dataimport torchvisionfrom data.apple_dataset import AppleDatasetfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictorfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictorimport utility.utils as utilsimport utility.transforms as Tfrom PIL import Imagefrom matplotlib import pyplot as plt%matplotlib inlinedef get_transform(train): transforms = [] transforms.append(T.ToTensor()) if train: transforms.append(T.RandomHorizontalFlip(0.5)) return T.Compose(transforms)def get_maskrcnn_model_instance(num_classes): # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False) # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # and replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes) return modelnum_classes = 2device = torch.device('cpu')model = get_maskrcnn_model_instance(num_classes)checkpoint = torch.load('model_49.pth', map_location=device)model.load_state_dict(checkpoint['model'], strict=False)dataset_test = AppleDataset(test_image_files_path, get_transform(train=False))img, _ = dataset_test[1]model.eval()with torch.no_grad(): prediction = model([img.to(device)])
1 回答
溫溫醬
TA貢獻1752條經驗 獲得超4個贊
來自掩碼 R-CNN 的預測具有以下結構：
在推理過程中，模型只需要輸入張量，并將后處理的預測作為 ，每個輸入圖像返回一個。的字段如下：List[Dict[Tensor]]Dict
boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values between 0 and H and 0 and W
labels (Int64Tensor[N]): the predicted labels for each image
scores (Tensor[N]): the scores or each prediction
masks (UInt8Tensor[N, 1, H, W]): the predicted masks for each instance, in 0-1 range.
您可以使用 OpenCV 和函數來繪制蒙版，如下所示：findContoursdrawContours
img_cv = cv2.imread('input.jpg', cv2.COLOR_BGR2RGB)
for i in range(len(prediction[0]['masks'])):
# iterate over masks
mask = prediction[0]['masks'][i, 0]
mask = mask.mul(255).byte().cpu().numpy()
contours, _ = cv2.findContours(
mask.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
cv2.drawContours(img_cv, contours, -1, (255, 0, 0), 2, cv2.LINE_AA)
cv2.imshow('img output', img_cv)
添加回答
舉報
0/150
提交
取消