def __init__(
self,
cfg,
confidence_thresholds_for_classes,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_thresholds_for_classes = torch.tensor(confidence_thresholds_for_classes)
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def prepare_for_coco_segmentation(predictions, dataset):
import pycocotools.mask as mask_util
import numpy as np
masker = Masker(threshold=0.5, padding=1)
# assert isinstance(dataset, COCODataset)
coco_results = []
for image_id, prediction in tqdm(enumerate(predictions)):
original_id = dataset.id_to_img_map[image_id]
if len(prediction) == 0:
continue
img_info = dataset.get_img_info(image_id)
image_width = img_info["width"]
image_height = img_info["height"]
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
# t = time.time()
# Masker is necessary only if masks haven't been already resized.
if list(masks.shape[-2:]) != [image_height, image_width]:
masks = masker(masks.expand(1, -1, -1, -1, -1), prediction)
masks = masks[0]
# logger.info('Time mask: {}'.format(time.time() - t))
# prediction = prediction.convert('xywh')
#import pdb
#pdb.set_trace()
# boxes = prediction.bbox.tolist()
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
# rles = prediction.get_field('mask')
rles = [
mask_util.encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
for mask in masks
]
for rle in rles:
rle["counts"] = rle["counts"].decode("utf-8")
mapped_labels = [
dataset.contiguous_category_id_to_json_id[i] for i in labels
]
coco_results.extend([{
"image_id": original_id,
"category_id": mapped_labels[k],
"segmentation": rle,
"score": scores[k],
} for k, rle in enumerate(rles)])
return coco_results
def __init__(
self,
cfg,
confidence_thresholds_for_classes,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
# self.CATEGORIES = ['__background', 'Person', 'Car', 'Train', 'Rider', 'Truck', 'Motorcycle', 'Bicycle', 'Bus']
self.CATEGORIES = ['__background', 'Car']
self.confidence_thresholds_for_classes = torch.tensor(
confidence_thresholds_for_classes)
self.cfg = cfg.clone()
self.min_image_size = min_image_size
self.model = {}
self.device = torch.device(cfg.MODEL.DEVICE)
backbone = build_backbone(cfg).to(self.device).eval()
fcos = build_rpn(cfg, backbone.out_channels).to(self.device).eval()
self.model["backbone"] = backbone
self.model["fcos"] = fcos
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.cpu_device = torch.device("cpu")
self.confidence_thresholds_for_classes = torch.tensor(
confidence_thresholds_for_classes)
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def prepare_for_coco_segmentation(predictions, dataset):
import pycocotools.mask as mask_util
import numpy as np
masker = Masker(threshold=0.5, padding=1)
# assert isinstance(dataset, COCODataset)
coco_results = []
for image_id, prediction in tqdm(enumerate(predictions)):
original_id = dataset.id_to_img_map[image_id]
if len(prediction) == 0:
continue
img_info = dataset.get_img_info(image_id)
image_width = img_info["width"]
image_height = img_info["height"]
input_w, input_h = prediction.size
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
if masks.dim() == 3:
# resize masks
stride_mask = prediction.get_field('stride')
h = (masks.shape[1] * stride_mask.float() * image_height / input_h).ceil().long()
w = (masks.shape[2] * stride_mask.float() * image_width / input_w).ceil().long()
mask_th = prediction.get_field('mask_th').cuda()
masks = masks.cuda()
masks = torch.nn.functional.interpolate(input=masks.unsqueeze(1).float(), size=(h, w), mode="bilinear", align_corners=False).gt(mask_th)
masks = masks[:, :, :image_height, :image_width]
masks = masks.cpu()
else:
# t = time.time()
# Masker is necessary only if masks haven't been already resized.
if list(masks.shape[-2:]) != [image_height, image_width]:
masks = masker(masks.expand(1, -1, -1, -1, -1), prediction)
masks = masks[0]
# logger.info('Time mask: {}'.format(time.time() - t))
# prediction = prediction.convert('xywh')
# boxes = prediction.bbox.tolist()
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
# rles = prediction.get_field('mask')
rles = [
mask_util.encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
for mask in masks
]
for rle in rles:
rle["counts"] = rle["counts"].decode("utf-8")
mapped_labels = [dataset.contiguous_category_id_to_json_id[i] for i in labels]
coco_results.extend(
[
{
"image_id": original_id,
"category_id": mapped_labels[k],
"segmentation": rle,
"score": scores[k],
}
for k, rle in enumerate(rles)
]
)
return coco_results