itr = iter(dataset_it)
# load model
model = get_model(args['model']['name'], args['model']['kwargs'])
model = torch.nn.DataParallel(model).to(device)
# load snapshot
if os.path.exists(args['checkpoint_path']):
state = torch.load(args['checkpoint_path'])
model.load_state_dict(state['model_state_dict'], strict=True)
#else:
# assert(False, 'checkpoint_path {} does not exist!'.format(args['checkpoint_path']))
model.eval()
# cluster module
cluster = Cluster()
# Visualizer
visualizer = Visualizer(('image', 'pred', 'sigma', 'seed'))
with torch.no_grad():
sample = next(itr)
im = sample['image']
instances = sample['instance'].squeeze()
output = model(im)
instance_map, predictions = cluster.cluster(output[0],
n_sigma=2,
threshold=0.9)
if args['display']:
optimizer = torch.optim.Adam(model.parameters(),
lr=args['lr'],
weight_decay=1e-4)
def lambda_(epoch):
return pow((1 - ((epoch) / args['n_epochs'])), 0.9)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda_,
)
# clustering
cluster = Cluster()
# Visualizer
visualizer = Visualizer(('image', 'pred', 'sigma', 'seed'))
# Logger
logger = Logger(('train', 'val', 'iou'), 'loss')
# resume
start_epoch = 0
best_iou = 0
if args['resume_path'] is not None and os.path.exists(args['resume_path']):
print('Resuming model from {}'.format(args['resume_path']))
state = torch.load(args['resume_path'])
start_epoch = state['epoch'] + 1
best_iou = state['best_iou']
model = get_model(args['model']['name'], args['model']['kwargs']['num_classes'], pretrainedEnc)
model = torch.nn.DataParallel(model).to(device)
# load snapshot
if os.path.exists(args['checkpoint_path']):
print("model_loaded")
state = torch.load(args['checkpoint_path'])
model.load_state_dict(state['model_state_dict'], strict=True)
else:
assert (False, 'checkpoint_path {} does not exist!'.format(args['checkpoint_path']))
model.eval()
batch_size = 1
cluster = Cluster()
time_train = []
it = iter(dataset_it)
for i in range(len(dataset_it)):
im = next(it)['image']
start_time = time.time()
with torch.no_grad():
output = model(im)
label = output[0][0].max(0)[1]
cluster.cluster(output[1][0], label, threshold=0.9)
torch.cuda.synchronize()
if i > 15:
fwt = time.time() - start_time
model = get_model(args['model']['name'],
args['model']['kwargs']['num_classes'], pretrainedEnc)
model = torch.nn.DataParallel(model).to(device)
if os.path.exists(args['checkpoint_path']):
print("model_loaded")
state = torch.load(args['checkpoint_path'])
model.load_state_dict(state['model_state_dict'], strict=True)
else:
assert (False, 'checkpoint_path {} does not exist!'.format(
args['checkpoint_path']))
model.eval()
# cluster module
cluster = Cluster()
with torch.no_grad():
for ind, sample in enumerate(tqdm(dataset_it)):
im = sample['image']
output = model(im)
im_name, _ = os.path.splitext(os.path.basename(
sample['im_name'][0]))[0].split("_leftImg8bit")
label = output[0][0].max(0)[1]
label = TrainIdToIDS(label)
instance_map = cluster.cluster(output[1][0], label, threshold=0.9)
ins = torchvision.transforms.ToPILImage()(instance_map.int().cpu())
ins.save(args['save_dir'] + im_name + "_gtFine_instanceIds" + ".png")
# dataloader
dataset = get_dataset(args['dataset']['name'], args['dataset']['kwargs'])
dataset_it = torch.utils.data.DataLoader(
dataset,
batch_size=1,
shuffle=False,
drop_last=False,
num_workers=4,
pin_memory=True if args['cuda'] else False)
# load model
model = get_model(args['model']['name'], args['model']['kwargs'])
model = torch.nn.DataParallel(model).to(device)
# cluster module
cluster = Cluster()
xm = torch.linspace(0, 2, 2048).view(1, 1, -1).expand(1, 1024, 2048)
ym = torch.linspace(0, 1, 1024).view(1, -1, 1).expand(1, 1024, 2048)
xym = torch.cat((xm, ym), 0).cuda()
def prepare_img(image):
if isinstance(image, Image.Image):
return image
if isinstance(image, torch.Tensor):
image.squeeze_()
image = image.numpy()
if isinstance(image, np.ndarray):