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()
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)
if weight_loading:
_ = checkpointer._load_model(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)
# 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,
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
