def validate(self):
print 'Validating'
self.model.temperature = 0
# Use temperature 0 while validating
self.model.evaluator.eval()
ious = []
pred_ious = []
loss = []
with torch.no_grad():
for step, data in enumerate(tqdm(self.val_loader)):
output = self.model(data['img'].to(device))
# Get full GT masks
gt_masks = []
for instance in data['instance']:
gt_masks.append(
utils.get_full_mask_from_instance(
self.opts['dataset']['train_val']['min_area'],
instance))
iou = np.zeros((len(gt_masks)), dtype=np.float32)
pred_masks = []
pred_polys = output['pred_polys'].cpu().numpy()
for i in range(pred_polys.shape[0]):
poly = pred_polys[i]
mask, poly = utils.get_full_mask_from_xy(
poly, self.grid_size, data['patch_w'][i],
data['starting_point'][i], instance['img_height'],
instance['img_width'])
pred_masks.append(mask)
for i, gt_mask in enumerate(gt_masks):
iou[i] = metrics.iou_from_mask(pred_masks[i], gt_mask)
l = losses.evaluator_loss(output['ious'],
torch.from_numpy(iou).to(device))
iou = np.mean(iou)
ious.append(iou)
pred_ious.append(float(torch.mean(output['ious'])))
loss.append(l)
del (output)
iou = np.mean(ious)
pred_iou = np.mean(pred_ious)
loss = np.mean(loss)
self.val_writer.add_scalar('loss', float(loss), self.global_step)
print '[VAL] GT IoU: %f, Pred IoU: %f, Loss: %f' % (iou, pred_iou,
loss)
# Reset
self.model.temperature = self.opts['temperature']
self.model.evaluator.train()
def validate(self):
print 'Validating'
self.model.first_v.eval()
# Leave LSTM in train mode, encoder always in eval mode for RL
self.model.temperature = 0
ious = []
with torch.no_grad():
for step, data in enumerate(tqdm(self.val_loader)):
output = self.model(data['img'].to(device))
iou = 0
# Get full GT masks
gt_masks = []
for instance in data['instance']:
gt_masks.append(
utils.get_full_mask_from_instance(
self.opts['dataset']['train_val']['min_area'],
instance))
pred_masks = []
pred_polys = output['pred_polys'].cpu().numpy()
for i in range(pred_polys.shape[0]):
poly = pred_polys[i]
mask, poly = utils.get_full_mask_from_xy(
poly, self.grid_size, data['patch_w'][i],
data['starting_point'][i], instance['img_height'],
instance['img_width'])
pred_masks.append(mask)
for i, gt_mask in enumerate(gt_masks):
iou += metrics.iou_from_mask(pred_masks[i], gt_mask)
iou = iou / pred_polys.shape[0]
ious.append(iou)
del (output)
iou = np.mean(ious)
self.val_writer.add_scalar('iou', float(iou), self.global_step)
print '[VAL] IoU: %f' % (iou)
# Reset
self.model.first_v.train()
def train(self, epoch):
print 'Starting training'
self.model.temperature = self.opts['temperature']
self.model.evaluator.train()
accum = defaultdict(float)
# To accumulate stats for printing
for step, data in enumerate(self.train_loader):
if self.global_step % self.opts['val_freq'] == 0:
self.validate()
self.save_checkpoint(epoch)
# Forward pass (Sampling)
output = self.model(data['img'].to(device))
# Get full GT masks
gt_masks = []
for instance in data['instance']:
gt_masks.append(
utils.get_full_mask_from_instance(
self.opts['dataset']['train_val']['min_area'],
instance))
# Get sampling masks
sampling_masks = []
pred_polys = output['pred_polys'].cpu().numpy()
for i in range(pred_polys.shape[0]):
poly = pred_polys[i]
mask, poly = utils.get_full_mask_from_xy(
poly, self.grid_size, data['patch_w'][i],
data['starting_point'][i], instance['img_height'],
instance['img_width'])
sampling_masks.append(mask)
# Get IoUs
sampling_ious = np.zeros((len(gt_masks)), dtype=np.float32)
for i, gt_mask in enumerate(gt_masks):
sampling_ious[i] = metrics.iou_from_mask(
sampling_masks[i], gt_mask)
# Get Loss
total_loss = losses.evaluator_loss(
output['ious'],
torch.from_numpy(sampling_ious).to(device))
# Backward pass
self.optimizer.zero_grad()
total_loss.backward()
if 'grad_clip' in self.opts.keys():
nn.utils.clip_grad_norm_(self.model.parameters(),
self.opts['grad_clip'])
self.optimizer.step()
accum['loss'] += float(total_loss)
accum['sampling_iou'] += np.mean(sampling_ious)
accum['pred_iou'] += torch.mean(output['ious'])
accum['length'] += 1
if step % self.opts['print_freq'] == 0:
# Mean of accumulated values
for k in accum.keys():
if k == 'length':
continue
accum[k] /= accum['length']
# Add summaries
for k in accum.keys():
if k == 'length':
continue
self.writer.add_scalar(k, accum[k], self.global_step)
print("[%s] Epoch: %d, Step: %d, Loss: %f, Sampling IOU: %f, Pred IOU: %f"\
%(str(datetime.now()), epoch, self.global_step, accum['loss'], accum['sampling_iou'], accum['pred_iou']))
accum = defaultdict(float)
del (output)
self.global_step += 1
def train(self, epoch):
print 'Starting training'
accum = defaultdict(float)
# To accumulate stats for printing
for step, data in enumerate(self.train_loader):
if self.global_step % self.opts['val_freq'] == 0:
self.validate()
self.save_checkpoint(epoch)
# Forward pass (Sampling)
self.model.temperature = self.opts['temperature']
output = self.model(data['img'].to(device))
# Forward pass (Greedy)
self.model.temperature = 0.0
with torch.no_grad():
output_greedy = self.model(data['img'].to(device))
# Get full GT masks
gt_masks = []
for instance in data['instance']:
gt_masks.append(
utils.get_full_mask_from_instance(
self.opts['dataset']['train_val']['min_area'],
instance))
# Get sampling masks
sampling_masks = []
pred_polys = output['pred_polys'].cpu().numpy()
for i in range(pred_polys.shape[0]):
poly = pred_polys[i]
mask, poly = utils.get_full_mask_from_xy(
poly, self.grid_size, data['patch_w'][i],
data['starting_point'][i], instance['img_height'],
instance['img_width'])
sampling_masks.append(mask)
# Get greedy masks
greedy_masks = []
pred_polys = output_greedy['pred_polys'].cpu().numpy()
for i in range(pred_polys.shape[0]):
poly = pred_polys[i]
mask, poly = utils.get_full_mask_from_xy(
poly, self.grid_size, data['patch_w'][i],
data['starting_point'][i], instance['img_height'],
instance['img_width'])
greedy_masks.append(mask)
# Get IoUs
sampling_ious = np.zeros((len(gt_masks)), dtype=np.float32)
greedy_ious = np.zeros((len(gt_masks)), dtype=np.float32)
for i, gt_mask in enumerate(gt_masks):
sampling_ious[i] = metrics.iou_from_mask(
sampling_masks[i], gt_mask)
greedy_ious[i] = metrics.iou_from_mask(greedy_masks[i],
gt_mask)
# Get Loss
total_loss = losses.self_critical_loss(
output['log_probs'], output['lengths'],
torch.from_numpy(sampling_ious).to(device),
torch.from_numpy(greedy_ious).to(device))
# Backward pass
self.optimizer.zero_grad()
total_loss.backward()
if 'grad_clip' in self.opts.keys():
nn.utils.clip_grad_norm_(self.model.parameters(),
self.opts['grad_clip'])
self.optimizer.step()
accum['loss'] += float(total_loss)
accum['sampling_iou'] += np.mean(sampling_ious)
accum['greedy_iou'] += np.mean(greedy_ious)
accum['length'] += 1
if step % self.opts['print_freq'] == 0:
# Mean of accumulated values
for k in accum.keys():
if k == 'length':
continue
accum[k] /= accum['length']
# Add summaries
img = (data['img'].cpu().numpy()[-1, ...] * 255).astype(
np.uint8)
img = np.transpose(img, [1, 2, 0]) # Make [H, W, 3]
vert_logits = np.reshape(
output['vertex_logits'][-1, ...].detach().cpu().numpy(),
(self.grid_size, self.grid_size, 1))
edge_logits = np.reshape(
output['edge_logits'][-1, ...].detach().cpu().numpy(),
(self.grid_size, self.grid_size, 1))
vert_logits = (1 / (1 + np.exp(-vert_logits)) * 255).astype(
np.uint8)
edge_logits = (1 / (1 + np.exp(-edge_logits)) * 255).astype(
np.uint8)
vert_logits = np.tile(vert_logits, [1, 1, 3]) # Make [H, W, 3]
edge_logits = np.tile(edge_logits, [1, 1, 3]) # Make [H, W, 3]
vertex_mask = np.tile(
np.expand_dims(
data['vertex_mask'][-1, ...].cpu().numpy().astype(
np.uint8) * 255, -1), (1, 1, 3))
edge_mask = np.tile(
np.expand_dims(
data['edge_mask'][-1, ...].cpu().numpy().astype(
np.uint8) * 255, -1), (1, 1, 3))
self.writer.add_image('image', img, self.global_step)
self.writer.add_image('vertex_logits', vert_logits,
self.global_step)
self.writer.add_image('edge_logits', edge_logits,
self.global_step)
self.writer.add_image('edge_mask', edge_mask, self.global_step)
self.writer.add_image('vertex_mask', vertex_mask,
self.global_step)
if self.opts['return_attention'] is True:
att = output['attention'][-1, 1:4,
...].detach().cpu().numpy()
att = np.transpose(att, [0, 2, 3, 1]) # Make [T, H, W, 1]
att = np.tile(att, [1, 1, 1, 3]) # Make [T, H, W, 3]
def _scale(att):
att = att / np.max(att)
return (att * 255).astype(np.int32)
self.writer.add_image(
'attention_1',
pyramid_expand(_scale(att[0]), upscale=8, sigma=10),
self.global_step)
self.writer.add_image(
'attention_2',
pyramid_expand(_scale(att[1]), upscale=8, sigma=10),
self.global_step)
self.writer.add_image(
'attention_3',
pyramid_expand(_scale(att[2]), upscale=8, sigma=10),
self.global_step)
for k in accum.keys():
if k == 'length':
continue
self.writer.add_scalar(k, accum[k], self.global_step)
print("[%s] Epoch: %d, Step: %d, Polygon Loss: %f, Sampling IOU: %f, Greedy IOU: %f"\
%(str(datetime.now()), epoch, self.global_step, accum['loss'], accum['sampling_iou'], accum['greedy_iou']))
accum = defaultdict(float)
del (output)
del (output_greedy)
self.global_step += 1
def process_outputs(self, data, output, save=True):
"""
Process outputs to get final outputs for the whole image
Optionally saves the outputs to a folder for evaluation
"""
instances = data['instance']
polys = []
for i, instance in enumerate(instances):
# Postprocess polygon
poly = output['pred_polys'][i]
_, poly = utils.get_full_mask_from_xy(poly, self.grid_size,
data['patch_w'][i],
data['starting_point'][i],
instance['img_height'],
instance['img_width'])
polys.append(poly)
if save:
img_h, img_w = instance['img_height'], instance['img_width']
predicted_poly = []
# Paint pred mask
pred_mask = np.zeros((img_h, img_w), dtype=np.uint8)
poly = poly.astype(np.int)
utils.draw_poly(pred_mask, poly)
predicted_poly.append(poly.tolist())
# Paint GT mask
gt_mask = utils.get_full_mask_from_instance(
self.opts['dataset']['train_val']['min_area'], instance)
instance['my_predicted_poly'] = predicted_poly
instance_id = instance['instance_id']
image_id = instance['image_id']
pred_mask_fname = os.path.join(
self.output_dir, '{}_pred.png'.format(instance_id))
instance['pred_mask_fname'] = os.path.relpath(
pred_mask_fname, self.output_dir)
gt_mask_fname = os.path.join(self.output_dir,
'{}_gt.png'.format(instance_id))
instance['gt_mask_fname'] = os.path.relpath(
gt_mask_fname, self.output_dir)
instance['n_corrections'] = 0
info_fname = os.path.join(self.output_dir,
'{}_info.json'.format(instance_id))
with warnings.catch_warnings():
warnings.simplefilter("ignore")
sio.imsave(pred_mask_fname, pred_mask)
sio.imsave(gt_mask_fname, gt_mask)
with open(info_fname, 'w') as f:
json.dump(instance, f, indent=2)
return polys