class Model:
def __init__(self, args):
# Define encoder and decoder
self.encoder = None
self.decoder = None
# Mode
self.video_mode = False
# Load model
self._init_model(args)
def _init_model(self, args):
print("Loading model: " + args.model_name)
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(
args.model_name, args.use_gpu)
load_args.use_gpu = args.use_gpu
self.encoder = FeatureExtractor(load_args)
if args.zero_shot:
self.decoder = RSIS(load_args)
else:
self.decoder = RSISMask(load_args)
print(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
encoder_dict, decoder_dict = check_parallel(encoder_dict, decoder_dict)
self.encoder.load_state_dict(encoder_dict)
to_be_deleted_dec = []
for k in decoder_dict.keys():
if 'fc_stop' in k:
to_be_deleted_dec.append(k)
for k in to_be_deleted_dec:
del decoder_dict[k]
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
if load_args.length_clip == 1:
self.video_mode = False
print('video mode not activated')
else:
self.video_mode = True
print('video mode activated')
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.dataset = args.dataset
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
if args.dataset == 'youtube':
dataset = get_dataset(args,
split=self.split,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment
and self.split == 'train',
inputRes=(256, 448),
video_mode=True,
use_prev_mask=False)
else:
dataset = get_dataset(args,
split=self.split,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment
and self.split == 'train',
video_mode=True,
use_prev_mask=False)
self.loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.args = args
print(args.model_name)
encoder_dict, decoder_dict, enc_opt_dict, dec_opt_dict, load_args = load_checkpoint(
args.model_name, args.use_gpu)
load_args.use_gpu = args.use_gpu
self.encoder = FeatureExtractor(load_args)
self.decoder = RSIS(load_args)
print(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
encoder_dict, decoder_dict = check_parallel(encoder_dict, decoder_dict)
self.encoder.load_state_dict(encoder_dict)
to_be_deleted_dec = []
for k in decoder_dict.keys():
if 'fc_stop' in k:
to_be_deleted_dec.append(k)
for k in to_be_deleted_dec:
del decoder_dict[k]
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
if load_args.length_clip == 1:
self.video_mode = False
print('video mode not activated')
else:
self.video_mode = True
print('video mode activated')
def run_eval(self):
print("Dataset is %s" % (self.dataset))
print("Split is %s" % (self.split))
if args.overlay_masks:
colors = []
palette = sequence_palette()
inv_palette = {}
for k, v in palette.items():
inv_palette[v] = k
num_colors = len(inv_palette.keys())
for id_color in range(num_colors):
if id_color == 0 or id_color == 21:
continue
c = inv_palette[id_color]
colors.append(c)
if self.split == 'val':
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results')
make_dir(results_dir)
json_data = open(
'../../databases/YouTubeVOS/train/train-val-meta.json')
data = json.load(json_data)
else:
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess-davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results-davis')
make_dir(results_dir)
for batch_idx, (inputs, targets, seq_name,
starting_frame) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
for ii in range(max_ii):
# x: input images (N consecutive frames from M different sequences)
# y_mask: ground truth annotations (some of them are zeros to have a fixed length in number of object instances)
# sw_mask: this mask indicates which masks from y_mask are valid
x, y_mask, sw_mask = batch_to_var(args, inputs[ii],
targets[ii])
print(seq_name[0] + '/' + '%05d' %
(starting_frame[0] + ii))
#from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test(
args, self.encoder, self.decoder, x,
prev_hidden_temporal_list)
if args.dataset == 'youtube':
num_instances = len(
data['videos'][seq_name[0]]['objects'])
else:
num_instances = 1 #int(torch.sum(sw_mask.data).data.cpu().numpy())
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
for t in range(10):
mask_pred = (torch.squeeze(outs[0,
t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
toimage(mask2assess, cmin=0,
cmax=255).save(base_dir_masks_sep +
'%05d_instance_%02d.png' %
(starting_frame[0] + ii, t))
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure()
plt.axis('off')
plt.figure()
plt.axis('off')
plt.imshow(frame_img)
for t in range(num_instances):
mask_pred = (torch.squeeze(
outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones(
(mask_pred.shape[0], mask_pred.shape[1], 3))
color_mask = np.array(colors[t]) / 255.0
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
figname = base_dir + 'frame_%02d.png' % (
starting_frame[0] + ii)
plt.savefig(figname, bbox_inches='tight')
plt.close()
if self.video_mode:
prev_hidden_temporal_list = hidden_temporal_list
else:
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess_val')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results_val')
make_dir(results_dir)
json_data = open('../../databases/YouTubeVOS/val/meta.json')
data = json.load(json_data)
else:
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess_val_davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results_val_davis')
make_dir(results_dir)
for batch_idx, (inputs, seq_name,
starting_frame) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
for ii in range(max_ii):
# x: input images (N consecutive frames from M different sequences)
x = batch_to_var_test(args, inputs[ii])
print(seq_name[0] + '/' + '%05d' %
(starting_frame[0] + ii))
if ii == 0:
if args.dataset == 'youtube':
num_instances = len(
data['videos'][seq_name[0]]['objects'])
else:
annotation = Image.open(
'../../databases/DAVIS2017/Annotations/480p/' +
seq_name[0] + '/00000.png')
instance_ids = sorted(np.unique(annotation))
instance_ids = instance_ids if instance_ids[
0] else instance_ids[1:]
if len(instance_ids) > 0:
instance_ids = instance_ids[:-1] if instance_ids[
-1] == 255 else instance_ids
num_instances = len(instance_ids)
#from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test(
args, self.encoder, self.decoder, x,
prev_hidden_temporal_list)
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
for t in range(10):
mask_pred = (torch.squeeze(outs[0,
t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
toimage(mask2assess, cmin=0,
cmax=255).save(base_dir_masks_sep +
'%05d_instance_%02d.png' %
(starting_frame[0] + ii, t))
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure()
plt.axis('off')
plt.figure()
plt.axis('off')
plt.imshow(frame_img)
for t in range(num_instances):
mask_pred = (torch.squeeze(
outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones(
(mask_pred.shape[0], mask_pred.shape[1], 3))
color_mask = np.array(colors[t]) / 255.0
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
figname = base_dir + 'frame_%02d.png' % (
starting_frame[0] + ii)
plt.savefig(figname, bbox_inches='tight')
plt.close()
if self.video_mode:
prev_hidden_temporal_list = hidden_temporal_list
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.display = args.display
self.dataset = args.dataset
self.all_classes = args.all_classes
self.T = args.maxseqlen
self.batch_size = args.batch_size
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
dataset = get_dataset(args,
self.split,
image_transforms,
augment=False)
self.loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.sample_list = dataset.get_sample_list()
self.args = args
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(
args.model_name)
self.args.use_feedback = load_args.use_feedback
self.args.base_model = load_args.base_model
self.hidden_size = load_args.hidden_size
self.args.nconvlstm = load_args.nconvlstm
self.encoder = FeatureExtractor(load_args)
self.decoder = RSIS(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
# check if the model was trained using multiple gpus
trained_parallel = False
for k, v in encoder_dict.items():
if k[:7] == "module.":
trained_parallel = True
break
if trained_parallel and not args.ngpus > 1:
# create new OrderedDict that does not contain "module."
new_encoder_state_dict = OrderedDict()
new_decoder_state_dict = OrderedDict()
for k, v in encoder_dict.items():
name = k[7:] # remove "module."
new_encoder_state_dict[name] = v
for k, v in decoder_dict.items():
name = k[7:] # remove "module."
new_decoder_state_dict[name] = v
encoder_dict = new_encoder_state_dict
decoder_dict = new_decoder_state_dict
self.encoder.load_state_dict(encoder_dict)
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
def create_figures(self):
acc_samples = 0
results_dir = os.path.join('../models', args.model_name,
args.model_name + '_results')
make_dir(results_dir)
masks_dir = os.path.join(args.model_name + '_masks')
abs_masks_dir = os.path.join(results_dir, masks_dir)
make_dir(abs_masks_dir)
print "Creating annotations for cityscapes validation..."
for batch_idx, (inputs, targets) in enumerate(self.loader):
x, y_mask, y_class, sw_mask, sw_class = batch_to_var(
self.args, inputs, targets)
out_masks, out_scores, stop_probs = test(self.args, self.encoder,
self.decoder, x)
class_ids = [24, 25, 26, 27, 28, 31, 32, 33]
for sample in range(self.batch_size):
sample_idx = self.sample_list[sample + acc_samples]
image_dir = os.path.join(sample_idx.split('.')[0] + '.png')
im = scipy.misc.imread(image_dir)
h = im.shape[0]
w = im.shape[1]
sample_idx = sample_idx.split('/')[-1].split('.')[0]
results_file = open(
os.path.join(results_dir, sample_idx + '.txt'), 'w')
img_masks = out_masks[sample]
instance_id = 0
class_scores = out_scores[sample]
stop_scores = stop_probs[sample]
for time_step in range(self.T):
mask = img_masks[time_step].cpu().numpy()
mask = (mask > args.mask_th)
h_mask = mask.shape[0]
w_mask = mask.shape[1]
mask = (mask > 0)
labeled_blobs = measure.label(mask, background=0).flatten()
# find the biggest one
count = Counter(labeled_blobs)
s = []
max_num = 0
for v, k in count.iteritems():
if v == 0:
continue
if k > max_num:
max_num = k
max_label = v
# build mask from the largest connected component
segmentation = (labeled_blobs == max_label).astype("uint8")
mask = segmentation.reshape([h_mask, w_mask]) * 255
mask = scipy.misc.imresize(mask, [h, w])
class_scores_mask = class_scores[time_step].cpu().numpy()
stop_scores_mask = stop_scores[time_step].cpu().numpy()
class_score = np.argmax(class_scores_mask)
for i in range(len(class_scores_mask) - 1):
name_instance = sample_idx + '_' + str(
instance_id) + '.png'
pred_class_score = class_scores_mask[i + 1]
objectness = stop_scores_mask[0]
pred_class_score *= objectness
scipy.misc.imsave(
os.path.join(abs_masks_dir, name_instance), mask)
results_file.write(masks_dir + '/' + name_instance +
' ' + str(class_ids[i]) + ' ' +
str(pred_class_score) + '\n')
instance_id += 1
results_file.close()
acc_samples += self.batch_size
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.dataset = args.dataset
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
if args.dataset == 'davis2017':
dataset = get_dataset(args,
split=self.split,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment
and self.split == 'train',
inputRes=(240, 427),
video_mode=True,
use_prev_mask=True)
else: #args.dataset == 'youtube'
dataset = get_dataset(args,
split=self.split,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment
and self.split == 'train',
inputRes=(256, 448),
video_mode=True,
use_prev_mask=True)
self.loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.args = args
print(args.model_name)
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(
args.model_name, args.use_gpu)
load_args.use_gpu = args.use_gpu
self.encoder = FeatureExtractor(load_args)
self.decoder = RSISMask(load_args)
print(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
encoder_dict, decoder_dict = check_parallel(encoder_dict, decoder_dict)
self.encoder.load_state_dict(encoder_dict)
to_be_deleted_dec = []
for k in decoder_dict.keys():
if 'fc_stop' in k:
to_be_deleted_dec.append(k)
for k in to_be_deleted_dec:
del decoder_dict[k]
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
if load_args.length_clip == 1:
self.video_mode = False
print('video mode not activated')
else:
self.video_mode = True
print('video mode activated')
def run_eval(self):
print("Dataset is %s" % (self.dataset))
print("Split is %s" % (self.split))
if args.overlay_masks:
colors = []
palette = sequence_palette()
inv_palette = {}
for k, v in palette.items():
inv_palette[v] = k
num_colors = len(inv_palette.keys())
for id_color in range(num_colors):
if id_color == 0 or id_color == 21:
continue
c = inv_palette[id_color]
colors.append(c)
if self.split == 'val':
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results')
make_dir(results_dir)
json_data = open(
'../../databases/YouTubeVOS/train/train-val-meta.json')
data = json.load(json_data)
else: #args.dataset == 'davis2017'
import lmdb
from misc.config import cfg
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess-davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results-davis')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
for batch_idx, (inputs, targets, seq_name,
starting_frame) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
if args.dataset == 'davis2017':
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(
key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
frame_names = sorted(_files)
for ii in range(max_ii):
#start_time = time.time()
# x: input images (N consecutive frames from M different sequences)
# y_mask: ground truth annotations (some of them are zeros to have a fixed length in number of object instances)
# sw_mask: this mask indicates which masks from y_mask are valid
x, y_mask, sw_mask = batch_to_var(args, inputs[ii],
targets[ii])
if ii == 0:
prev_mask = y_mask
#from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test_prev_mask(
args, self.encoder, self.decoder, x,
prev_hidden_temporal_list, prev_mask)
#end_inference_time = time.time()
#print("inference time: %.3f" %(end_inference_time-start_time))
if args.dataset == 'youtube':
num_instances = len(
data['videos'][seq_name[0]]['objects'])
else:
num_instances = int(
torch.sum(sw_mask.data).data.cpu().numpy())
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
for t in range(num_instances):
mask_pred = (torch.squeeze(outs[0,
t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
if args.dataset == 'youtube':
toimage(mask2assess, cmin=0,
cmax=255).save(base_dir_masks_sep +
'%05d_instance_%02d.png' %
(starting_frame[0] + ii, t))
else:
toimage(mask2assess, cmin=0,
cmax=255).save(base_dir_masks_sep +
frame_names[ii] +
'_instance_%02d.png' % (t))
#end_saving_masks_time = time.time()
#print("inference + saving masks time: %.3f" %(end_saving_masks_time - start_time))
if args.dataset == 'youtube':
print(seq_name[0] + '/' + '%05d' %
(starting_frame[0] + ii))
else:
print(seq_name[0] + '/' + frame_names[ii])
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure()
plt.axis('off')
plt.figure()
plt.axis('off')
plt.imshow(frame_img)
for t in range(num_instances):
mask_pred = (torch.squeeze(
outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones(
(mask_pred.shape[0], mask_pred.shape[1], 3))
color_mask = np.array(colors[t]) / 255.0
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
if args.dataset == 'youtube':
figname = base_dir + 'frame_%02d.png' % (
starting_frame[0] + ii)
else:
figname = base_dir + frame_names[ii] + '.png'
plt.savefig(figname, bbox_inches='tight')
plt.close()
if self.video_mode:
if args.only_spatial == False:
prev_hidden_temporal_list = hidden_temporal_list
if ii > 0:
prev_mask = outs
else:
prev_mask = y_mask
del outs, hidden_temporal_list, x, y_mask, sw_mask
else:
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess_val')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results_val')
make_dir(results_dir)
json_data = open('../../databases/YouTubeVOS/valid/meta.json')
data = json.load(json_data)
else: #args.dataset == 'davis2017'
import lmdb
from misc.config import cfg
masks_sep_dir = os.path.join('../models', args.model_name,
'masks_sep_2assess_val_davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name,
'results_val_davis')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
for batch_idx, (inputs, seq_name,
starting_frame) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
if args.dataset == 'youtube':
seq_data = data['videos'][seq_name[0]]['objects']
frame_names = []
frame_names_with_new_objects = []
instance_ids = []
for obj_id in seq_data.keys():
instance_ids.append(int(obj_id))
frame_names_with_new_objects.append(
seq_data[obj_id]['frames'][0])
for frame_name in seq_data[obj_id]['frames']:
if frame_name not in frame_names:
frame_names.append(frame_name)
frame_names.sort()
frame_names_with_new_objects_idxs = []
for kk in range(len(frame_names_with_new_objects)):
new_frame_idx = frame_names.index(
frame_names_with_new_objects[kk])
frame_names_with_new_objects_idxs.append(new_frame_idx)
else: #davis2017
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(
key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
frame_names = sorted(_files)
for ii in range(max_ii):
# x: input images (N consecutive frames from M different sequences)
# y_mask: ground truth annotations (some of them are zeros to have a fixed length in number of object instances)
# sw_mask: this mask indicates which masks from y_mask are valid
x = batch_to_var_test(args, inputs[ii])
print(seq_name[0] + '/' + frame_names[ii])
if ii == 0:
frame_name = frame_names[0]
if args.dataset == 'youtube':
annotation = Image.open(
'../../databases/YouTubeVOS/valid/Annotations/'
+ seq_name[0] + '/' + frame_name + '.png')
annot = imresize(annotation, (256, 448),
interp='nearest')
else: #davis2017
annotation = Image.open(
'../../databases/DAVIS2017/Annotations/480p/' +
seq_name[0] + '/' + frame_name + '.png')
instance_ids = sorted(np.unique(annotation))
instance_ids = instance_ids if instance_ids[
0] else instance_ids[1:]
if len(instance_ids) > 0:
instance_ids = instance_ids[:-1] if instance_ids[
-1] == 255 else instance_ids
annot = imresize(annotation, (240, 427),
interp='nearest')
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
annot = annot.numpy().squeeze()
annot = annot_from_mask(annot, instance_ids)
prev_mask = annot
prev_mask = np.expand_dims(prev_mask, axis=0)
prev_mask = torch.from_numpy(prev_mask)
y_mask = Variable(prev_mask.float(),
requires_grad=False)
prev_mask = y_mask.cuda()
del annot
if args.dataset == 'youtube':
if ii > 0 and ii in frame_names_with_new_objects_idxs:
frame_name = frame_names[ii]
annotation = Image.open(
'../../databases/YouTubeVOS/valid/Annotations/'
+ seq_name[0] + '/' + frame_name + '.png')
annot = imresize(annotation, (256, 448),
interp='nearest')
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
annot = annot.numpy().squeeze()
new_instance_ids = np.unique(annot)[1:]
annot = annot_from_mask(annot, new_instance_ids)
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = Variable(annot.float(),
requires_grad=False)
annot = annot.cuda()
for kk in new_instance_ids:
prev_mask[:, int(kk - 1), :] = annot[:,
int(kk -
1), :]
del annot
#from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test_prev_mask(
args, self.encoder, self.decoder, x,
prev_hidden_temporal_list, prev_mask)
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
for t in range(len(instance_ids)):
mask_pred = (torch.squeeze(outs[0,
t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
toimage(mask2assess, cmin=0,
cmax=255).save(base_dir_masks_sep +
frame_names[ii] +
'_instance_%02d.png' % (t))
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure()
plt.axis('off')
plt.figure()
plt.axis('off')
plt.imshow(frame_img)
for t in range(len(instance_ids)):
mask_pred = (torch.squeeze(
outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones(
(mask_pred.shape[0], mask_pred.shape[1], 3))
color_mask = np.array(colors[t]) / 255.0
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
figname = base_dir + frame_names[ii] + '.png'
plt.savefig(figname, bbox_inches='tight')
plt.close()
if self.video_mode:
if args.only_spatial == False:
prev_hidden_temporal_list = hidden_temporal_list
if ii > 0:
prev_mask = outs
del x, hidden_temporal_list, outs
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.display = args.display
self.no_display_text = args.no_display_text
self.dataset = args.dataset
self.all_classes = args.all_classes
self.use_cats = args.use_cats
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
dataset = get_dataset(args,
self.split,
image_transforms,
augment=False,
imsize=args.imsize)
self.sample_list = dataset.get_sample_list()
self.class_names = dataset.get_classes()
if args.dataset == 'pascal':
self.gt_file = pickle.load(
open(
os.path.join(args.pascal_dir,
'VOCGT_%s.pkl' % (self.split)), 'rb'))
self.key_to_anns = dict()
self.ignoremasks = {}
for ann in self.gt_file:
if ann['ignore'] == 1:
if type(ann['segmentation']['counts']) == list:
im_height = ann['segmentation']['size'][0]
im_width = ann['segmentation']['size'][1]
rle = mask.frPyObjects([ann['segmentation']],
im_height, im_width)
else:
rle = [ann['segmentation']]
m = mask.decode(rle)
self.ignoremasks[ann['image_id']] = m
if ann['image_id'] in self.key_to_anns.keys():
self.key_to_anns[ann['image_id']].append(ann)
else:
self.key_to_anns[ann['image_id']] = [ann]
self.coco = create_coco_object(args, self.sample_list,
self.class_names)
self.loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.args = args
self.colors = []
palette = sequence_palette()
inv_palette = {}
for k, v in palette.iteritems():
inv_palette[v] = k
num_colors = len(inv_palette.keys())
for i in range(num_colors):
if i == 0 or i == 21:
continue
c = inv_palette[i]
self.colors.append(c)
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(
args.model_name, args.use_gpu)
load_args.use_gpu = args.use_gpu
self.encoder = FeatureExtractor(load_args)
self.decoder = RSIS(load_args)
print(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
encoder_dict, decoder_dict = check_parallel(encoder_dict, decoder_dict)
self.encoder.load_state_dict(encoder_dict)
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
def _create_json(self):
predictions = list()
acc_samples = 0
print "Creating annotations..."
for batch_idx, (inputs, targets) in enumerate(self.loader):
x, y_mask, y_class, sw_mask, sw_class = batch_to_var(
self.args, inputs, targets)
num_objects = np.sum(sw_mask.data.float().cpu().numpy(), axis=-1)
out_masks, out_scores, stop_probs = test(self.args, self.encoder,
self.decoder, x)
out_scores = out_scores.cpu().numpy()
stop_scores = stop_probs.cpu().numpy()
out_masks = out_masks.cpu().numpy()
out_classes = np.argmax(out_scores, axis=-1)
w = x.size()[-1]
h = x.size()[-2]
#out_masks, out_classes, y_mask, y_class = outs_perms_to_cpu(self.args,outs,true_perms,h,w)
for s in range(out_masks.shape[0]):
this_pred = list()
sample_idx = self.sample_list[s + acc_samples]
if self.args.dataset == 'pascal':
ignore_mask = self.ignoremasks[sample_idx]
else:
ignore_mask = None
if self.dataset == 'pascal':
image_dir = os.path.join(args.pascal_dir, 'JPEGImages',
sample_idx + '.jpg')
elif self.dataset == 'cityscapes':
sample_idx = sample_idx.split('.')[0]
image_dir = sample_idx + '.png'
elif self.dataset == 'leaves':
image_dir = sample_idx
im = imread(image_dir)
h = im.shape[0]
w = im.shape[1]
objectness_scores = []
class_scores = []
reached_end = False
for i in range(out_masks.shape[1]):
if reached_end:
break
objectness = stop_scores[s][i][0]
if objectness < args.stop_th:
continue
pred_mask = out_masks[s][i]
# store class with max confidence for display
if args.class_th == 0.0:
max_class = 1
else:
max_class = out_classes[s][i]
# process mask to create annotation
pred_mask, is_valid, raw_pred_mask = resize_mask(
args, pred_mask, h, w, ignore_mask)
# for evaluation we repeat the mask with all its class probs
for cls_id in range(len(self.class_names)):
if cls_id == 0:
# ignore eos
continue
pred_class_score = out_scores[s][i][cls_id]
pred_class_score_mod = pred_class_score * objectness
ann = create_annotation(self.args, sample_idx,
pred_mask, cls_id,
pred_class_score_mod,
self.class_names, is_valid)
if ann is not None:
if self.dataset == 'leaves':
if objectness > args.stop_th:
this_pred.append(ann)
else:
# for display we only take the mask with max confidence
if cls_id == max_class and pred_class_score_mod >= self.args.class_th:
ann_save = create_annotation(
self.args, sample_idx, raw_pred_mask,
cls_id, pred_class_score_mod,
self.class_names, is_valid)
this_pred.append(ann_save)
predictions.append(ann)
if self.display:
figures_dir = os.path.join(
'../models', args.model_name,
args.model_name + '_figs_' + args.eval_split)
make_dir(figures_dir)
plt.figure()
plt.axis('off')
plt.figure()
plt.axis('off')
plt.imshow(im)
display_masks(this_pred,
self.colors,
im_height=im.shape[0],
im_width=im.shape[1],
no_display_text=self.args.no_display_text)
if self.dataset == 'cityscapes':
sample_idx = sample_idx.split('/')[-1]
if self.dataset == 'leaves':
sample_idx = sample_idx.split('/')[-1]
figname = os.path.join(figures_dir, sample_idx)
plt.savefig(figname, bbox_inches='tight')
plt.close()
acc_samples += np.shape(out_masks)[0]
return predictions
def run_eval(self):
print "Dataset is %s" % (self.dataset)
print "Split is %s" % (self.split)
print "Evaluating for %d images" % (len(self.sample_list))
print "Number of classes is %d" % (len(self.class_names))
if self.dataset == 'pascal':
cocoGT = self.coco.loadRes(self.gt_file)
predictions = self._create_json()
if not args.no_run_coco_eval:
cocoP = self.coco.loadRes(predictions)
cocoEval = COCOeval(cocoGT, cocoP, 'segm')
cocoEval.params.maxDets = [1, args.max_dets, 100]
cocoEval.params.useCats = args.use_cats
if not args.cat_id == -1:
cocoEval.params.catIds = [args.cat_id]
cocoEval.params.imgIds = sorted(self.sample_list)
cocoEval.params.catIds = range(1, len(self.class_names))
print("Results for all the classes together")
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
if self.all_classes:
for actual_class in cocoEval.params.catIds:
print("Testing class dataset_id: " + str(actual_class))
print("Which corresponds to name: " +
self.class_names[actual_class])
cocoEval.params.catIds = [actual_class]
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.dataset = args.dataset
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
if args.dataset == 'davis2017':
dataset = get_dataset(args,
split=self.split,
e=0,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment and self.split == 'train',
inputRes=(240, 427),
video_mode=True,
use_prev_mask=True,
eval=True)
else: # args.dataset == 'youtube' or kittimots
dataset = get_dataset(args,
split=self.split,
e=0,
image_transforms=image_transforms,
target_transforms=None,
augment=args.augment and self.split == 'train',
#inputRes=(256, 448),
inputRes=(287, 950),
#inputRes=(412,723),
#inputRes=(178,590),
video_mode=True,
use_prev_mask=True,
eval=True)
self.loader = data.DataLoader(dataset, batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.args = args
print(args.model_name)
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(args.model_name, args.use_gpu)
load_args.use_gpu = args.use_gpu
self.encoder = FeatureExtractor(load_args)
self.decoder = RSISMask(load_args)
print(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder, device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder, device_ids=range(args.ngpus))
encoder_dict, decoder_dict = check_parallel(encoder_dict, decoder_dict)
self.encoder.load_state_dict(encoder_dict)
to_be_deleted_dec = []
for k in decoder_dict.keys():
if 'fc_stop' in k:
to_be_deleted_dec.append(k)
for k in to_be_deleted_dec:
del decoder_dict[k]
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
if load_args.length_clip == 1:
self.video_mode = False
print('video mode not activated')
else:
self.video_mode = True
print('video mode activated')
def run_eval(self):
print("Dataset is %s" % (self.dataset))
print("Split is %s" % (self.split))
if args.overlay_masks:
colors = []
palette = sequence_palette()
inv_palette = {}
for k, v in palette.items():
inv_palette[v] = k
num_colors = len(inv_palette.keys())
for id_color in range(num_colors):
if id_color == 0 or id_color == 21:
continue
c = inv_palette[id_color]
colors.append(c)
if self.split == 'val-inference':
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results')
make_dir(results_dir)
json_data = open('../../databases/YouTubeVOS/train/train-val-meta.json')
data = json.load(json_data)
elif args.dataset == 'davis2017':
import lmdb
from misc.config import cfg
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess-davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results-davis')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
else:
import lmdb
from misc.config_kittimots import cfg
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess-kitti')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results-kitti')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
for batch_idx, (inputs, targets, seq_name, starting_frame, frames_with_new_ids) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
frames_with_new_ids = np.array(frames_with_new_ids)
#print('Variable max_ii')
#print(max_ii)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
if args.dataset == 'davis2017':
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
frame_names = sorted(_files)
if args.dataset == 'kittimots':
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
frame_names = sorted(_files) # llistat de frames d'una seqüència de video
dict_outs = {}
#print("NEW OBJECTS FRAMES", frames_with_new_ids)
# make a dir of results for each instance
'''for t in range(args.maxseqlen):
base_dir_2 = results_dir + '/' + seq_name[0] + '/' + str(t)
make_dir(base_dir_2)'''
for ii in range(max_ii): # iteració sobre els frames/clips amb dimensio lenght_clip
# start_time = time.time()
# x: input images (N consecutive frames from M different sequences)
# y_mask: ground truth annotations (some of them are zeros to have a fixed length in number of object instances)
# sw_mask: this mask indicates which masks from y_mask are valid
x, y_mask, sw_mask = batch_to_var(args, inputs[ii], targets[ii])
#print(seq_name[0] + '/' + frame_names[ii])
if ii == 0:
#one-shot approach, information about the first frame of the sequende
prev_mask = y_mask
#list of the first instances that appear on the sequence and update of the dictionary
annotation = Image.open(
'../../databases/KITTIMOTS/Annotations/' + seq_name[0] + '/' + frame_names[
ii] + '.png').convert('P')
annot = np.expand_dims(annotation, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
instance_ids = np.unique(annot)
for i in instance_ids[1:]:
dict_outs.update({str(int(i-1)):int(i)})
#instances = len(instance_ids)-1
#one-shot approach, add GT information when a new instance appears on the video sequence
if args.dataset == 'kittimots':
if ii > 0 and ii in frames_with_new_ids:
frame_name = frame_names[ii]
annotation = Image.open(
'../../databases/KITTIMOTS/Annotations/' + seq_name[
0] + '/' + frame_name + '.png').convert('P')
#annot = imresize(annotation, (256, 448), interp='nearest')
annot = imresize(annotation, (287,950), interp='nearest')
#annot = imresize(annotation, (412, 723), interp='nearest')
#annot = imresize(annotation, (178, 590), interp='nearest')
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
annot = annot.numpy().squeeze()
new_instance_ids = np.setdiff1d(np.unique(annot), instance_ids)
annot = annot_from_mask(annot, new_instance_ids)
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = Variable(annot.float(), requires_grad=False)
annot = annot.cuda()
#adding only the information of the new instance after the last active branch
for kk in new_instance_ids:
if dict_outs:
last = int(list(dict_outs.keys())[-1])
else:
#if the dictionary is empty
last = -1
prev_mask[:, int(last+1), :] = annot[:, int(kk - 1), :]
dict_outs.update({str(last+1):int(kk)})
del annot
#update the list of instances that have appeared on the video sequence
if len(new_instance_ids) > 0:
instance_ids = np.append(instance_ids, new_instance_ids)
#instances = instances + len(new_instance_ids)
# from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test_prev_mask(args, self.encoder, self.decoder, x,
prev_hidden_temporal_list, prev_mask)
# end_inference_time = time.time()
# print("inference time: %.3f" %(end_inference_time-start_time))
if args.dataset == 'youtube':
num_instances = len(data['videos'][seq_name[0]]['objects'])
else:
num_instances = int(torch.sum(sw_mask.data).data.cpu().numpy())
num_instances = args.maxseqlen
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
'''out_stop = outs[1]
outs = outs[0]
for m in range(len(out_stop[0])):
print(m)
print(out_stop[0][m])'''
#print("OUT STOP: ", out_stop[0])
outs_masks = np.zeros((args.maxseqlen,), dtype=int)
for t in range(num_instances):
mask_pred = (torch.squeeze(outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
'''indxs_instance = np.where((0.6 > mask_pred) & (mask_pred>= 0.5))
indxs_instance_1 = np.where((0.7 > mask_pred) & (mask_pred>= 0.6))
indxs_instance_2 = np.where((0.8 > mask_pred) & (mask_pred>= 0.7))
indxs_instance_3 = np.where((0.9 > mask_pred) & (mask_pred>= 0.8))
indxs_instance_4 = np.where((0.999> mask_pred) & (mask_pred>= 0.9))
indxs_instance_5 = np.where((0.9999 > mask_pred) & (mask_pred>= 0.999))'''
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
''' mask2assess[indxs_instance] = 40
mask2assess[indxs_instance_1] = 80
mask2assess[indxs_instance_2] = 120
mask2assess[indxs_instance_3] = 160
mask2assess[indxs_instance_4] = 200
mask2assess[indxs_instance_5] = 255'''
if str(t) in dict_outs:
i = dict_outs[str(t)]
else:
break
if args.dataset == 'youtube':
toimage(mask2assess, cmin=0, cmax=255).save(
base_dir_masks_sep + '%05d_instance_%02d.png' % (starting_frame[0] + ii, i))
else:
toimage(mask2assess, cmin=0, cmax=255).save(
base_dir_masks_sep + frame_names[ii] + '_instance_%02d.png' % (i))
#create vector of predictions, gives information about which branches are active
if len(indxs_instance[0]) != 0:
outs_masks[t] = 1
else:
outs_masks[t] = 0
outs = outs.cpu().numpy()
#print("INS: ", outs_masks)
#print(json.dumps(dict_outs))
#delete spurious branches
last_position = last_ocurrence(outs_masks, 1) + 1
while len(dict_outs) < last_position:
for n in range(args.maxseqlen):
if outs_masks[n] == 1 and str(n) not in dict_outs:
outs = np.delete(outs, n, axis=1)
outs_masks = np.delete(outs_masks, n)
del hidden_temporal_list[n]
z = np.zeros((height * width))
outs = np.insert(outs, args.maxseqlen - 1, z, axis=1)
hidden_temporal_list.append(None)
outs_masks = np.append(outs_masks, 0)
last_position = last_ocurrence(outs_masks, 1) + 1
instances = sum(outs_masks) # number of active branches
#delete branches of instances that disappear and rearrange
for n in range(args.maxseqlen):
while outs_masks[n] == 0 and n < instances:
outs = np.delete(outs, n, axis=1 )
outs_masks = np.delete(outs_masks, n)
del hidden_temporal_list[n]
z = np.zeros((height * width))
outs = np.insert(outs, args.maxseqlen-1, z, axis=1)
hidden_temporal_list.append(None)
outs_masks = np.append(outs_masks, 0)
#update dictionary by shifting entries
for m in range(len(dict_outs)-(n+1)):
value = dict_outs[str(m + n + 1)]
dict_outs.update({str(n + m): value})
#print(json.dumps(dict_outs))
last = int(list(dict_outs.keys())[-1])
del dict_outs[str(last)]
#an instance has disappeared, update dictionary
while len(dict_outs) > sum(outs_masks):
last = int(list(dict_outs.keys())[-1])
del dict_outs[str(last)]
outs = torch.from_numpy(outs)
outs = outs.cuda()
#print("OUTS: ", outs_masks)
#print(json.dumps(dict_outs))
# end_saving_masks_time = time.time()
# print("inference + saving masks time: %.3f" %(end_saving_masks_time - start_time))
if args.dataset == 'youtube':
print(seq_name[0] + '/' + '%05d' % (starting_frame[0] + ii))
else:
print(seq_name[0] + '/' + frame_names[ii])
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure();
plt.axis('off')
plt.figure();
plt.axis('off')
plt.imshow(frame_img)
#print("INSTANCES: ", instances)
for t in range(instances):
mask_pred = (torch.squeeze(outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones((mask_pred.shape[0], mask_pred.shape[1], 3))
if str(t) in dict_outs:
color_mask = np.array(colors[dict_outs[str(t)]]) / 255.0
else:
#color_mask = np.array(colors[0]) / 255.0
break
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
if args.dataset == 'youtube':
figname = base_dir + 'frame_%02d.png' % (starting_frame[0] + ii)
else:
figname = base_dir + frame_names[ii] + '.png'
plt.savefig(figname, bbox_inches='tight')
plt.close()
#Print a video for each instance
'''for t in range(instances):
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure();
plt.axis('off')
plt.figure();
plt.axis('off')
plt.imshow(frame_img)
mask_pred = (torch.squeeze(outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones((mask_pred.shape[0], mask_pred.shape[1], 3))
if str(t) in dict_outs:
color_mask = np.array(colors[dict_outs[str(t)]]) / 255.0
else:
#color_mask = np.array(colors[0]) / 255.0
break
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
figname = base_dir + '/' + str(dict_outs[str(t)]) + '/' + frame_names[ii] + '.png'
plt.savefig(figname, bbox_inches='tight')
plt.close()'''
if self.video_mode:
if args.only_spatial == False:
prev_hidden_temporal_list = hidden_temporal_list
if ii > 0:
prev_mask = outs
else:
prev_mask = y_mask
del outs, hidden_temporal_list, x, y_mask, sw_mask
else:
if args.dataset == 'youtube':
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess_val')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results_val')
make_dir(results_dir)
json_data = open('../../databases/YouTubeVOS/val/meta.json')
data = json.load(json_data)
elif args.dataset == 'davis2017':
import lmdb
from misc.config import cfg
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess_val_davis')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results_val_davis')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
else:
import lmdb
from misc.config_kittimots import cfg
masks_sep_dir = os.path.join('../models', args.model_name, 'masks_sep_2assess-kitti')
make_dir(masks_sep_dir)
if args.overlay_masks:
results_dir = os.path.join('../models', args.model_name, 'results-kitti')
make_dir(results_dir)
lmdb_env_seq_dir = osp.join(cfg.PATH.DATA, 'lmdb_seq')
if osp.isdir(lmdb_env_seq_dir):
lmdb_env_seq = lmdb.open(lmdb_env_seq_dir)
else:
lmdb_env_seq = None
for batch_idx, (inputs, seq_name, starting_frame) in enumerate(self.loader):
prev_hidden_temporal_list = None
max_ii = min(len(inputs), args.length_clip)
if args.overlay_masks:
base_dir = results_dir + '/' + seq_name[0] + '/'
make_dir(base_dir)
if args.dataset == 'youtube':
seq_data = data['videos'][seq_name[0]]['objects']
frame_names = []
frame_names_with_new_objects = []
instance_ids = []
for obj_id in seq_data.keys():
instance_ids.append(int(obj_id))
frame_names_with_new_objects.append(seq_data[obj_id]['frames'][0])
for frame_name in seq_data[obj_id]['frames']:
if frame_name not in frame_names:
frame_names.append(frame_name)
frame_names.sort()
frame_names_with_new_objects_idxs = []
for kk in range(len(frame_names_with_new_objects)):
new_frame_idx = frame_names.index(frame_names_with_new_objects[kk])
frame_names_with_new_objects_idxs.append(new_frame_idx)
elif args.dataset == 'davis2017':
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
frame_names = sorted(_files)
else:
key_db = osp.basename(seq_name[0])
if not lmdb_env_seq == None:
with lmdb_env_seq.begin() as txn:
_files_vec = txn.get(key_db.encode()).decode().split('|')
_files = [osp.splitext(f)[0] for f in _files_vec]
else:
seq_dir = osp.join(cfg['PATH']['SEQUENCES'], key_db)
_files_vec = os.listdir(seq_dir)
_files = [osp.splitext(f)[0] for f in _files_vec]
# frame_names_with_new_objects_idxs = [3,6,9]
frame_names = sorted(_files)
for ii in range(max_ii):
# x: input images (N consecutive frames from M different sequences)
# y_mask: ground truth annotations (some of them are zeros to have a fixed length in number of object instances)
# sw_mask: this mask indicates which masks from y_mask are valid
x = batch_to_var_test(args, inputs[ii])
print(seq_name[0] + '/' + frame_names[ii])
if ii == 0:
frame_name = frame_names[0]
if args.dataset == 'youtube':
annotation = Image.open(
'../../databases/YouTubeVOS/val/Annotations/' + seq_name[0] + '/' + frame_name + '.png')
annot = imresize(annotation, (256, 448), interp='nearest')
elif args.dataset == 'davis2017':
annotation = Image.open(
'../../databases/DAVIS2017/Annotations/480p/' + seq_name[0] + '/' + frame_name + '.png')
instance_ids = sorted(np.unique(annotation))
instance_ids = instance_ids if instance_ids[0] else instance_ids[1:]
if len(instance_ids) > 0:
instance_ids = instance_ids[:-1] if instance_ids[-1] == 255 else instance_ids
annot = imresize(annotation, (240, 427), interp='nearest')
else: # kittimots
annotation = Image.open(
'../../databases/KITTIMOTS/Annotations/' + seq_name[0] + '/' + frame_name + '.png')
instance_ids = sorted(np.unique(annotation))
instance_ids = instance_ids if instance_ids[0] else instance_ids[1:]
print("IDS instances: ", instance_ids)
if len(instance_ids) > 0:
instance_ids = instance_ids[:-1] if instance_ids[-1] == 255 else instance_ids
annot = imresize(annotation, (256, 448), interp='nearest')
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
annot = annot.numpy().squeeze()
annot = annot_from_mask(annot, instance_ids)
prev_mask = annot
prev_mask = np.expand_dims(prev_mask, axis=0)
prev_mask = torch.from_numpy(prev_mask)
y_mask = Variable(prev_mask.float(), requires_grad=False)
prev_mask = y_mask.cuda()
del annot
if args.dataset == 'youtube':
if ii > 0 and ii in frame_names_with_new_objects_idxs:
frame_name = frame_names[ii]
annotation = Image.open(
'../../databases/YouTubeVOS/val/Annotations/' + seq_name[0] + '/' + frame_name + '.png')
annot = imresize(annotation, (256, 448), interp='nearest')
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = annot.float()
annot = annot.numpy().squeeze()
new_instance_ids = np.unique(annot)[1:]
annot = annot_from_mask(annot, new_instance_ids)
annot = np.expand_dims(annot, axis=0)
annot = torch.from_numpy(annot)
annot = Variable(annot.float(), requires_grad=False)
annot = annot.cuda()
for kk in new_instance_ids:
prev_mask[:, int(kk - 1), :] = annot[:, int(kk - 1), :]
del annot
# from one frame to the following frame the prev_hidden_temporal_list is updated.
outs, hidden_temporal_list = test_prev_mask(args, self.encoder, self.decoder, x,
prev_hidden_temporal_list, prev_mask)
base_dir_masks_sep = masks_sep_dir + '/' + seq_name[0] + '/'
make_dir(base_dir_masks_sep)
x_tmp = x.data.cpu().numpy()
height = x_tmp.shape[-2]
width = x_tmp.shape[-1]
for t in range(len(instance_ids)):
mask_pred = (torch.squeeze(outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
indxs_instance = np.where(mask_pred > 0.5)
mask2assess = np.zeros((height, width))
mask2assess[indxs_instance] = 255
toimage(mask2assess, cmin=0, cmax=255).save(
base_dir_masks_sep + frame_names[ii] + '_instance_%02d.png' % (t))
if args.overlay_masks:
frame_img = x.data.cpu().numpy()[0, :, :, :].squeeze()
frame_img = np.transpose(frame_img, (1, 2, 0))
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
frame_img = std * frame_img + mean
frame_img = np.clip(frame_img, 0, 1)
plt.figure();
plt.axis('off')
plt.figure();
plt.axis('off')
plt.imshow(frame_img)
for t in range(len(instance_ids)):
mask_pred = (torch.squeeze(outs[0, t, :])).cpu().numpy()
mask_pred = np.reshape(mask_pred, (height, width))
ax = plt.gca()
tmp_img = np.ones((mask_pred.shape[0], mask_pred.shape[1], 3))
color_mask = np.array(colors[t]) / 255.0
for i in range(3):
tmp_img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((tmp_img, mask_pred * 0.7)))
figname = base_dir + frame_names[ii] + '.png'
plt.savefig(figname, bbox_inches='tight')
plt.close()
if self.video_mode:
if args.only_spatial == False:
prev_hidden_temporal_list = hidden_temporal_list
if ii > 0:
prev_mask = outs
del x, hidden_temporal_list, outs
class Evaluate():
def __init__(self, args):
self.split = args.eval_split
self.display = args.display
self.dataset = args.dataset
self.all_classes = args.all_classes
self.T = args.maxseqlen
self.batch_size = args.batch_size
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
image_transforms = transforms.Compose([to_tensor, normalize])
dataset = get_dataset(args,
self.split,
image_transforms,
augment=False,
imsize=args.imsize)
self.loader = data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
drop_last=False)
self.sample_list = dataset.get_sample_list()
self.args = args
encoder_dict, decoder_dict, _, _, load_args = load_checkpoint(
args.model_name)
self.args.use_feedback = load_args.use_feedback
self.args.base_model = load_args.base_model
self.hidden_size = load_args.hidden_size
self.args.nconvlstm = load_args.nconvlstm
self.encoder = FeatureExtractor(load_args)
self.decoder = RSIS(load_args)
if args.ngpus > 1 and args.use_gpu:
self.decoder = torch.nn.DataParallel(self.decoder,
device_ids=range(args.ngpus))
self.encoder = torch.nn.DataParallel(self.encoder,
device_ids=range(args.ngpus))
# check if the model was trained using multiple gpus
trained_parallel = False
for k, v in encoder_dict.items():
if k[:7] == "module.":
trained_parallel = True
break
if trained_parallel and not args.ngpus > 1:
# create new OrderedDict that does not contain "module."
new_encoder_state_dict = OrderedDict()
new_decoder_state_dict = OrderedDict()
for k, v in encoder_dict.items():
name = k[7:] # remove "module."
new_encoder_state_dict[name] = v
for k, v in decoder_dict.items():
name = k[7:] # remove "module."
new_decoder_state_dict[name] = v
encoder_dict = new_encoder_state_dict
decoder_dict = new_decoder_state_dict
self.encoder.load_state_dict(encoder_dict)
self.decoder.load_state_dict(decoder_dict)
if args.use_gpu:
self.encoder.cuda()
self.decoder.cuda()
self.encoder.eval()
self.decoder.eval()
def create_figures(self):
acc_samples = 0
results_root_dir = os.path.join('../models', args.model_name,
args.model_name + '_results')
make_dir(results_root_dir)
results_dir = os.path.join(results_root_dir, 'A1')
make_dir(results_dir)
print "Creating annotations for leaves validation..."
for batch_idx, (inputs, targets) in enumerate(self.loader):
x, y_mask, y_class, sw_mask, sw_class = batch_to_var(
self.args, inputs, targets)
out_masks, _, stop_probs = test(self.args, self.encoder,
self.decoder, x)
for sample in range(self.batch_size):
sample_idx = self.sample_list[sample + acc_samples]
image_dir = os.path.join(sample_idx.split('.')[0] + '.png')
im = scipy.misc.imread(image_dir)
h = im.shape[0]
w = im.shape[1]
mask_sample = np.zeros([h, w])
sample_idx = sample_idx.split('/')[-1].split('.')[0]
img_masks = out_masks[sample]
instance_id = 0
class_scores = stop_probs[sample]
for time_step in range(self.T):
mask = img_masks[time_step].cpu().numpy()
mask = scipy.misc.imresize(mask, [h, w])
class_scores_mask = class_scores[time_step].cpu().numpy()
class_score = class_scores_mask[0]
if class_score > args.class_th:
mask_sample[mask > args.mask_th * 255] = time_step
instance_id += 1
file_name = os.path.join(results_dir, sample_idx + '.png')
file_name_prediction = file_name.replace(
'rgb.png', 'label.png')
im = Image.fromarray(mask_sample).convert('L')
im.save(file_name_prediction)
acc_samples += self.batch_size
