def _create_mask(_line, width=8, guard=False):
sz = mask.shape
line_mask = lines.create_grid(sz, [_line], width=width).astype(int)
if guard:
line_mask = 2 * line_mask - lines.create_grid(sz, [_line],
width=2 * width)
return line_mask * mask
def create_mask(self, label_test, weight, width):
_, rot_angle = lines.extract_lines((label_test == 1),
self.angle_range_v)
rot_angle = np.rad2deg(rot_angle)
angle_range = np.array([rot_angle - 2.5, rot_angle + 2.5])
lines_v, _ = lines.extract_lines((label_test == 1),
angle_range,
tresh=0.2)
mask_v = lines.create_grid(label_test.shape, lines_v,
width=width).astype(np.float32) * weight
lines_h, _ = lines.extract_lines((label_test == 1),
angle_range + 90.0,
tresh=0.2)
mask_h = lines.create_grid(label_test.shape, lines_h,
width=width).astype(np.float32) * weight
mask = np.clip(mask_v + mask_h, a_min=0.0, a_max=1.0)
return mask_v, mask_h, mask
def __getitem__(self, index):
data = super(HistDataset, self).__getitem__(index)
label_test = data['label_test']
label = data['label']
weights = data['weights']
angle_range_label = data['angle_range_label']
if angle_range_label == 255:
return data
###############################################################################
lines_v, _rot_angle = lines.extract_lines((label_test == 1),
self.angle_range_v)
lines_h, _ = lines.extract_lines((label_test == 1), self.angle_range_h)
lines_v_mask = lines.create_grid(label.shape, lines_v, width=16)
lines_h_mask = lines.create_grid(label.shape, lines_h, width=16)
""" plt.figure()
plt.imshow(lines_v_mask)
plt.figure()
plt.imshow(lines_h_mask)
plt.show() """
###############################################################################
sz = (len(self.rot_angles), ) + label.shape
bin_label_v = np.zeros(sz, dtype=np.float32)
bin_label_v[angle_range_label] = lines_v_mask.astype(np.float32)
bin_label_h = np.zeros(sz, dtype=np.float32)
bin_label_h[angle_range_label] = lines_h_mask.astype(np.float32)
bin_label = np.stack((bin_label_v, bin_label_h), 0)
weights_v = np.repeat(weights[np.newaxis, ...], len(self.rot_angles),
0)
weights_h = np.repeat(weights[np.newaxis, ...], len(self.rot_angles),
0)
close_idx_0 = angle_range_label - 1
close_idx_1 = angle_range_label + 1
if close_idx_0 > 0:
weights_v[close_idx_0] *= (lines_v_mask == 0)
weights_h[close_idx_0] *= (lines_h_mask == 0)
if close_idx_1 < len(self.rot_angles):
weights_v[close_idx_1] *= (lines_v_mask == 0)
weights_h[close_idx_1] *= (lines_h_mask == 0)
pdb.set_trace()
weights = np.stack((weights_v, weights_h), 0)
###############################################################################
softmax_label_v = 255 * np.ones(label.shape, dtype=np.int64)
softmax_label_v[lines_v_mask.astype(bool)] = angle_range_label
softmax_label_h = 255 * np.ones(label.shape, dtype=np.int64)
softmax_label_h[lines_h_mask.astype(bool)] = angle_range_label
softmax_label = np.stack((softmax_label_v, softmax_label_h), 0)
###############################################################################
data.update(bin_label=bin_label,
softmax_label=softmax_label,
weights=weights)
return data
def _create_mask(_lines, width=8):
sz = mask.shape
line_mask = lines.create_grid(sz, _lines, width=width).astype(int)
return line_mask * mask
def forward(self, inputs):
def _gabor_bank(x):
bs = x.shape[0]
x = self.gabor_bank(x)
x = x.transpose(0, 1)
x = x[:-1] + x[1:]
n_intervals = x.shape[0] // 2
x_v = x[:n_intervals]
x_h = x[n_intervals:]
x_vh = torch.max(x_v, x_h).transpose(0, 1)
hist = self.relu(x_vh).view(bs, n_intervals, -1).mean(2)
return (x_v, x_h), hist
input_shape = inputs["image"].shape[-2:]
features = super(AngleNet, self).forward(inputs,
return_intermediate=True)
x = F.interpolate(features["decoder"],
size=input_shape,
mode='bilinear',
align_corners=False)
seg = self.classifier(x)
seg = seg.transpose(0, 1)
x_0, x_1 = seg[:2]
x_diff = x_1 - x_0
(x_gabor_v,
x_gabor_h), hist = _gabor_bank(x_diff.unsqueeze(0).transpose(0, 1))
x_max = torch.max(x_0, x_1).unsqueeze(0)
x_max = torch.cat([x_max, seg[2:]], 0).transpose(0, 1)
if self.line_clf is not None:
idx = inputs["angle_range_label"]
x_gabor_v = x_gabor_v[idx]
x_gabor_h = x_gabor_h[idx]
lines_scores_v = self.lines_detect(
x_gabor_v, inputs['lines_endpoints_v'].squeeze(0))
lines_scores_h = self.lines_detect(
x_gabor_h, inputs['lines_endpoints_h'].squeeze(0))
lines_scores = torch.cat([lines_scores_v, lines_scores_h])
result = OrderedDict()
if self.training:
result['out'] = OrderedDict()
if self.line_clf is not None:
result['out']["lines_scores"] = lines_scores
else:
result['out']["x_diff"] = x_diff
result['out']["x_max"] = x_max
result['out']["hist"] = hist
else:
result["hist"] = hist
lines_v = self.lines_detect_test(lines_scores_v,
inputs['proposed_lines_v'])
lines_h = self.lines_detect_test(lines_scores_h,
inputs['proposed_lines_h'])
lines_mask = lines.create_grid(tuple(x_gabor_v.shape[-2:]),
lines_v + lines_h)
result["seg"] = lines_mask
return result
def __getitem__(self, index): data = super(HistDataset, self).__getitem__(index) label_test = data['label_test'] label = data['label'] weights = data['weights'] angle_range_label = data['angle_range_label'] if angle_range_label == 255: return data ############################################################################### lines_v, _rot_angle = lines.extract_lines((label_test == 1), self.angle_range_v) lines_h, _ = lines.extract_lines((label_test == 1), self.angle_range_h) lines_v_mask = lines.create_grid(label.shape, lines_v, width=16) * (label == 0).astype(int) lines_h_mask = lines.create_grid(label.shape, lines_h, width=16) * (label == 0).astype(int) """ plt.figure() plt.imshow(lines_v_mask) plt.figure() plt.imshow(lines_h_mask) plt.show() """ _rot_angle = np.rad2deg(_rot_angle) angle_dist = np.abs(self.rot_angles - _rot_angle) ############################################################################### if self.combine: idx = angle_range_label n_angles = len(self.rot_angles) - 1 else: idx = np.argmin(angle_dist) n_angles = len(self.rot_angles) sz = (n_angles,) + label.shape bin_label_v = np.zeros(sz, dtype=np.float32) bin_label_v[idx] = lines_v_mask.astype(np.float32) bin_label_h = np.zeros(sz, dtype=np.float32) bin_label_h[idx] = lines_h_mask.astype(np.float32) bin_label = np.stack((bin_label_v, bin_label_h), 0) lines_v_mask_inv = (lines_v_mask != 1).astype(np.float32) lines_h_mask_inv = (lines_h_mask != 1).astype(np.float32) weights_v = np.repeat(weights[np.newaxis,...] * lines_v_mask_inv, n_angles, 0) weights_v[idx] = weights weights_h = np.repeat(weights[np.newaxis,...] * lines_h_mask_inv, n_angles, 0) weights_h[idx] = weights weights = np.stack((weights_v, weights_h), 0) ############################################################################### softmax_label_v = 255 * np.ones(label.shape, dtype=np.int64) softmax_label_v[lines_v_mask.astype(bool)] = angle_range_label softmax_label_h = 255 * np.ones(label.shape, dtype=np.int64) softmax_label_h[lines_h_mask.astype(bool)] = angle_range_label softmax_label = np.stack((softmax_label_v, softmax_label_h), 0) ############################################################################### data.update(bin_label=bin_label, softmax_label=softmax_label, weights=weights) return data