def main():
charts_cfg = cfg.clone()
tables_cfg = cfg.clone()
charts_cfg.merge_from_file(charts_path)
tables_cfg.merge_from_file(tables_path)
model_charts = ChataDemo(charts_cfg, "charts")
model_tables = ChataDemo(tables_cfg, "tables")
client = ApiClient()
client.login()
newest_ids = client.get_new_publications()
for pub_id in newest_ids:
pages = client.get_pages(pub_id)
print(pub_id)
for id, url in pages:
print(id)
print(url)
resp = urllib.request.urlopen(url)
image = np.asarray(bytearray(resp.read()), dtype="uint8")
image = cv2.imdecode(image, cv2.IMREAD_COLOR)
predictions_charts = model_charts.get_predictions(image)
predictions_tables = model_tables.get_predictions(image)
annotations_charts = annotate(predictions_charts)
annotations_tables = annotate(predictions_tables)
annotations = annotations_charts+annotations_tables
print(annotations)
if annotations:
client.add_annotation(id, annotations)
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 __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 __init__(self, cfg):
super(REID_Baseline, self).__init__()
self.cfg = cfg.clone()
self.base = ResNet(cfg.REID.MODEL.LAST_STRIDE)
self.base.load_param(cfg.REID.MODEL.PRETRAIN_PATH)
self.gap = nn.AdaptiveAvgPool2d(1)
self.num_classes = cfg.REID.DATASETS.ID_NUM
# if cfg.MODEL.MASK_ON:
# self.in_planes = 4096
# else:
# self.in_planes = 2048
self.in_planes = 2048
self.reid_feat_dim = 256
self.bottleneck = nn.BatchNorm1d(self.in_planes)
self.bottleneck.bias.requires_grad_(False) # no shift
self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)
self.bottleneck.apply(weights_init_kaiming)
self.classifier.apply(weights_init_classifier)
self.queue_size = 5000
self.margin = self.cfg.TRIPLET.MARGIN
self.lut_momentum = 0.0 ###
self.reid_feat_dim = self.in_planes
# self.register_buffer('lut', torch.zeros(self.num_pid, self.reid_feat_dim).cuda())
self.register_buffer('lut1', torch.zeros(self.num_classes, self.reid_feat_dim).cuda())
self.register_buffer('lut2', torch.zeros(self.num_classes, self.reid_feat_dim).cuda())
self.register_buffer('queue', torch.zeros(self.queue_size, self.reid_feat_dim).cuda())
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)
for n, m in self.model.named_modules():
if n == "roi_heads":
m.register_forward_hook(hook)
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,
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)
logging.info('model loaded from: {}'.format(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):
super(RPNModule, self).__init__()
self.cfg = cfg.clone()
anchor_generator = make_anchor_generator(cfg)
in_channels = cfg.MODEL.BACKBONE.OUT_CHANNELS
rpn_head = registry.RPN_HEADS[cfg.MODEL.RPN.RPN_HEAD]
head = rpn_head(cfg, in_channels,
anchor_generator.num_anchors_per_location()[0])
rpn_box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))
box_selector_train = make_rpn_postprocessor(cfg,
rpn_box_coder,
is_train=True)
box_selector_test = make_rpn_postprocessor(cfg,
rpn_box_coder,
is_train=False)
loss_evaluator = make_rpn_loss_evaluator(cfg, rpn_box_coder)
self.anchor_generator = anchor_generator
self.head = head
self.box_selector_train = box_selector_train
self.box_selector_test = box_selector_test
self.loss_evaluator = loss_evaluator
def __init__(self, cfg_path, weights_path, input_shape=(608, 608)):
cfg.merge_from_file(cfg_path)
cfg.merge_from_list(['DTYPE', 'float16'])
self._cfg = cfg.clone()
self._model = build_detection_model(self._cfg)
self._model.eval()
self._device = 'cuda'
self._model.to(self._device)
self.shape = input_shape
save_dir = cfg.OUTPUT_DIR
checkpoint = torch.load(weights_path, map_location=torch.device("cpu"))
load_state_dict(self._model, checkpoint.pop("model"))
self._transform = self._build_transform()
def __init__(
self,
cfg,
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.min_image_size = min_image_size
checkpointer = DetectronCheckpointer(cfg,self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.model.to(self.device)
self.transforms = self.build_trainsform()
self.cpu_device = torch.device("cpu")
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 __init__(self, cfg, image_size=(100, 100)):
super(OPN, self).__init__()
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.model.eval()
self.backbone = self.model.backbone
self.rpn = self.model.rpn
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.image_size = image_size # (h, w)
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
self.cpu_device = torch.device("cpu")
def __init__(self, cfg):
super(RetinaNetModule, self).__init__()
self.cfg = cfg.clone()
anchor_generator = make_anchor_generator_retinanet(cfg)
head = RetinaNetHead(cfg)
box_coder = BoxCoder(weights=(10., 10., 5., 5.))
if self.cfg.MODEL.SPARSE_MASK_ON:
raise NotImplementedError
else:
box_selector_test = make_retinanet_postprocessor(
cfg, 100, box_coder)
self.anchor_generator = anchor_generator
self.head = head
self.box_selector_test = box_selector_test
def __init__(self,
cfg,
confidence_threshold=0.7,
min_image_size=224,
output_polygon=True):
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()
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
self.output_polygon = output_polygon
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 __init__(
self,
config_file,
checkpoint_path=None
):
cfg.merge_from_file(config_file)
cfg.defrost()
if checkpoint_path:
cfg.MODEL.WEIGHT = checkpoint_path
cfg.freeze()
self.cfg = cfg.clone()
self.model = build_detection_model(cfg)
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.cpu_device = torch.device("cpu")
self.checkpointer = DetectronCheckpointer(cfg, self.model)
_ = self.checkpointer.load(cfg.MODEL.WEIGHT)
if hasattr(self.checkpointer, "classes"):
self.classes = self.checkpointer.classes
def __init__(self, cfg: CfgNode, args):
self.args = args
self.cfg = cfg.clone()
self.model: nn.Module = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
# Load weight
save_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=save_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.cpu_device = torch.device('cpu')
# Keep all result
self.confidence_threshold = 0.
self.datasets = GIFDataset(cfg, args)
self.results = defaultdict(list)
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 main(args, skip_test=False):
cfg = c.clone()
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
args.distributed = num_gpus > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl", init_method="env://"
)
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.OUTPUT_DIR = os.path.join("output", cfg.OUTPUT_DIR, "train")
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir:
mkdir(output_dir)
logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(args)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
logger.info("Loaded configuration file {}".format(args.config_file))
with open(args.config_file, "r") as cf:
config_str = "\n" + cf.read()
logger.info(config_str)
logger.info("Running with config:\n{}".format(cfg))
model = train(cfg, args.local_rank, args.distributed)
if (not args.skip_test) and (not skip_test):
test(cfg, model, args.distributed)
def __init__(self,
config_file,
confidence_threshold=0.3,
mask_threshold=0.5,
mask_on=True):
cfg.merge_from_file(config_file)
cfg.MODEL.MASK_ON = mask_on
cfg.freeze()
self.cfg = cfg.clone()
self.model = build_detection_model(cfg).eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
checkpointer = DetectronCheckpointer(cfg,
self.model,
save_dir=cfg.OUTPUT_DIR)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
self.masker = Masker(threshold=mask_threshold, padding=1)
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
def __init__(self,
cfg,
conf_thresh=0.7,
mask_thresh=0.5,
masks_per_dim=2,
min_image_size=224,
ckpt=None):
assert 0.0 <= conf_thresh <= 1.0 and isinstance(conf_thresh, float)
assert 0.0 <= mask_thresh <= 1.0 and isinstance(mask_thresh, float)
assert masks_per_dim > 0 and isinstance(masks_per_dim, int)
assert min_image_size > 0 and isinstance(min_image_size, int)
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)
if ckpt:
checkpointer.load_model(checkpointer.load_file(ckpt))
else:
_ = checkpointer.load()
self.transforms = self.build_transform()
self.masker = Masker(threshold=mask_thresh, 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.conf_thresh = conf_thresh
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,
cache_loc='cache/unknown_object_detection', image_extension='.jpg',
overwrite=False):
self.cache_loc = cache_loc
if not os.path.exists(self.cache_loc):
os.makedirs(self.cache_loc)
self.extension_length = len(image_extension)
self.overwrite = overwrite
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,
config,
):
self.confidence_threshold = 0.5
self.min_image_size = 800
self.mask_shrink = 4
# set cpu/gpu
self.cpu_only = True
visible_device_list = []
for handle in config.devices:
if int(handle.type) == DeviceType.GPU.value:
visible_device_list.append(handle.id)
self.cpu_only = False
if not self.cpu_only:
print('Using GPU: {}'.format(visible_device_list[0]))
torch.cuda.set_device(visible_device_list[0])
else:
print('Using CPU')
# set maskrcnn config
cfg.merge_from_file(CONFIG_FILE)
if self.cpu_only:
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
else:
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
cfg.merge_from_list(["INPUT.TO_BGR255", False])
self.cfg = cfg.clone()
# build model and transform
self.model = build_detection_model(cfg)
self.model.eval()
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.transforms = self.build_transform()
self.masker = Masker(threshold=0.5, padding=1)
self.cpu_device = torch.device("cpu")
def __init__(self, cfg, confidence_threshold=0.5):
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)
checkpointer = DetectronCheckpointer(cfg, self.model)
_ = checkpointer.load(cfg.MODEL.WEIGHT)
self.transforms = self.build_transform()
self.cpu_device = torch.device("cpu")
self.confidence_threshold = confidence_threshold
show_mask_heatmaps = True
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.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
masks_per_dim = 2
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 Training")
parser.add_argument(
"--config-file",
default="",
metavar="FILE",
help="path to config file",
type=str,
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"--skip-test",
dest="skip_test",
help="Do not test the final model",
action="store_true",
)
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
args.distributed = num_gpus > 1
if args.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.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 31
cfg.SOLVER.BASE_LR = 0.01
cfg_origial = cfg.clone()
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir:
mkdir(output_dir)
logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(args)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
logger.info("Loaded configuration file {}".format(args.config_file))
with open(args.config_file, "r") as cf:
config_str = "\n" + cf.read()
logger.info(config_str)
logger.info("Running with config:\n{}".format(cfg))
model = train(cfg, cfg_origial, args.local_rank, args.distributed)
if not args.skip_test:
run_test(cfg, model, args.distributed)
def load_cfg(conf):
config = cfg.clone()
config.merge_from_file(conf.config_file)
config.freeze()
return config
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Inference")
parser.add_argument(
"--config",
default="maskrcnn_benchmark/data/datasets/predictor.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
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)
# cfg.merge_from_list(args.opts)
# cfg.freeze()
mycfg = cfg.clone()
save_dir = ""
logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
model = build_detection_model(mycfg)
model.eval()
model.to(mycfg.MODEL.DEVICE)
output_dir = cfg.OUTPUT_DIR
# checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
# _ = checkpointer.load(cfg.MODEL.WEIGHT)
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
if cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference",
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()
def __init__(self, cfg):
self.cfg = cfg.clone()
self.transforms = self.build_transform()
def train_with_validation(cfg, local_rank, distributed, test_weights=None):
arguments = {}
arguments["iteration"] = 0
if test_weights:
cfg.MODEL.WEIGHT = test_weights
cfg.SOLVER.MAX_ITER = 0
ignore_labels = (cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES == 0)
# prepare training data
root_path = os.path.expanduser(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "data")))
data_loader, class_ids = make_data_loader(
root_path,
cfg,
is_train=True,
is_distributed=distributed,
start_iter=arguments["iteration"],
ignore_labels=ignore_labels,
)
# overwrite the number of classes by considering the training set
if cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES <= 0: # if we have binary classification or unknown number of classes
cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = len(class_ids)
# prepare model
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# prepare optimizer
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
)
# prepare validation
run_validation_for_model = partial(run_validation, root_path=root_path, cfg=cfg.clone(), class_ids=class_ids, ignore_labels=ignore_labels, distributed=distributed)
# setup checkpointer
output_dir = cfg.OUTPUT_DIR
save_to_disk = get_rank() == 0
checkpointer = DetectronCheckpointer(
cfg, model, optimizer, scheduler, output_dir, save_to_disk
)
extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, use_latest=False if test_weights else True)
arguments.update(extra_checkpoint_data)
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
validation_period = cfg.SOLVER.VALIDATION_PERIOD
# start training
do_train(
model,
data_loader,
optimizer,
scheduler,
checkpointer,
device,
validation_period,
checkpoint_period,
arguments,
run_validation_for_model)
return model
