def _load_model(self, cfg_file, gpu):
if gpu is not None:
torch.cuda.set_device(0)
basepath = rospkg.RosPack().get_path('object_segmentation')
cfg.merge_from_file(basepath + "/" + cfg_file)
logger = setup_logger(distributed_rank=0)
logger.info(f"Loaded configuration file {cfg_file}")
logger.info("Running with config:\n{}".format(cfg))
cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()
# absolute paths of model weights
cfg.MODEL.weights_encoder = (
Path(basepath) / cfg.DIR /
('encoder_' + cfg.TEST.checkpoint)).as_posix()
cfg.MODEL.weights_decoder = (
Path(basepath) / cfg.DIR /
('decoder_' + cfg.TEST.checkpoint)).as_posix()
if not os.path.exists(cfg.MODEL.weights_encoder) or not os.path.exists(
cfg.MODEL.weights_decoder):
download.ycb(Path(basepath) / 'ckpt')
assert os.path.exists(
cfg.MODEL.weights_encoder
), f"checkpoint {cfg.MODEL.weights_encoder} does not exitst!"
assert os.path.exists(
cfg.MODEL.weights_decoder
), f"checkpoint {cfg.MODEL.weights_decoder} does not exitst!"
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
if self.gpu is not None:
segmentation_module.cuda()
segmentation_module.eval()
self.model = segmentation_module
def main():
args = build_argparser().parse_args()
cfg.merge_from_file(args.cfg)
#cfg.merge_from_list(args.opts)
logger = setup_logger(distributed_rank=0) # TODO
logger.info("Loaded configuration file {}".format(args.cfg))
logger.info("Running with config:\n{}".format(cfg))
cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()
# absolute paths of model weights
cfg.MODEL.weights_encoder = os.path.join(
cfg.DIR, 'encoder_' + cfg.TEST.checkpoint)
cfg.MODEL.weights_decoder = os.path.join(
cfg.DIR, 'decoder_' + cfg.TEST.checkpoint)
assert os.path.exists(cfg.MODEL.weights_encoder) and \
os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
seg_model = SegmentationModule(net_encoder, net_decoder, None)
input = torch.randn(1, 3, 320, 320)
#input_decoder = torch.randn(1, 180, 720, 3, 3)
segSize = (256, 256)
#torch.onnx.export(net_encoder, input, "encoder.onnx", opset_version=9, verbose=True)
torch.onnx.export(seg_model, input, "seg.onnx", opset_version=9, verbose=True)
def main():
args = build_argparser().parse_args()
cfg.merge_from_file(args.cfg)
#cfg.merge_from_list(args.opts)
logger = setup_logger(distributed_rank=0) # TODO
logger.info("Loaded configuration file {}".format(args.cfg))
logger.info("Running with config:\n{}".format(cfg))
cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()
# absolute paths of model weights
cfg.MODEL.weights_encoder = os.path.join(cfg.DIR,
'encoder_' + cfg.TEST.checkpoint)
cfg.MODEL.weights_decoder = os.path.join(cfg.DIR,
'decoder_' + cfg.TEST.checkpoint)
assert os.path.exists(cfg.MODEL.weights_encoder) and \
os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"
# Network Builders
net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
'''
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
'''
input = torch.randn(1, 3, 320, 320)
input_decoder = torch.randn(1, 512, 2048, 40, 40)
print(input_decoder[-1])
segSize = (256, 256)
default=0,
type=int,
help="gpu id for evaluation")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg.merge_from_file(args.cfg)
cfg.merge_from_list(args.opts)
# cfg.freeze()
logger = setup_logger(distributed_rank=0) # TODO
logger.info("Loaded configuration file {}".format(args.cfg))
logger.info("Running with config:\n{}".format(cfg))
cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()
# absolute paths of model weights
cfg.MODEL.weights_encoder = os.path.join(cfg.DIR,
'encoder_' + cfg.TEST.checkpoint)
cfg.MODEL.weights_decoder = os.path.join(cfg.DIR,
'decoder_' + cfg.TEST.checkpoint)
assert os.path.exists(cfg.MODEL.weights_encoder) and \
os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"
def inference_prob(
img,
device,
select_model_option="ade20k-resnet50dilated-ppm_deepsup"
): # select_model_option = "ade20k-mobilenetv2dilated-c1_deepsup" / "ade20k-hrnetv2"
'''Load the data and preprocess settings
Input:
img - the path of our target image
device - Current device running
select_model_option - name of NN we use
'''
cfg_ss.merge_from_file("ss/config/" + select_model_option + ".yaml")
logger = setup_logger(distributed_rank=0) # TODO
cfg_ss.MODEL.arch_encoder = cfg_ss.MODEL.arch_encoder.lower()
cfg_ss.MODEL.arch_decoder = cfg_ss.MODEL.arch_decoder.lower()
# absolute paths of model weights
cfg_ss.MODEL.weights_encoder = os.path.join(
'ss/' + cfg_ss.DIR, 'encoder_' + cfg_ss.TEST.checkpoint)
cfg_ss.MODEL.weights_decoder = os.path.join(
'ss/' + cfg_ss.DIR, 'decoder_' + cfg_ss.TEST.checkpoint)
assert os.path.exists(cfg_ss.MODEL.weights_encoder) and os.path.exists(
cfg_ss.MODEL.weights_decoder), "checkpoint does not exist!"
# generate testing image list
imgs = [img]
assert len(imgs), "imgs should be a path to image (.jpg) or directory."
cfg_ss.list_test = [{'fpath_img': x} for x in imgs]
if not os.path.isdir(cfg_ss.TEST.result):
os.makedirs(cfg_ss.TEST.result)
if torch.cuda.is_available():
torch.cuda.set_device(device)
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg_ss.MODEL.arch_encoder,
fc_dim=cfg_ss.MODEL.fc_dim,
weights=cfg_ss.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg_ss.MODEL.arch_decoder,
fc_dim=cfg_ss.MODEL.fc_dim,
num_class=cfg_ss.DATASET.num_class,
weights=cfg_ss.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_test = TestDataset(cfg_ss.list_test, cfg_ss.DATASET)
loader_test = torch.utils.data.DataLoader(
dataset_test,
batch_size=cfg_ss.TEST.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.to(device)
# Main loop
return segmentation_module, loader_test
