def main(cfg, gpu):
torch.cuda.set_device(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 = TestDataset(cfg.list_test, cfg.DATASET)
loader_test = torch.utils.data.DataLoader(
dataset_test,
batch_size=cfg.TEST.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test, gpu)
print('Inference done!')
def worker(cfg, gpu_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(gpu_id)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
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)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu_id, result_queue)
def __init__(self):
"""
extract semantic map from pretrained model, for knowing where to ignore in the image,
since we only have depth info on mountains.
"""
self.names = {}
with open('semseg/object150_info.csv') as f:
reader = csv.reader(f)
next(reader)
for row in reader:
self.names[int(row[0])] = row[5].split(";")[0]
# Network Builders
self.net_encoder = ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights=
'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
)
self.net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights=
'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
use_softmax=True)
self.crit = torch.nn.NLLLoss(ignore_index=-1)
self.segmentation_module = SegmentationModule(self.net_encoder,
self.net_decoder,
self.crit)
self.segmentation_module.eval()
self.segmentation_module.to(device=defs.get_dev())
def initialize():
colors = scipy.io.loadmat('data/color150.mat')['colors']
names = {}
with open('data/object150_info.csv') as f:
reader = csv.reader(f)
next(reader)
for row in reader:
names[int(row[0])] = row[5].split(";")[0]
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth')
net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module = segmentation_module.eval()
x = torch.rand(3,672,504)
output_size = x.shape
# segmentation_module.cuda()
with torch.no_grad():
with torch.jit.optimized_execution(True, {'target_device': 'eia:0'}):
segmentation_module = torch.jit.trace(segmentation_module, x[None])
# torch.jit.save(segmentation_module, 'traced.pt')
return segmentation_module
def segmentation_model_init():
if not os.path.exists(model_folder):
os.makedirs(model_folder)
ENCODER_NAME = 'resnet101'
DECODER_NAME = 'upernet'
PRETRAINED_ENCODER_MODEL_URL = 'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-%s-%s/encoder_epoch_50.pth' % (
ENCODER_NAME, DECODER_NAME)
PRETRAINED_DECODER_MODEL_URL = 'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-%s-%s/decoder_epoch_50.pth' % (
ENCODER_NAME, DECODER_NAME)
pretrained_encoder_file = ENCODER_NAME + basename(
PRETRAINED_ENCODER_MODEL_URL)
pretrained_decoder_file = DECODER_NAME + basename(
PRETRAINED_DECODER_MODEL_URL)
encoder_path = os.path.join(model_folder, pretrained_encoder_file)
decoder_path = os.path.join(model_folder, pretrained_decoder_file)
if not os.path.exists(encoder_path):
urllib.request.urlretrieve(PRETRAINED_ENCODER_MODEL_URL, encoder_path)
if not os.path.exists(decoder_path):
urllib.request.urlretrieve(PRETRAINED_DECODER_MODEL_URL, decoder_path)
# options
options = SimpleNamespace(fc_dim=2048,
num_class=150,
imgSizes=[300, 400, 500, 600],
imgMaxSize=1000,
padding_constant=8,
segm_downsampling_rate=8)
# create model
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=ENCODER_NAME,
weights=encoder_path,
fc_dim=options.fc_dim)
net_decoder = builder.build_decoder(arch=DECODER_NAME,
weights=decoder_path,
fc_dim=options.fc_dim,
num_class=options.num_class,
use_softmax=True)
segmentation_module = SegmentationModule(net_encoder, net_decoder,
torch.nn.NLLLoss(ignore_index=-1))
segmentation_module = segmentation_module.eval()
torch.set_grad_enabled(False)
if torch.cuda.is_available():
segmentation_module = segmentation_module.cuda()
return segmentation_module, options
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 load_model():
net_encoder = ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights='models/encoder_epoch_20.pth')
net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights='models/decoder_epoch_20.pth',
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
#segmentation_module.cuda()
return segmentation_module
def setup(gpu):
global segmentation_module
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights=downloads.download_data_file(
'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth',
'semantic-segmentation-pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
))
net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights=downloads.download_data_file(
'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
'semantic-segmentation-pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth'
),
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()
def model_builder(self, imode="sphe"):
if imode == "sphe":
# Network Builders
net_encoder = seg_sphe.ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth')
net_decoder = seg_sphe.ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
use_softmax=True)
elif imode == "persp":
# encoder_epoch = 'ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
# decoder_epoch = 'ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth'
encoder_epoch = 'ckpt_nef/r50d_ppm_rot_e40_nef_30/encoder_epoch_40.pth'
decoder_epoch = 'ckpt_nef/r50d_ppm_rot_e40_nef_30/decoder_epoch_40.pth'
net_encoder = seg_persp.ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights=encoder_epoch)
net_decoder = seg_persp.ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights=decoder_epoch,
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
semseg_model = SegmentationModule(net_encoder, net_decoder, crit)
semseg_model.eval()
semseg_model.cuda()
return semseg_model
def __init__(self, encoder_weights, decoder_weights):
super().__init__()
self.impl = SegmentationModule(
net_enc=ModelBuilder.build_encoder(arch="hrnetv2",
fc_dim=720,
weights=encoder_weights),
net_dec=ModelBuilder.build_decoder(arch="c1",
fc_dim=720,
num_class=150,
weights=decoder_weights,
use_softmax=True),
crit=None,
)
def main(cfg, gpus):
# Network Builders
# net_encoder = ModelBuilder.build_encoder(
# arch=cfg.MODEL.arch_encoder.lower(),
# fc_dim=cfg.MODEL.fc_dim,
# weights=cfg.MODEL.weights_encoder)
# net_decoder = ModelBuilder.build_decoder(
# arch=cfg.MODEL.arch_decoder.lower(),
# fc_dim=cfg.MODEL.fc_dim,
# num_class=cfg.DATASET.num_class,
# weights=cfg.MODEL.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit, cfg.TRAIN.deep_sup_scale)
elif cfg.Model.arch_decoder.endswith('deeplab'):
net = DeepLabV3(num_classes=7)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit, cfg.TRAIN.deep_sup_scale))
else:
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit)
def create_network(cfg, gpu):
torch.cuda.set_device(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)
segmentation_module.cuda()
# Main loop
return segmentation_module, gpu
print('Inference done!')
def load_model_from_cfg(cfg):
model_config, encoder_path, decoder_path = parse_model_config(cfg)
net_encoder = ModelBuilder.build_encoder(
arch=model_config["MODEL"]['arch_encoder'],
fc_dim=model_config['MODEL']['fc_dim'],
weights=encoder_path)
net_decoder = ModelBuilder.build_decoder(
arch=model_config["MODEL"]['arch_decoder'],
fc_dim=model_config['MODEL']['fc_dim'],
num_class=model_config['DATASET']['num_class'],
weights=decoder_path,
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
return segmentation_module
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 __init__(self, model_name, model_path):
self.model_name = model_name
self.model_path = model_path
#
torch.cuda.set_device(0)
# Network Builders
net_encoder = ModelBuilder.build_encoder(arch="hrnetv2",
fc_dim=720,
weights=self.model_path)
net_decoder = ModelBuilder.build_decoder(arch="c1",
fc_dim=720,
num_class=2,
weights=self.model_path,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
self.segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
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 semantic_segmentation(image_folder):
download_segmentation()
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch='resnet101dilated',
fc_dim=2048,
weights='data/segm_data/encoder_epoch_25.pth')
net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights='data/segm_data/decoder_epoch_25.pth',
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()
# Load and normalize one image as a singleton tensor batch
pil_to_tensor = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], # These are RGB mean+std values
std=[0.229, 0.224, 0.225]) # across a large photo dataset.
])
pil_image = PIL.Image.open(os.path.join(image_folder,
'000000.jpg')).convert('RGB')
img_original = np.array(pil_image)
img_data = pil_to_tensor(pil_image)
singleton_batch = {'img_data': img_data[None].cuda()}
output_size = img_data.shape[1:]
# Run the segmentation at the highest resolution.
with torch.no_grad():
scores = segmentation_module(singleton_batch, segSize=output_size)
# Get the predicted scores for each pixel
_, pred = torch.max(scores, dim=1)
pred = pred.cpu()[0].numpy()
del net_encoder, net_decoder, segmentation_module
return pred
def load_model_from_url(semsegPath):
model_urls = {
'encoder':
'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth',
'decoder':
'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth'
}
if not os.path.exists(os.path.join(semsegPath, 'ckpt')):
os.makedirs(os.path.join(semsegPath, 'ckpt'))
r = requests.get(model_urls['encoder'])
with open(os.path.join(semsegPath, 'ckpt/encoder.pth'), 'wb') as f:
f.write(r.content)
r = requests.get(model_urls['decoder'])
with open(os.path.join(semsegPath, 'ckpt/decoder.pth'), 'wb') as f:
f.write(r.content)
net_encoder = ModelBuilder.build_encoder(arch='resnet50dilated',
fc_dim=2048,
weights=os.path.join(
semsegPath,
'ckpt/encoder.pth'))
net_decoder = ModelBuilder.build_decoder(arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights=os.path.join(
semsegPath,
'ckpt/decoder.pth'),
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()
return segmentation_module
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
def main(cfg, gpus):
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder)
if cfg.MODEL.arch_decoder.endswith('regression'):
crit = nn.MSELoss(
reduction="sum"
) # Sum for multi-output learning, need to sum across all labels
else:
crit = nn.NLLLoss(ignore_index=-1) # negative log likelihood loss
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg.DATASET.classes,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg.DATASET.classes)
print("net_encoder")
print(type(net_encoder))
print(net_encoder)
print("net_decoder")
print(type(net_decoder))
print(net_decoder)
# Dataset and Loader
if cfg.MODEL.arch_decoder.endswith('regression'):
print("performing regression")
dataset_train = TrainDatasetRegression(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET.classes,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
dataset_val = TrainDatasetRegression(
cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET.classes,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
else:
dataset_train = TrainDataset(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
loader_val = torch.utils.data.DataLoader(
dataset_val,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
iterator_val = iter(loader_val)
# load nets into gpu
if len(gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, cfg)
# Main loop
history = {
'train': {
'epoch': [],
'loss': []
},
'val': {
'epoch': [],
'loss': []
}
}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
val(segmentation_module, iterator_val, optimizers, history, epoch + 1,
cfg)
# checkpointing every 5th epoch
if (epoch % 5 == 0):
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
def main(cfg, gpus):
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_train = TrainDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
print(gpus)
# if len(gpus) > 1:
if True:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, cfg)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
# checkpointing
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
# weights='ckpt/ade20k-resnet101-upernet/decoder_epoch_50.pth',
# use_softmax=True)
# hrnet
net_encoder = ModelBuilder.build_encoder(
arch='hrnetv2',
fc_dim=720,
weights='../../ckpt/ade20k-hrnetv2-c1/encoder_epoch_30.pth')
net_decoder = ModelBuilder.build_decoder(
arch='c1',
fc_dim=720,
num_class=150,
weights='../../ckpt/ade20k-hrnetv2-c1/decoder_epoch_30.pth',
use_softmax=True)
crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()
pil_to_tensor = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], # These are RGB mean+std values
std=[0.229, 0.224, 0.225]) # across a large photo dataset.
])
pil_image = PIL.Image.open('sample.png').convert('RGB')
img_original = numpy.array(pil_image)
img_data = pil_to_tensor(pil_image)
singleton_batch = {'img_data': img_data[None].cuda()}
output_size = img_data.shape[1:]
class ExtractSegmentationMask:
def __init__(self):
"""
extract semantic map from pretrained model, for knowing where to ignore in the image,
since we only have depth info on mountains.
"""
self.names = {}
with open('semseg/object150_info.csv') as f:
reader = csv.reader(f)
next(reader)
for row in reader:
self.names[int(row[0])] = row[5].split(";")[0]
# Network Builders
self.net_encoder = ModelBuilder.build_encoder(
arch='resnet50dilated',
fc_dim=2048,
weights=
'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
)
self.net_decoder = ModelBuilder.build_decoder(
arch='ppm_deepsup',
fc_dim=2048,
num_class=150,
weights=
'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
use_softmax=True)
self.crit = torch.nn.NLLLoss(ignore_index=-1)
self.segmentation_module = SegmentationModule(self.net_encoder,
self.net_decoder,
self.crit)
self.segmentation_module.eval()
self.segmentation_module.to(device=defs.get_dev())
def visualize_result(self, img, pred, index=None):
# filter prediction class if requested
colors = scipy.io.loadmat('semseg/color150.mat')['colors']
if index is not None:
pred = pred.copy()
pred[pred != index] = -1
print(f'{self.names[index + 1]}:')
# colorize prediction
pred_color = colorEncode(pred, colors).astype(np.float32) / 255
# aggregate images and save
im_vis = np.concatenate((img, pred_color), axis=1)
plt.imshow(im_vis)
plt.show()
def __call__(self, sample):
img_data = sample['image']
old_shape = img_data.shape[1], img_data.shape[2]
img_data = ResizeToAlmostResolution(180, 224)({
'image': img_data
})['image']
singleton_batch = {'img_data': img_data[None]}
output_size = img_data.shape[1:]
# Run the segmentation at the highest resolution.
with torch.no_grad():
scores = self.segmentation_module(singleton_batch,
segSize=output_size)
# Get the predicted scores for each pixel
_, pred = torch.max(scores, dim=1)
pred = TF.resize(pred, old_shape, Image.NEAREST)
# other irrelevant classes: 1, 4, 12, 20, 25, 83, 116, 126, 127.
# see csv in semseg folder.
bad_classes = torch.Tensor([2]).to(device=defs.get_dev())
mask = torch.full_like(pred, True, dtype=torch.bool)
mask[(pred[..., None] == bad_classes).any(-1)] = False
if 'mask' in sample:
sample['mask'] = sample['mask'] & mask
else:
sample['mask'] = mask
return sample
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
# net_weight_1='/home/admins/qyl/huawei_compete/semantic-segmentation-pytorch/ckpt/ade20k-hrnetv2-c1/encoder_epoch_20.pth'
# net_weight_2='/home/admins/qyl/huawei_compete/semantic-segmentation-pytorch/ckpt/ade20k-hrnetv2-c1/decoder_epoch_20.pth'
# dic
#
# ["encoder"]=torch.load(net_weight_1)
# dic["decoder"]=torch.load(net_weight_2)
# torch.save(dic,"model_best_20.pth")
# print("merge model finished..........")
#
# Network Builders
model_path = 'model_best.pth'
net_encoder = ModelBuilder.build_encoder(arch="hrnetv2",
fc_dim=720,
weights=model_path)
net_decoder = ModelBuilder.build_decoder(arch="c1",
fc_dim=720,
num_class=2,
weights=model_path,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
image_demo_dir = "./demo/182_28_9.png"
imags = Image.open(image_demo_dir)
imags = np.array(imags)
imags = imags[512:1024, 0:512, :]
print(imags.shape)
imags = Image.fromarray(np.uint8(imags))
pred = huawei_seg(imags, segmentation_module)
print(pred)