def _load_mistcsail_model(encoder, decoder, encoder_path=None, decoder_path=None):
from mit_semseg.models import ModelBuilder, SegmentationModule
from mit_semseg.config import cfg
if encoder_path is None:
encoder_path = str(Path(__file__).parents[0].joinpath('weights', f'encoder_{encoder}.pth'))
if decoder_path is None:
decoder_path = str(Path(__file__).parents[0].joinpath('weights', f'decoder_{decoder}.pth'))
config_file = str(Path(__file__).parents[0].joinpath('config', f'ade20k-{encoder}-{decoder}.yaml'))
cfg.merge_from_file(config_file)
net_encoder = ModelBuilder.build_encoder(
arch=encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=encoder_path)
net_decoder = ModelBuilder.build_decoder(
arch=decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=decoder_path,
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
model = SegmentationModule(net_encoder, net_decoder, crit)
return model
def predict(imgPath):
args = Namespace(imgs=imgPath,
cfg="config/ade20k-resnet50dilated-ppm_deepsup.yaml",
opts=None)
cfg.merge_from_file(args.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!"
imgs = [args.imgs]
assert len(imgs), "imgs should be a path to image (.jpg) or directory."
cfg.list_test = [{'fpath_img': x} for x in imgs]
if not os.path.isdir(cfg.TEST.result):
os.makedirs(cfg.TEST.result)
prediction = main(cfg, None)
return prediction
def huawei_seg(imgs, segmentation_module):
#
cfg.merge_from_file("/home/mind/model/config/ade20k-hrnetv2-huawei.yaml")
imgs = [imgs]
cfg.list_test = [{'fpath_img': x} for x in imgs]
# Dataset and Loader
dataset_test = InferDataset(cfg.list_test, cfg.DATASET)
loader_test = torch.utils.data.DataLoader(
dataset_test,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
loader = loader_test
# Main loop
segmentation_module.eval()
pbar = tqdm(total=len(loader))
for batch_data in loader:
# process data
batch_data = batch_data[0]
segSize = (batch_data['img_ori'].shape[0],
batch_data['img_ori'].shape[1])
img_resized_list = batch_data['img_data']
with torch.no_grad():
scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0],
segSize[1])
scores = async_copy_to(scores, 0)
for img in img_resized_list:
feed_dict = batch_data.copy()
feed_dict['img_data'] = img
del feed_dict['img_ori']
del feed_dict['info']
feed_dict = async_copy_to(feed_dict, 0)
# forward pass
pred_tmp = segmentation_module(feed_dict, segSize=segSize)
# print(pred_tmp.shape)#torch.Size([1, 2, 1024, 1024])
scores = scores + pred_tmp / len(cfg.DATASET.imgSizes)
_, pred = torch.max(scores, dim=1)
pred = as_numpy(pred.squeeze(0).cpu())
#
# visualize_result(
# (batch_data['img_ori'], batch_data['info']),
# pred,
# cfg
# )
pbar.update(1)
#
return pred
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)
def huawei_seg(imgs):
parser = argparse.ArgumentParser(
description="PyTorch Semantic Segmentation Testing")
parser.add_argument(
"--cfg",
default="config/ade20k-hrnetv2-huawei.yaml",
metavar="FILE",
help="path to config file",
type=str,
)
parser.add_argument("--gpu",
default=0,
type=int,
help="gpu id for evaluation")
args = parser.parse_args()
cfg.merge_from_file(args.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)
#
imgs = [imgs]
cfg.list_test = [{'fpath_img': x} for x in imgs]
torch.cuda.set_device(args.gpu)
# 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)
# Dataset and Loader
dataset_test = InferDataset(cfg.list_test, cfg.DATASET)
loader_test = torch.utils.data.DataLoader(
dataset_test,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
loader = loader_test
# Main loop
segmentation_module.eval()
pbar = tqdm(total=len(loader))
for batch_data in loader:
# process data
batch_data = batch_data[0]
segSize = (batch_data['img_ori'].shape[0],
batch_data['img_ori'].shape[1])
img_resized_list = batch_data['img_data']
with torch.no_grad():
scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0],
segSize[1])
scores = async_copy_to(scores, args.gpu)
for img in img_resized_list:
feed_dict = batch_data.copy()
feed_dict['img_data'] = img
del feed_dict['img_ori']
del feed_dict['info']
feed_dict = async_copy_to(feed_dict, args.gpu)
# forward pass
pred_tmp = segmentation_module(feed_dict, segSize=segSize)
# print(pred_tmp.shape)#torch.Size([1, 2, 1024, 1024])
scores = scores + pred_tmp / len(cfg.DATASET.imgSizes)
_, pred = torch.max(scores, dim=1)
pred = as_numpy(pred.squeeze(0).cpu())
# visualization
visualize_result((batch_data['img_ori'], batch_data['info']), pred,
cfg)
pbar.update(1)
#
return pred
help="path to config file",
type=str,
)
parser.add_argument("--gpu",
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)
parser.add_argument("--input", default=f'{TEST_DIR}\\input.mp4', help="Path to input video file")
parser.add_argument("--output", default=f'{TEST_DIR}\\output.mp4', help="Path to output style video file")
parser.add_argument("--label", default="car", help="Path to output style video file")
parser.add_argument("--use-cpu", action='store_true', help="Use CPU instead of GPU")
parser.add_argument("--gpu-device", type=int, default=None, help="GPU device index")
parser.add_argument("--fps", type=int, default=None, help="FPS of output video")
args = parser.parse_args()
label_index = -1
with open(LABEL_CSV, 'r') as csv_file:
csv_reader = csv.reader(csv_file)
for row in csv_reader:
if label_index > 0: break
for column in row:
if (args.label in column): label_index = int(row[0]) - 1
style_model = StyleModel(DEFAULT_MODEL, use_gpu=not args.use_cpu, gpu_device=args.gpu_device)
cfg.merge_from_file("config/ade20k-resnet50dilated-ppm_deepsup.yaml")
segmentation_model = SegmentationModel(
DEFAULT_MODEL,
use_gpu=not args.use_cpu,
gpu_device=args.gpu_device,
cfg=cfg)
fix_path = lambda x: x.replace('\\', r'\\')
reader = VideoReader(fix_path(args.input), fps=args.fps)
output_writer = VideoWriter(fix_path(args.output), reader.width, reader.height, reader.fps)
with output_writer:
frame_number = 0
for frame in reader:
print(f"Processing frame {frame_number}")
frame_number+=1
image = np.array(frame)
style = style_model.run(image)
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
