def generate_davis_final_annotation(predictions, dataset, output_file, annotation_type):
masker = Masker(padding=1, keep_score=True)
# assert isinstance(dataset, COCODataset)
davis_results = []
for id, prediction in tqdm(enumerate(predictions)):
if len(prediction) == 0:
continue
video_id = dataset.get_annotation_video_id(id)
img_id = dataset.get_annotation_img_id(id)
if video_id == "scooter-black" and img_id == "00012":
a = 1
palette = dataset.get_annotation_palette(id)
image_width = dataset.get_img_width(id)
image_height = dataset.get_img_height(id)
prediction = prediction.resize((image_width, image_height))
prediction = prediction.convert('xyxy')
masks = prediction.get_field("mask")
# visualize_batch_mask_for_debug(masks, output_file)
masks = masker(masks, prediction)
# save_final_annotation(masks, video_id, img_id, palette, output_file)
final_predict_annotation = vote_pixel_of_mask_for_annotation(masks, prediction, threshold=0.5)
# Vote
# annotation = ((annotation > 0.5) + 0).astype(int)
save_final_annotation(final_predict_annotation, video_id, img_id, palette, output_file, annotation_type)
def compute_prediction(output):
if output.has_field("mask"):
masks = output.get_field("mask")
masker = Masker(threshold=0.5, padding=1)
masks = masker([masks], [output])[0]
output.add_field("mask", masks)
return output
def extract_information_one(detections, image_shape, original_filenames):
masker = Masker(threshold=0.5, padding=1)
# The resize does not handle crop: the padding area was still there so the bounding box shrink to the left!?
# It seems like the original implementation stored un-padded size!
# Problem solved. See the wiki https://github.com/nncrystals/maskrcnn-benchmark/wiki/Inference-procedures
detections = detections.resize(image_shape)
masks = detections.get_field("mask")
# Caution: if the mask offsets, it is because the bounding box!
masks = masker(masks.expand(1, -1, -1, -1, -1), detections)
masks = masks[0]
num_detections = len(detections)
bbox = detections.bbox.tolist()
scores = detections.extra_fields["scores"].tolist()
labels = detections.extra_fields["labels"].tolist()
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")
results = []
for i in range(num_detections):
results.append({
"name": original_filenames,
"bbox": bbox[i],
"score": scores[i],
"label": labels[i],
"rle": rles[i],
"is_cropped": None
})
return results
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=800,
):
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
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = 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.color_table = np.array([[162, 109, 35], [69, 94, 183], [72, 161, 198],
[82, 158, 127], [120, 72, 122], [105, 124, 135]]*10)
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def predict(
self,
image,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2
):
"""
Arguments:
image (np.ndarray): an image as returned by OpenCV
Returns:
prediction (BoxList): the detected objects. Additional information
of the detection properties can be found in the fields of
the BoxList via `prediction.fields()`
"""
self.model.eval()
self.min_image_size = self.cfg.INPUT.MIN_SIZE_TEST
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
predictions = self.compute_prediction(image)
top_predictions = self.select_top_predictions(predictions, confidence_threshold)
return top_predictions
def __init__(self,
model_path,
config_file,
categories,
show_mask_heatmaps=False,
iou_thr=0.5,
score_thr=0,
device=None):
super(MbDetector, self).__init__(categories=categories,
iou_thr=iou_thr,
score_thr=score_thr,
device=device)
assert is_file(model_path), "model path does not exist!"
assert is_file(config_file), "config path does not exist!"
self.model_path = model_path
self.config_file = config_file
cfg.merge_from_file(config_file)
self.cfg = cfg.clone()
self.cpu_device = torch.device("cpu")
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.show_mask_heatmaps = show_mask_heatmaps
self._build_detector()
def __init__(self,
cfg,
confidence_threshold=-1,
show_mask_heatmaps=False,
show_mask_montage=True,
masks_per_dim=2,
categories=None):
self.categories = categories
self.cfg = cfg.clone()
mask_threshold = -1 if show_mask_montage else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
if confidence_threshold >= 0:
self.confidence_threshold = confidence_threshold
else:
self.confidence_threshold = cfg.MODEL.ROI_HEADS.SCORE_THRESH_VISUALIZATION
self.show_heatmap = show_mask_heatmaps
self.show_mask_montage = show_mask_montage
self.masks_per_dim = masks_per_dim
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.my_palette = []
colors = create_palette()
self.my_palette = []
for color in colors:
self.my_palette.append(color['rgb'])
def prepare_for_pap_segmentation(predictions, dataset):
masker = Masker(threshold=0.5, padding=1)
pap_results = []
pap_gt = []
for idx, (image_id, prediction) in enumerate(predictions.items()):
original_id = dataset.id_to_img_map[image_id]
if len(prediction) == 0:
continue
# TODO replace with get_img_info?
image_width = dataset.maxWS
image_height = dataset.maxWS
target = dataset.get_ground_truth(original_id)
gt_labels = target.get_field("labels").tolist()
gt_masks = target.get_field("masks")
gt_boxes = target.bbox.tolist()
gt_mapped_labels = [
dataset.contiguous_category_id_to_json_id[i] for i in gt_labels
]
pap_gt.extend([{
"image_id": original_id,
"category_id": gt_mapped_labels[k],
"segmentation": rle,
"bbox": gt_boxes[k]
} for k, rle in enumerate(gt_masks)])
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
# 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]
# pdb.set_trace()
# prepare for bbox
# prediction = prediction.convert("xywh")
boxes = prediction.bbox.tolist()
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
# blank_rle = mask_util.encode(np.array(blank_mask[0, :, :], order="F"))
rles = [
maskUtils.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")
# pdb.set_trace()
try:
mapped_labels = [
dataset.contiguous_category_id_to_json_id[i] for i in labels
]
except:
# pdb.set_trace()
print('bug')
pap_results.extend([{
"image_id": original_id,
"category_id": mapped_labels[k],
"segmentation": rle,
"score": scores[k],
"bbox": boxes[k]
} for k, rle in enumerate(rles)])
return pap_gt, pap_results
def __init__(
self,
cfg,
confidence_threshold=0.7,
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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
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)
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
def __init__(self):
self.outputdir = cfg.outputdir + '/maskrcnn'
os.system('mkdir -p ' + self.outputdir)
confidence_threshold = 0.8
show_mask_heatmaps = False
masks_per_dim = 1
min_image_size = 224
mcfg.merge_from_file(cfg.maskrcnnpath)
mcfg.freeze()
self.cfg = mcfg.clone()
self.model = build_detection_model(mcfg)
self.model.eval()
self.device = torch.device(mcfg.MODEL.DEVICE)
self.model.to(self.device)
self.min_image_size = min_image_size
checkpointer = DetectronCheckpointer(mcfg, self.model)
_ = checkpointer.load(mcfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.4
self.masker = Masker(threshold=mask_threshold, padding=100)
# 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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def prepare_for_coco_segmentation(predictions, dataset, maskiou_on):
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
# TODO replace with get_img_info?
# image_width = dataset.coco.imgs[original_id]["width"]
# image_height = dataset.coco.imgs[original_id]["height"]
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')
# boxes = prediction.bbox.tolist()
if maskiou_on:
scores = prediction.get_field("mask_scores").tolist()
else:
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_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
weight_loading=None):
self.cfg = cfg.clone()
# dynamically load labels.json in log directory
self.CATEGORIES = ["__background"]
if 'wolf' in self.cfg.DATASETS.TEST[0]:
with open('../log/wolf_labels.json') as f:
labels = json.load(f)
else:
with open('../log/coco_labels.json') as f:
labels = json.load(f)
for id in labels:
self.CATEGORIES.append(labels[id])
print(self.CATEGORIES)
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
print('self.device: {}'.format(self.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)
if weight_loading:
print('Loading weight from {}.'.format(weight_loading))
_ = checkpointer._load_model(torch.load(weight_loading))
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.cpu_device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# self.cpu_device = torch.device("cpu")
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1,
2**21 - 1]).to(self.cpu_device)
self.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self, cfg, confidence_threshold=0.7, show_mask_heatmaps=False, masks_per_dim=2
):
self.cfg = cfg.clone()
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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def prepare_for_sunspot_segmentation(predictions, dataset):
masker = Masker(threshold=0.5, padding=1)
# assert isinstance(dataset, COCODataset)
results = []
for id, prediction in enumerate(predictions):
original_id = dataset.split_index[id]
if len(prediction) == 0:
continue
image_id = int(original_id.split("_")[1])
img_info = dataset.coco.imgs[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')
# boxes = prediction.bbox.tolist()
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
# rles = prediction.get_field('mask')
rles = [
maskUtils.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]
results.extend(
[
{
"ann_id": original_id.split('_', 1)[1],
"category_id": mapped_labels[k],
"segmentation": rle,
"score": scores[k],
}
for k, rle in enumerate(rles)
]
)
return results
def prepare_for_davis_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 enumerate(predictions):
original_id = dataset.id_to_img_map[image_id]
if len(prediction) == 0:
continue
# TODO replace with get_img_info?
image_width = dataset.get_img_width(image_id)
image_height = dataset.get_img_height(image_id)
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
# t = time.time()
masks = masker(masks, prediction)
# 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
def _buildCOCOPredictions(self):
dt_cache_name = 'len_{}.pkl'.format(len(self.dataset))
dt_cache_path = os.path.join(self.cache_path, dt_cache_name)
if os.path.exists(dt_cache_path):
print("Loading COCO DT annos from: ", dt_cache_path)
coco_preds = pkl.load(open(dt_cache_path, 'rb'))
else:
print("Building COCO Predictions", flush=True)
desc = "Parsing images"
masker = Masker(threshold=0.5, padding=1)
coco_preds = []
for image_id, predictions in tqdm(enumerate(self.every_prediction),
desc=desc):
if len(predictions) == 0:
continue
img_info = self.dataset.get_img_info(image_id)
width = img_info["width"]
height = img_info["height"]
if predictions.size[0] != width or predictions.size[
1] != height:
predictions = predictions.resize(size=(width, height))
for inst_idx in range(len(predictions)):
pred = predictions[inst_idx:inst_idx + 1]
mask = pred.get_field('mask')
if list(mask.shape[-2:]) != [height, width]:
mask = masker(mask.expand(1, -1, -1, -1, -1), pred)
mask = mask[0].squeeze()
segm = SegmentationMask([mask], pred.size, mode='mask')
pred = {
"id": 1 + inst_idx,
"image_id": image_id,
"size": pred.size,
"bbox": pred.bbox[0].tolist(),
"area": pred.area().item(),
"segmentation": segm,
"category_id": pred.get_field("labels").item(),
"score":
pred.get_field("scores").item(), # preds differ here
"iscrowd": 0
}
pred["segmentation"] = self.annToRLE(pred)
coco_preds.append(pred)
if not os.path.exists(self.cache_path):
os.makedirs(self.cache_path)
print("Dumping COCO DT annots", flush=True)
pkl.dump(coco_preds, open(dt_cache_path, 'wb'))
return coco_preds
def __init__(
self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=800,
categories=None,
topk_inplane_rotations=9,
topk_viewpoints=9,
# viewpoints_xyz=None,
# inplane_rotations=None,
# fixed_transforms_dict=None,
# camera_intrinsics=None
):
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
self.CATEGORIES += categories
self.topk_inplane_rotations = topk_inplane_rotations
self.topk_viewpoints = topk_viewpoints
# self.viewpoints_xyz = viewpoints_xyz
# self.inplane_rotations = inplane_rotations
# self.fixed_transforms_dict = fixed_transforms_dict
# self.camera_intrinsics = camera_intrinsics
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
print("Using model at : {}".format(cfg.MODEL.WEIGHT))
_ = 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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
cfg,
model_type,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
self.model_type = model_type
if model_type=="charts":
self.CATEGORIES = (
"__background__",
'other_object',
'chart'
)
elif model_type=="tables":
self.CATEGORIES = (
"__background__",
'table'
)
else:
raise Exception("Unsupported model_type. (either 'charts' or 'tables'")
self.cfg = cfg.clone()
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
model,
CATEGORIES,
dataset,
confidence_threshold=0.5,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=224,
):
if model == 'faster':
config_file = "faster-retina/configs/e2e_faster_rcnn_R_50_FPN_1x_{}_test.yaml".format(
dataset)
if model == 'retinanet':
config_file = 'faster-retina/configs/retinanet_R-50-FPN_1x-{}.yaml'.format(
dataset)
if model == 'maskrcnn':
config_file = 'faster-retina/configs/e2e_mask_rcnn_R_50_FPN_1x-{}.yaml'.format(
dataset)
cfg.merge_from_file(config_file)
self.cfg = cfg.clone()
self.CATEGORIES = CATEGORIES
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device('cuda')
self.model.to(self.device)
self.min_image_size = min_image_size
self.feat_extractor = FeatureExtractorFromBoxes(self.model)
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_threshold = confidence_threshold
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
# TODO replace with get_img_info?
image_width = dataset.coco.imgs[original_id]["width"]
image_height = dataset.coco.imgs[original_id]["height"]
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
# 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]
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
rects = [mask_to_roRect(mask, [image_height, image_width]) for mask in masks]
mapped_labels = [dataset.contiguous_category_id_to_json_id[i] for i in labels]
esd = []
for k, rect in enumerate(rects):
if rect.all() == 0:
continue
else:
esd.append(
{
"image_id": original_id,
"category_id": mapped_labels[k],
# "segmentation": rle,
"seg_rorect": rect.tolist(),
"score": scores[k],
}
)
coco_results.extend(esd)
return coco_results
def prepare_for_segmentation(predictions):
import pycocotools.mask as mask_util
import numpy as np
masker = Masker(threshold=0.5, padding=1)
coco_results = []
for image_id, prediction in tqdm(enumerate(predictions)):
original_id = image_id
if len(prediction) == 0:
continue
image_width = 640
image_height = 480
prediction = prediction.resize((image_width, image_height))
masks = prediction.get_field("mask")
# 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]
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
rles = [
mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0]
for mask in masks
]
for rle in rles:
rle["counts"] = rle["counts"].decode("utf-8")
coco_results.extend(
[
{
"image_id": original_id,
"category_id": labels[k],
"segmentation": rle,
"score": scores[k],
}
for k, rle in enumerate(rles)
]
)
return coco_results
def __init__(self, cfg, confidence_threshold, models_dir, dataset,
max_training_images, num_classes_trained_together):
self.cfg = cfg.clone()
self.cfg.MINIBOOTSTRAP.DETECTOR.NUM_CLASSES = num_classes_trained_together
self.cfg.NUM_IMAGES = max_training_images
self.model = build_detection_model(self.cfg)
try:
self.model.rpn.head.classifiers = torch.load(
os.path.join(models_dir, 'classifier_rpn'))
self.model.rpn.head.regressors = torch.load(
os.path.join(models_dir, 'regressor_rpn'))
self.model.rpn.head.stats = torch.load(
os.path.join(models_dir, 'stats_rpn'))
except:
pass
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform(cfg)
self.masker = Masker(threshold=0.5, 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_threshold = confidence_threshold
if dataset == 'iCWT_TT':
self.CATEGORIES = self.CATEGORIES_iCWT_TT
elif dataset == 'iCWT_TT_TABLETOP':
self.CATEGORIES = self.CATEGORIES_iCWT_TT_21
elif dataset == 'ycbv':
self.CATEGORIES = self.CATEGORIES_YCBV
else:
self.CATEGORIES = None
def __init__(self,
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
model_path=None):
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)
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
if model_path:
logging.info('Loading model from model-path: %s', model_path)
load_path = model_path
else:
if checkpointer.has_checkpoint():
load_path = checkpointer.get_checkpoint_file()
logging.info('Loading model from latest checkpoint: %s',
load_path)
else:
load_path = cfg.MODEL.WEIGHT
logging.info('Loading model from cfg.MODEL.WEIGHT: %s',
load_path)
checkpointer.load(load_path, use_latest=False)
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_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def __init__(
self,
device,
confidence_threshold=0.7,
masks_per_dim=2,
min_image_size=224,
):
self.device = device
self.imgCnt = 0
self.min_image_size = min_image_size
mask_threshold = 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_threshold = confidence_threshold
self.masks_per_dim = masks_per_dim
def add_annotations(orig_image, prediction, mask_on=True):
show_mask_heatmaps = False
mask_threshold = -1 if show_mask_heatmaps else 0.5
masker = Masker(threshold=mask_threshold, padding=1)
# reshape prediction (a BoxList) into the original image size
height, width = orig_image.shape[:-1]
prediction = prediction.resize((width, height))
if prediction.has_field("mask"):
# if we have masks, paste the masks in the right position
# in the image, as defined by the bounding boxes
masks = prediction.get_field("mask")
# always single image is passed at a time
masks = masker([masks], [prediction])[0]
prediction.add_field("mask", masks)
top_preds = select_top_predictions(prediction)
result = orig_image.copy()
result = overlay_boxes(result, top_preds)
if mask_on: # cfg.MODEL.MASK_ON:
result = overlay_mask(result, top_preds)
return overlay_class_names(result, top_preds)
def __init__(
self,
cfg,
show_mask_heatmaps=False,
min_image_size=224,
):
self.cfg = cfg.clone()
self.min_image_size = min_image_size
self.transforms = self.build_transform()
mask_threshold = -1 if show_mask_heatmaps else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
self.cpu_device = torch.device("cpu")
self.visualizer = Visualizer(categories=ro_categories.CATEGORIES,
cfg=cfg,
confidence_threshold=0.8,
show_mask_heatmaps=show_mask_heatmaps)
def prepare_for_coco_segmentation(predictions, dataset):
masker = Masker(threshold=0.5, padding=1)
# assert isinstance(dataset, COCODataset)
coco_results = []
for image_id, prediction in tqdm(enumerate(predictions),
total=len(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("masks")
if isinstance(masks, SegmentationMask):
masks = masks.get_mask_tensor(do_squeeze=False)[:, None]
scores = prediction.get_field("scores").tolist()
labels = prediction.get_field("labels").tolist()
# 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]
rles = convert_binary_to_rle(masks.cpu())
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,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=3,
min_image_size=224,
):
self.res_label_mask_scorse = []
self.res_dir = './res_person/'
config_file = "../configs/e2e_fashion_mask_rcnn_R_50_FPN_1x.yaml"
cfg.merge_from_file(config_file) # 设置配置文件
cfg.merge_from_list(["MODEL.MASK_ON", True])
# cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) # 指定为CPU
cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) # 指定为GPU
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.confidence_threshold = confidence_threshold
self.show_mask_heatmaps = show_mask_heatmaps
self.masks_per_dim = masks_per_dim
def prepare_for_coco_segmentation(predictions, dataset):
masker = Masker(threshold=0.5, padding=1)
# assert isinstance(dataset, COCODataset)
result_dict = {}
total_threads = 5
threads = []
indices = list(range(len(dataset)))
for rank in range(total_threads):
thread = Thread(target=_run_coco_segmentation,
args=(predictions, dataset, masker, rank,
indices[rank::total_threads], result_dict))
thread.start()
threads.append(thread)
for i in range(total_threads):
threads[i].join()
coco_results = []
for rank in result_dict:
coco_results.extend(result_dict[rank])
return coco_results
