def _get_multilabel_prediction(dir_logits, no_offset_mask=None, topk=8):
h, w, _ = dir_logits.shape
dir_logits = torch.from_numpy(dir_logits).unsqueeze(0).permute(
0, 3, 1, 2)
offsets = []
if topk == dir_logits.shape[1]:
for i in range(topk):
offset_i = DTOffsetHelper.label_to_vector(
torch.tensor([i]).view(1, 1, 1)).repeat(1, 1, h, w)
offset_i = offset_i.float() * dir_logits[:, i:i + 1, :, :]
offsets.append(offset_i)
else:
dir_logits, dir_pred = torch.topk(dir_logits, topk, dim=1)
for i in range(topk):
dir_pred_i = dir_pred[:, i, :, :]
offset_i = DTOffsetHelper.label_to_vector(dir_pred_i)
offset_i = offset_i.float() * dir_logits[:, i:i + 1, :, :]
offsets.append(offset_i)
offset = sum(offsets)
dir_pred = DTOffsetHelper.vector_to_label(offset.permute(0, 2, 3, 1),
num_classes=8,
return_tensor=True)
dir_pred = dir_pred.squeeze(0).numpy()
if no_offset_mask is not None:
dir_pred[no_offset_mask] = 8
return dir_pred
def forward(self, inputs, targets, **kwargs):
from lib.utils.helpers.offset_helper import DTOffsetHelper
pred_mask, pred_direction = inputs
seg_label_map, distance_map, angle_map = targets[0], targets[
1], targets[2]
gt_mask = DTOffsetHelper.distance_to_mask_label(distance_map,
seg_label_map,
return_tensor=True)
gt_size = gt_mask.shape[1:]
mask_weights = self.calc_weights(gt_mask, 2)
pred_direction = F.interpolate(pred_direction,
size=gt_size,
mode="bilinear",
align_corners=True)
pred_mask = F.interpolate(pred_mask,
size=gt_size,
mode="bilinear",
align_corners=True)
mask_loss = F.cross_entropy(pred_mask,
gt_mask,
weight=mask_weights,
ignore_index=-1)
mask_threshold = float(os.environ.get('mask_threshold', 0.5))
binary_pred_mask = torch.softmax(pred_mask,
dim=1)[:, 1, :, :] > mask_threshold
gt_direction = DTOffsetHelper.angle_to_direction_label(
angle_map,
seg_label_map=seg_label_map,
extra_ignore_mask=(binary_pred_mask == 0),
return_tensor=True)
direction_loss_mask = gt_direction != -1
direction_weights = self.calc_weights(
gt_direction[direction_loss_mask], pred_direction.size(1))
direction_loss = F.cross_entropy(pred_direction,
gt_direction,
weight=direction_weights,
ignore_index=-1)
if self.training \
and self.configer.get('iters') % self.configer.get('solver', 'display_iter') == 0 \
and torch.cuda.current_device() == 0:
Log.info('mask loss: {} direction loss: {}.'.format(
mask_loss, direction_loss))
mask_weight = float(os.environ.get('mask_weight', 1))
direction_weight = float(os.environ.get('direction_weight', 1))
return mask_weight * mask_loss + direction_weight * direction_loss
def eval(outputs, meta, running_scores):
distance_map = meta['ori_distance_map']
seg_label_map = meta['ori_target']
gt_mask_label = DTOffsetHelper.distance_to_mask_label(
distance_map, seg_label_map, return_tensor=False)
mask_pred = MaskTask.get_mask_pred(outputs['mask'])
running_scores['mask'].update((mask_pred == 1).astype(np.int)[None],
gt_mask_label[None])
def eval(outputs, meta, running_scores):
distance_map = meta['ori_distance_map']
seg_label_map = meta['ori_target']
dir_label_map = meta['ori_multi_label_direction_map']
dir_label_map = DTOffsetHelper.encode_multi_labels(dir_label_map)
dir_label_map[seg_label_map == -1, :] = -1
gt_mask_label = DTOffsetHelper.distance_to_mask_label(
distance_map, seg_label_map, return_tensor=False)
mask_pred = MaskTask.get_mask_pred(outputs['mask'])
dir_pred = MLDirectionTask._get_multilabel_prediction(
outputs['ml_dir'], no_offset_mask=mask_pred == 0, topk=8)
running_scores['ML dir (mask)'].update(
dir_pred, dir_label_map, (mask_pred == 1) & (seg_label_map != -1))
running_scores['ML dir (GT)'].update(dir_pred, dir_label_map,
gt_mask_label == 1)
def _get_offset(self, mask_logits, dir_logits):
edge_mask = mask_logits[:, 1] > 0.5
dir_logits = torch.softmax(dir_logits, dim=1)
n, _, h, w = dir_logits.shape
keep_mask = edge_mask
dir_label = torch.argmax(dir_logits, dim=1).float()
offset = DTOffsetHelper.label_to_vector(dir_label)
offset = offset.permute(0, 2, 3, 1)
offset[~keep_mask, :] = 0
return offset
def eval(outputs, meta, running_scores):
distance_map = meta['ori_distance_map']
angle_map = meta['ori_angle_map']
seg_label_map = meta['ori_target']
mask_pred = MaskTask.get_mask_pred(outputs['mask'])
dir_pred = np.argmax(outputs['dir'], axis=-1)
gt_mask_label = DTOffsetHelper.distance_to_mask_label(
distance_map, seg_label_map, return_tensor=False)
gt_dir_label = DTOffsetHelper.angle_to_direction_label(
angle_map,
seg_label_map=seg_label_map,
extra_ignore_mask=mask_pred != 1)
running_scores['dir (mask)'].update(dir_pred[None], gt_dir_label[None])
dir_pred[mask_pred != 1] = DTOffsetConfig.num_classes
gt_dir_label = DTOffsetHelper.angle_to_direction_label(
angle_map,
seg_label_map=seg_label_map,
extra_ignore_mask=(gt_mask_label == 0),
)
running_scores['dir (GT)'].update(dir_pred[None], gt_dir_label[None])
def load_boundary(self, fn):
if fn.endswith('mat'):
mat = io.loadmat(fn)
if 'depth' in mat:
dist_map, _ = self._load_maps(fn, None)
boundary_map = DTOffsetHelper.distance_to_mask_label(
dist_map, np.zeros_like(dist_map)).astype(np.float32)
else:
boundary_map = mat['mat'].transpose(1, 2, 0)
else:
boundary_map = ImageHelper.read_image(fn,
tool=self.configer.get(
'data', 'image_tool'),
mode='P')
boundary_map = boundary_map.astype(np.float32) / 255
return boundary_map
def _flip_offset(self, x):
x = self._flip(x, dim=-1)
if len(x.shape) == 4:
return x[:, DTOffsetHelper.flipping_indices()]
else:
return x[DTOffsetHelper.flipping_indices()]