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 main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
required=True,
metavar="FILE",
help="path to config file",
)
parser.add_argument(
'--model-path',
type=Path,
help=('Path to model pickle file. If not specified, the latest '
'checkpoint, if it exists, or cfg.MODEL.WEIGHT is loaded.'))
parser.add_argument(
'--output-dir',
default='{cfg_OUTPUT_DIR}/inference-{model_stem}',
help=('Output directory. Can use variables {cfg_OUTPUT_DIR}, which is '
'replaced by cfg.OUTPUT_DIR, and {model_stem}, which is '
'replaced by the stem of the file used to load weights.'))
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"--ckpt",
help=
"The path to the checkpoint for test, default is the latest checkpoint.",
default=None,
)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(
os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
assert cfg.OUTPUT_DIR, 'cfg.OUTPUT_DIR must not be empty.'
checkpointer = DetectronCheckpointer(cfg, model, save_dir=cfg.OUTPUT_DIR)
if args.model_path:
load_path = str(args.model_path.resolve())
load_msg = 'Loading model from --model-path: %s' % load_path
else:
if checkpointer.has_checkpoint():
load_path = checkpointer.get_checkpoint_file()
load_msg = 'Loading model from latest checkpoint: %s' % load_path
else:
load_path = cfg.MODEL.WEIGHT
load_msg = 'Loading model from cfg.MODEL.WEIGHT: %s' % load_path
output_dir = Path(
args.output_dir.format(cfg_OUTPUT_DIR=cfg.OUTPUT_DIR,
model_stem=Path(load_path).stem))
output_dir.mkdir(exist_ok=True, parents=True)
file_logger = common_setup(__file__, output_dir, args)
# We can't log the load_msg until we setup the output directory, but we
# can't get the output directory until we figure out which model to load.
# So we save load_msg and log it here.
logging.info(load_msg)
logging.info('Output inference results to: %s' % output_dir)
logger = logging.getLogger("maskrcnn_benchmark")
logger.info("Using {} GPUs".format(num_gpus))
file_logger.info('Config:')
file_logger.info(cfg)
file_logger.info("Collecting env info (might take some time)")
file_logger.info("\n" + collect_env_info())
# Initialize mixed-precision if necessary
use_mixed_precision = cfg.DTYPE == 'float16'
amp_handle = amp.init(enabled=use_mixed_precision, verbose=cfg.AMP_VERBOSE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt
_ = checkpointer.load(ckpt, use_latest=args.ckpt is None)
iou_types = ("bbox", )
if cfg.MODEL.MASK_ON:
iou_types = iou_types + ("segm", )
if cfg.MODEL.KEYPOINT_ON:
iou_types = iou_types + ("keypoints", )
output_folders = [None] * len(cfg.DATASETS.TEST)
dataset_names = cfg.DATASETS.TEST
for idx, dataset_name in enumerate(dataset_names):
output_folder = output_dir / dataset_name
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders_val = make_data_loader(cfg,
is_train=False,
is_distributed=distributed)
for output_folder, dataset_name, data_loader_val in zip(
output_folders, dataset_names, data_loaders_val):
inference(
model,
data_loader_val,
dataset_name=dataset_name,
iou_types=iou_types,
box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=output_folder,
)
synchronize()
