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