def show_patch_warpped_t(self, patch_template, patch_search):
cv2.namedWindow("patch_template", cv2.WND_PROP_FULLSCREEN)
cv2.namedWindow("patch_search", cv2.WND_PROP_FULLSCREEN)
img_w = patch_template.shape[-1]
cv2.imshow('patch_template', kornia.tensor_to_image(patch_template.byte()).reshape(-1, img_w, 3))
cv2.imshow('patch_search', kornia.tensor_to_image(patch_search.byte()).reshape(-1, img_w, 3))
cv2.waitKey(0)
def track(model, p, template_img, template_bbox, search_img, search_bbox):
pos_z, size_z = bbox2center_sz(template_bbox)
pos_x, size_x = bbox2center_sz(search_bbox)
model.template(template_img.to(device), pos_z.to(device),
size_z.to(device))
pscore, delta, pscore_size = model.track(search_img.to(device),
pos_x.to(device),
size_x.to(device))
scale_x = model.penalty.get_scale_x(size_x)
assert pscore.shape[0] == 1
tuple(
map(lambda x: x.squeeze_().numpy(),
[pos_x, size_x, template_bbox, search_bbox]))
template_img = np.ascontiguousarray(
kornia.tensor_to_image(template_img.byte()))
search_img = np.ascontiguousarray(kornia.tensor_to_image(
search_img.byte()))
best_pscore_id = np.argmax(pscore.squeeze().detach().cpu().numpy())
pred_in_img = delta.squeeze().detach().cpu().numpy(
)[:, best_pscore_id] / scale_x
lr = pscore_size.squeeze().detach().cpu().numpy(
)[best_pscore_id] * p.lr # lr for OTB
res_x = pred_in_img[0] + pos_x[0]
res_y = pred_in_img[1] + pos_x[1]
res_w = size_x[0] * (1 - lr) + pred_in_img[2] * lr
res_h = size_x[1] * (1 - lr) + pred_in_img[3] * lr
target_pos = np.array([res_x, res_y])
target_sz = np.array([res_w, res_h])
im_h, im_w = template_img.shape[0], template_img.shape[1]
target_pos[0] = max(0, min(im_w, target_pos[0]))
target_pos[1] = max(0, min(im_h, target_pos[1]))
target_sz[0] = max(10, min(im_w, target_sz[0]))
target_sz[1] = max(10, min(im_h, target_sz[1]))
x, y, w, h = template_bbox
x2, y2 = x + w, y + h
cv2.rectangle(template_img, (x, y), (x2, y2), (0, 255, 0), 4)
cv2.imshow('template', template_img)
x, y = (target_pos - target_sz / 2).astype(int)
x2, y2 = (target_pos + target_sz / 2).astype(int)
cv2.rectangle(search_img, (x, y), (x2, y2), (0, 255, 0), 8)
cv2.imshow('SiamMask', search_img)
key = cv2.waitKey(1)
return x, y, x2 - x, y2 - y
def gen_xcrop(self):
# state = self.state
# im_pert_template = state['im_pert_template']
# pert_sz_ratio = state['pert_sz_ratio']
# p = state['p']
# im_shape = state['im'].shape[0:2]
# bbox_pert_xcrop = scale_bbox(state['gts'][state['n_frame']], pert_sz_ratio)
# mask_xcrop = get_bbox_mask(shape=im_shape, bbox=bbox_pert_xcrop, mode='tensor').to(state['device'])
# bbox_pert_temp = scale_bbox(state['gts'][0], pert_sz_ratio)
# bbox_pert_xcrop = scale_bbox(state['gts'][state['n_frame']], pert_sz_ratio)
# im_pert_warped = warp(im_pert_template, bbox_pert_temp, bbox_pert_xcrop)
# im_xcrop = kornia.image_to_tensor(state['im']).to(state['device'])
# im_pert_xcrop = im_xcrop * (1-mask_xcrop) + im_pert_warped * mask_xcrop
# x_crop = test.get_subwindow_tracking_(im_pert_xcrop, state['target_pos'], p.instance_size, round(self.s_x), 0)
state = self.state
pert = state['pert']
pert_sz_ratio = state['pert_sz_ratio']
p = state['p']
device = state['device']
im_shape = state['im'].shape[0:2]
bbox_pert_xcrop = scale_bbox(state['gts'][state['n_frame']],
pert_sz_ratio)
mask_xcrop = get_bbox_mask(shape=im_shape,
bbox=bbox_pert_xcrop,
mode='tensor').to(device)
bbox_pert_xcrop = torch.tensor(bbox_pert_xcrop).unsqueeze(
dim=0).to(device)
im_xcrop = kornia.image_to_tensor(state['im']).to(
torch.float).unsqueeze(dim=0).to(device)
patch_warped_search = warp_patch(pert, im_xcrop, bbox_pert_xcrop)
patch_search = torch.where(mask_xcrop == 1, patch_warped_search,
im_xcrop)
x_crop = test.get_subwindow_tracking_(patch_search,
state['target_pos'],
p.instance_size, round(self.s_x),
0)
cv2.imshow('template',
kornia.tensor_to_image(state['pert_template'].byte()))
cv2.imshow('x_crop', kornia.tensor_to_image(x_crop.byte()))
cv2.waitKey(1)
return state['pert_template'], x_crop
def forward(self, x: torch.Tensor, mask=None):
max_val = x.max()
if max_val < 2.0:
img_np = (255 * K.tensor_to_image(x)).astype(np.uint8)
else:
img_np = K.tensor_to_image(x).astype(np.uint8)
if mask is not None:
mask = K.tensor_to_image(x).astype(np.uint8)
kpts, descs = self.features.detectAndCompute(img_np, mask)
lafs, resp = laf_from_opencv_kpts(kpts,
mrSize=self.mrSize,
with_resp=True,
device=x.device)
return lafs, resp, torch.from_numpy(descs).to(device=x.device)[None]
def test_batch_hsv_to_rgb(self, device):
data = torch.rand(3, 5, 5).to(device) # 3x5x5
# OpenCV
data_cv = kornia.tensor_to_image(data.clone())
data_cv[:, :, 0] = 360 * data_cv[:, :, 0]
expected = cv2.cvtColor(data_cv, cv2.COLOR_HSV2RGB)
expected = kornia.image_to_tensor(expected, False).to(device)
# Kornia
f = kornia.color.HsvToRgb()
data[0] = 2 * pi * data[0]
data = data.repeat(2, 1, 1, 1) # 2x3x5x5
expected = expected.repeat(2, 1, 1, 1) # 2x3x5x5
assert_allclose(f(data), expected)
data[:, 0] += 2 * pi
assert_allclose(f(data), expected)
data[:, 0] -= 4 * pi
assert_allclose(f(data), expected)
def test_hsv_to_rgb(self, device):
data = torch.rand(3, 5, 5).to(device) # 3x5x5
# OpenCV
data_cv = kornia.tensor_to_image(data.clone())
data_cv[:, :, 0] = 360 * data_cv[:, :, 0]
expected = cv2.cvtColor(data_cv, cv2.COLOR_HSV2RGB)
expected = kornia.image_to_tensor(expected, True).to(device)
r_expected = expected[0]
g_expected = expected[1]
b_expected = expected[2]
# Kornia
f = kornia.color.HsvToRgb()
data[0] = 2 * pi * data[0]
result = f(data)
r = result[0, :, :]
g = result[1, :, :]
b = result[2, :, :]
assert_allclose(r, r_expected)
assert_allclose(g, g_expected)
assert_allclose(b, b_expected)
def _convert(self):
src_img = kornia.image_to_tensor(self.src_img, keepdim=False)
dst_h, dst_w = 800, 800
# Compute perspective transform
M = kornia.get_perspective_transform(self.points_src, self.points_dst)
# Image to BEV transformation
dst_img = kornia.warp_perspective(src_img.float(),
M,
dsize=(dst_h, dst_w),
flags='bilinear',
border_mode='zeros')
# remove unwanted portion of BEV image. e.g for FRONT view dst point should not be higher than 450.
if 'FRONT' in self.path:
dst_img[:, :, 400:, :] = 0
if 'BACK' in self.path:
dst_img[:, :, :400, :] = 0
if 'LEFT' in self.path:
dst_img[:, :, :, 400:] = 0
if 'RIGHT' in self.path:
dst_img[:, :, :, :400] = 0
dst_img = kornia.tensor_to_image(dst_img.byte())
return dst_img
def __call__(self, image):
random_kernel = random.randrange(25, 35, 2)
random_kernel = (random_kernel, random_kernel)
random_sigma = random.uniform(5, 10)
random_sigma = (random_sigma, random_sigma)
image = torch.unsqueeze(kornia.image_to_tensor(image).float(), dim=0)
image = gaussian_blur2d(image, random_kernel, random_sigma)
return kornia.tensor_to_image(image)
def imshow(input: torch.Tensor):
fig = plt.figure(figsize=(8,6), dpi=150)
out: torch.Tensor = torchvision.utils.make_grid(input, nrow=2, padding=1)
out_np: np.ndarray = kornia.tensor_to_image(out)
plt.imshow(out_np)
plt.axis('off')
plt.show()
plt.close(fig)
def forward(self, x: torch.Tensor, mask=None):
self.affnet = self.affnet.to(x.device)
max_val = x.max()
if max_val < 2.0:
img_np = (255 * K.tensor_to_image(x)).astype(np.uint8)
else:
img_np = K.tensor_to_image(x).astype(np.uint8)
if mask is not None:
mask = K.tensor_to_image(x).astype(np.uint8)
kpts = self.features.detect(img_np, mask)
lafs, resp = laf_from_opencv_kpts(kpts,
mrSize=self.mrSize,
with_resp=True,
device=x.device)
ori = KF.get_laf_orientation(lafs)
lafs = self.affnet(lafs, x.mean(dim=1, keepdim=True))
lafs = KF.set_laf_orientation(lafs, ori)
return lafs, resp
def to_numpy_image(img: Union[str, np.array, torch.Tensor]):
if type(img) is str:
img_out = cv2.cvtColor(cv2.imread(img), cv2.COLOR_BGR2RGB)
elif isinstance(img, torch.Tensor):
img_out = K.tensor_to_image(img)
elif isinstance(img, np.ndarray):
img_out = img
else:
raise TypeError('img should be str, np.array or torch.Tensor')
return img_out
def tensor_to_images(tensor):
out = []
for t in tensor:
arr = (kornia.tensor_to_image(t) * 255).astype('uint8')
img = Image.fromarray(arr.astype("uint8"))
rawBytes = io.BytesIO()
img.save(rawBytes, "PNG")
rawBytes.seek(0) # return to the start of the file
out.append(base64.b64encode(rawBytes.read()).decode('utf-8'))
return out
def gaussian_blur_pytorch(img_gray: np.ndarray, sigma: float = 1.0):
t_gray = (kornia.utils.image_to_tensor(
np.array(img_gray)).float().unsqueeze(0).unsqueeze(0))
smoothing = GaussianSmoothing(1, 5, sigma, dim=3)
input_t = F.pad(t_gray, (2, 2, 2, 2, 2, 2))
output = smoothing(input_t)
output: np.ndarray = kornia.tensor_to_image(
output.squeeze(0).squeeze(0).float())
return output
def apply_blur_face(img: torch.Tensor, img_vis: np.ndarray,
det: FaceDetectorResult):
# crop the face
x1, y1 = det.xmin.int(), det.ymin.int()
x2, y2 = det.xmax.int(), det.ymax.int()
roi = img[..., y1:y2, x1:x2]
# apply blurring and put back to the visualisation image
roi = K.filters.gaussian_blur2d(roi, (21, 21), (35., 35.))
roi = K.color.rgb_to_bgr(roi)
img_vis[y1:y2, x1:x2] = K.tensor_to_image(roi)
def show_attack_plt(self, pscore, bbox, bbox_src, patch):
fig, axes = plt.subplots(1,3,num='attacking')
ax = axes[0]
ax.set_title('patch')
ax.imshow(kornia.tensor_to_image(patch.byte()))
ax = axes[1]
ax.set_title('template')
ax.imshow(kornia.tensor_to_image(self.model.template_cropped.byte()).reshape(-1, 127, 3))
ax = axes[2]
ax.set_title('result')
ax.imshow(kornia.tensor_to_image(self.model.search_cropped.byte()).reshape(-1, 255, 3))
for i, xywh in enumerate(bbox):
x, y, w, h = xywh
rect = patches.Rectangle((x, y+i*255), w, h, linewidth=1, edgecolor='b', facecolor='none')
ax.add_patch(rect)
for i, xywh in enumerate(bbox_src):
x, y, w, h = xywh
rect = patches.Rectangle((x, y+i*255), w, h, linewidth=1, edgecolor='r', facecolor='none')
ax.add_patch(rect)
plt.pause(0.01)
def get_images(self, uint=False):
images = []
with torch.no_grad():
for n in range(self.n_classes):
img_np = kornia.tensor_to_image(self.param_fn(self.input_tensor[n]))
if uint == True:
scaled_img_np = img_np * 255
img_ = scaled_img_np.astype(int)
else:
img_ = img_np
images.append(img_)
return images
def init(model, template_img, template_bbox):
pos_z, size_z = bbox2center_sz(template_bbox)
model.template(template_img.to(device), pos_z.to(device),
size_z.to(device))
template_img = np.ascontiguousarray(
kornia.tensor_to_image(template_img.byte()))
x, y, w, h = template_bbox.squeeze().cpu().numpy()
x2, y2 = x + w, y + h
cv2.rectangle(template_img, (x, y), (x2, y2), (0, 255, 0), 4)
cv2.imshow('template', template_img)
def test_rgb_to_hls(self, device):
data = torch.rand(3, 5, 5).to(device)
# OpenCV
data_cv = kornia.tensor_to_image(data.clone())
expected = cv2.cvtColor(data_cv, cv2.COLOR_RGB2HLS)
expected = kornia.image_to_tensor(expected, True).to(device)
expected[0] = 2 * math.pi * expected[0] / 360.
f = kornia.color.RgbToHls()
assert_allclose(f(data), expected)
def track(model, p, search_img, search_bbox):
pos_x, size_x = bbox2center_sz(search_bbox)
pscore, delta, pscore_size = model.track(search_img.to(device),
pos_x.to(device),
size_x.to(device))
scale_x = model.penalty.get_scale_x(size_x)
assert pscore.shape[0] == 1
tuple(
map(lambda x: x.squeeze_().cpu().numpy(),
[pos_x, size_x, search_bbox]))
search_img = np.ascontiguousarray(kornia.tensor_to_image(
search_img.byte()))
best_pscore_id = np.argmax(pscore.squeeze().detach().cpu().numpy())
pred_in_img = delta.squeeze().detach().cpu().numpy(
)[:, best_pscore_id] / scale_x.cpu().numpy()
lr = pscore_size.squeeze().detach().cpu().numpy(
)[best_pscore_id] * p.lr # lr for OTB
res_x = pred_in_img[0] + pos_x[0]
res_y = pred_in_img[1] + pos_x[1]
res_w = size_x[0] * (1 - lr) + pred_in_img[2] * lr
res_h = size_x[1] * (1 - lr) + pred_in_img[3] * lr
target_pos = np.array([res_x, res_y])
target_sz = np.array([res_w, res_h])
im_h, im_w = search_img.shape[0], search_img.shape[1]
target_pos[0] = max(0, min(im_w, target_pos[0]))
target_pos[1] = max(0, min(im_h, target_pos[1]))
target_sz[0] = max(10, min(im_w, target_sz[0]))
target_sz[1] = max(10, min(im_h, target_sz[1]))
x, y = (target_pos - target_sz / 2).astype(int)
x2, y2 = (target_pos + target_sz / 2).astype(int)
cv2.rectangle(search_img, (x, y), (x2, y2), (0, 255, 0), 8)
cv2.imshow('SiamMask', search_img)
key = cv2.waitKey(1)
global i, save_img
i += 1
if save_img:
status = cv2.imwrite('./results/res_{:03d}.jpg'.format(i),
cv2.resize(search_img, (384, 216)))
print(status, 'results//res_{:03d}.jpg'.format(i))
return x, y, x2 - x, y2 - y
def tvdenoise_kornia(img, regularization_amount=0.001, max_iter=50):
img_t = kornia.utils.image_to_tensor(np.array(img)).float().unsqueeze(0)
tv_denoiser = TVDenoise(img_t, regularization_amount)
optimizer = torch.optim.SGD(tv_denoiser.parameters(), lr=0.1, momentum=0.9)
for i in range(max_iter):
optimizer.zero_grad()
loss = tv_denoiser()
if i % 25 == 0:
logger.debug("Loss in iteration {} of {}: {:.3f}".format(
i, max_iter, loss.item()))
loss.backward()
optimizer.step()
img_clean: np.ndarray = kornia.tensor_to_image(
tv_denoiser.get_denoised_image())
return img_clean
def test_batch_rgb_to_hls(self, device):
data = torch.rand(3, 5, 5).to(device)
# OpenCV
data_cv = kornia.tensor_to_image(data.clone())
expected = cv2.cvtColor(data_cv, cv2.COLOR_RGB2HLS)
expected = kornia.image_to_tensor(expected, False).to(device)
expected[:, 0] = 2 * math.pi * expected[:, 0] / 360.
# Kornia
f = kornia.color.RgbToHls()
data = data.repeat(2, 1, 1, 1) # 2x3x5x5
expected = expected.repeat(2, 1, 1, 1) # 2x3x5x5
assert_allclose(f(data), expected)
def laplacian(img: np.ndarray, kernel_size) -> np.ndarray:
"""Laplacian filter a numpy array
Arguments: np.ndarray (D,H,W)
Input image
Returns:
np.ndarray -- filtered array
"""
img_clean = rescale_denan(img_as_float(np.clip(img, 0.0, 1.0)))
img_clean_t = kornia.utils.image_to_tensor(
np.array(img_clean)).float().unsqueeze(0)
kernel_size = int(kernel_size)
if kernel_size % 2 == 0:
kernel_size += 1
laplacian: torch.Tensor = kornia.filters.laplacian(img_clean_t,
kernel_size=kernel_size)
laplacian_img: np.ndarray = kornia.tensor_to_image(laplacian.float())
return np.nan_to_num(laplacian_img)
def spatial_gradient_3d(vol_gray: np.ndarray, dim=0) -> np.ndarray:
"""Spatial gradient of a array of intensity values
Arguments:
vol_gray {np.ndarray} -- input array
Returns:
np.ndarray -- filtered array
"""
img_gray = img_as_float(np.clip(vol_gray, 0.0, 1.0))
t_gray = (kornia.utils.image_to_tensor(
np.array(img_gray)).float().unsqueeze(0).unsqueeze(0))
spatialgradient3d = kornia.filters.SpatialGradient3d(mode="diff")
result = spatialgradient3d(t_gray)
result = result[0, 0, dim, :]
result_arr: np.ndarray = kornia.tensor_to_image(result.float())
logger.debug(f"Calculated gradient of shape {result_arr.shape}")
return result_arr
def hessian_eigvals_image(data, sigma, correct=False):
"""Hessian eigenvalues using Pytorch
Parameters
----------
data : np.ndarray(D,H,W)
Input image
sigma : Vector of 3 float
Kernel size
correct : bool, optional
Adjust sigma to square of input sigma, by default False
Returns
-------
np.array
(D,H,W,3) Hessian eigenvalues
"""
H = compute_hessian(data, sigma)
if correct:
if not isinstance(sigma, numbers.Number):
s = max(sigma) ** 2
else:
s = sigma ** 2
Hxx, Hxy, Hxz, Hyy, Hyz, Hzz = [h * s for h in H]
else:
Hxx, Hxy, Hxz, Hyy, Hyz, Hzz = H
Hs = torch.FloatTensor([[Hzz, Hyz, Hxz], [Hyz, Hyy, Hxy], [Hxz, Hxy, Hxx]])
Hs_batch = Hs.reshape((3, 3, Hs.shape[-3] * Hs.shape[-2] * Hs.shape[-1]))
e, v = torch.symeig(Hs_batch.permute((2, 0, 1)))
e = e.reshape((Hs.shape[-3], Hs.shape[-2], Hs.shape[-1], 3))
img: np.ndarray = kornia.tensor_to_image(e)
img = np.transpose(img, (0, 3, 1, 2))
img = img[:, :, :, 0] # return primay eigenvalue
return np.nan_to_num(img)
def gaussian_blur_kornia(img: np.ndarray, sigma):
"""Gaussian blur using Kornia Filter3D
Parameters
----------
img : np.ndarray
Input image array
sigma : float, optional
sigma value, by default 1.0
Returns
-------
output
filtered numpy array
"""
img_t = (kornia.utils.image_to_tensor(
np.array(img)).float().unsqueeze(0).unsqueeze(0))
output = gaussian_blur_t(img_t, sigma)
output: np.ndarray = kornia.tensor_to_image(
output.squeeze(0).squeeze(0).float())
return output
def test_rgb_to_hsv(self, device):
data = torch.rand(3, 5, 5).to(device)
# OpenCV
data_cv = kornia.tensor_to_image(data.clone())
expected = cv2.cvtColor(data_cv, cv2.COLOR_RGB2HSV)
expected = kornia.image_to_tensor(expected, True).to(device)
h_expected = 2 * math.pi * expected[0] / 360.
s_expected = expected[1]
v_expected = expected[2]
f = kornia.color.RgbToHsv()
result = f(data)
h = result[0, :, :]
s = result[1, :, :]
v = result[2, :, :]
assert_allclose(h, h_expected)
assert_allclose(s, s_expected)
assert_allclose(v, v_expected)
def train(arg):
log_writer = None
if arg.save_logs:
log_path = './logs/align_' + arg.dataset + '_' + arg.split
if not os.path.exists(log_path):
os.makedirs(log_path)
log_writer = SummaryWriter(log_dir=log_path)
def log(tag, scalar, step):
log_writer.add_scalar(tag, scalar, step)
def log_img(tag, img, step):
log_writer.add_image(tag, img, step)
else:
def log(tag, scalar, step):
pass
def log_img(tag, img, step):
pass
epoch = None
devices = get_devices_list(arg)
print('***** Training decoder *****')
print('Training parameters:\n' + '# Dataset: ' + arg.dataset +
'\n' + '# Dataset split: ' + arg.split + '\n' +
'# Batchsize: ' + str(arg.batch_size) + '\n' +
'# Num workers: ' + str(arg.workers) + '\n' +
'# PDB: ' + str(arg.PDB) + '\n' +
'# Use GPU: ' + str(arg.cuda) + '\n' +
'# Start lr: ' + str(arg.lr) + '\n' +
'# Max epoch: ' + str(arg.max_epoch) + '\n' +
'# Loss type: ' + arg.loss_type + '\n' +
'# Resumed model: ' + str(arg.resume_epoch > 0))
if arg.resume_epoch > 0:
print('# Resumed epoch: ' + str(arg.resume_epoch))
print('Creating networks ...')
estimator = create_model_estimator(arg, devices, eval=True)
estimator.eval()
regressor = None
if arg.regressor_loss:
regressor = create_model_regressor(arg, devices, eval=True)
regressor.eval()
align = create_model_align(arg, devices)
align.train()
edge = create_model_edge(arg, devices, eval=True)
edge.eval()
print('Creating networks done!')
optimizer_align, scheduler_align = create_optimizer(arg,
align.parameters(),
create_scheduler=True)
if arg.loss_type == 'L2':
criterion = nn.MSELoss()
elif arg.loss_type == 'L1':
criterion = nn.L1Loss()
elif arg.loss_type == 'smoothL1':
criterion = nn.SmoothL1Loss()
elif arg.loss_type == 'wingloss':
criterion = WingLoss(omega=arg.wingloss_omega,
epsilon=arg.wingloss_epsilon)
else:
criterion = AWingLoss(arg.wingloss_omega,
theta=arg.wingloss_theta,
epsilon=arg.wingloss_epsilon,
alpha=arg.wingloss_alpha)
if arg.cuda:
criterion = criterion.cuda(device=devices[0])
print('Loading dataset ...')
trainset = GeneralDataset(arg, dataset=arg.dataset, split=arg.split)
dataloader = torch.utils.data.DataLoader(trainset,
batch_size=arg.batch_size,
shuffle=arg.shuffle,
num_workers=arg.workers,
pin_memory=True)
steps_per_epoch = len(dataloader)
print('Loading dataset done!')
# evolving training
print('Start training ...')
for epoch in range(arg.resume_epoch, arg.max_epoch):
global_step_base = epoch * steps_per_epoch
forward_times_per_epoch, sum_loss = 0, 0.
for data in tqdm.tqdm(dataloader):
forward_times_per_epoch += 1
global_step = global_step_base + forward_times_per_epoch
_, input_images, _, gt_coords_xy, _, _, _, _ = data
if arg.cuda:
input_images = input_images.cuda(device=devices[0])
gt_coords_xy = gt_coords_xy.cuda(device=devices[0])
with torch.no_grad():
heatmaps = estimator(input_images)[-1]
# heatmaps[heatmaps < arg.boundary_cutoff_lambda * heatmaps.max()] = 0
heatmaps = edge(heatmaps)
optimizer_align.zero_grad()
coords_predict = align(heatmaps)
gt_coords_xy = gt_coords_xy / 4
loss = criterion(coords_predict, gt_coords_xy)
log('loss', loss.item(), global_step)
loss.backward()
optimizer_align.step()
sum_loss += loss.item()
show_img(tensor_to_image(input_images[0]))
heatmap_sum = np.uint8(
get_heatmap_gray(heatmaps[0].unsqueeze(0),
denorm=True).detach().squeeze().cpu().numpy())
heatmap_sum = np.moveaxis(np.tile(heatmap_sum, [3, 1, 1]), 0, -1)
heatmap_sum[:, :, 1:2] = 0
gt_coords_xy = gt_coords_xy[0].detach().cpu().squeeze().numpy()
for i in range(0, 2 * kp_num[arg.dataset], 2):
draw_circle(heatmap_sum,
(int(gt_coords_xy[i]), int(gt_coords_xy[i + 1])),
color=(255, 255, 255))
show_img(cv2.resize(heatmap_sum, (256, 256)))
heatmap_sum = np.uint8(
get_heatmap_gray(heatmaps[0].unsqueeze(0),
denorm=True).detach().squeeze().cpu().numpy())
heatmap_sum = np.moveaxis(np.tile(heatmap_sum, [3, 1, 1]), 0, -1)
heatmap_sum[:, :, 1:2] = 0
coords_predict = coords_predict[0].detach().cpu().squeeze().numpy()
for i in range(0, 2 * kp_num[arg.dataset], 2):
draw_circle(
heatmap_sum,
(int(coords_predict[i]), int(coords_predict[i + 1])),
color=(255, 255, 255))
show_img(cv2.resize(heatmap_sum, (256, 256)))
mean_sum_loss = sum_loss / forward_times_per_epoch
if scheduler_align is not None:
scheduler_align.step(mean_sum_loss)
if (epoch + 1) % arg.save_interval == 0:
torch.save(
align.state_dict(), arg.save_folder + 'align_' + arg.dataset +
'_' + str(epoch + 1) + '.pth')
print('\nepoch: {:0>4d} | loss: {:.10f}'.format(
epoch,
mean_sum_loss,
))
torch.save(
align.state_dict(), arg.save_folder + 'align_' + arg.dataset + '_' +
str(epoch + 1) + '.pth')
print('Training done!')
if __name__ == "__main__":
import glob
from os.path import join
import cv2
# setup device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
# load images and gt
fpath = '/DataServer/car/car-1/img/'
N_IMS = 2
img_files = sorted(glob.glob(join(fpath, '*.jp*')))
ims = [cv2.imread(imf) for imf in img_files[:N_IMS]]
ims = torch.tensor(ims, device=device, dtype=float).permute(0,3,1,2)
with open(fpath + '/../groundtruth.txt', "r") as f:
gts = f.readlines()
split_flag = ',' if ',' in gts[0] else '\t'
gts = list(map(lambda x: list(map(int, x.rstrip().split(split_flag))), gts))
gts = np.array(gts)[:N_IMS]
pos = torch.tensor(gts[:,:2] + gts[:,2:]/2, device=device)
sz = torch.tensor(gts[:,2:], device=device).max(dim=1)[0]
# test SubWindow
model = SubWindow()
out = model(ims, pos, sz)
for i in range(out.shape[0]):
cv2.imshow('im_ori', kornia.tensor_to_image(ims[i].byte()))
cv2.imshow('im', kornia.tensor_to_image(out[i].byte()))
cv2.waitKey(0)
import matplotlib.pyplot as plt
# read the image with OpenCV
img: np.ndarray = cv2.imread('./download.png')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# convert to torch tensor
data: torch.tensor = kornia.image_to_tensor(img, keepdim=False) # BxCxHxW
# create the operator
gauss = kornia.filters.GaussianBlur2d((11, 11), (10.5, 10.5))
# blur the image
x_blur: torch.tensor = gauss(data.float())
# convert back to numpy
img_blur: np.ndarray = kornia.tensor_to_image(x_blur.byte())
# Create the plot
fig, axs = plt.subplots(1, 2, figsize=(16, 10))
axs = axs.ravel()
axs[0].axis('off')
axs[0].set_title('image source')
axs[0].imshow(img)
axs[1].axis('off')
axs[1].set_title('image blurred')
axs[1].imshow(img_blur)
pass
]])
# the destination points are the image vertexes
h, w = 64, 128 # destination size
points_dst = torch.tensor([[
[0., 0.], [w - 1., 0.], [w - 1., h - 1.], [0., h - 1.],
]])
# compute perspective transform
M: torch.tensor = kornia.get_perspective_transform(points_src, points_dst)
# warp the original image by the found transform
data_warp: torch.tensor = kornia.warp_perspective(data.float(), M, dsize=(h, w))
# convert back to numpy
img_warp: np.ndarray = kornia.tensor_to_image(data_warp.byte())
# draw points into original image
for i in range(4):
center = tuple(points_src[0, i].long().numpy())
img = cv2.circle(img.copy(), center, 5, (0, 255, 0), -1)
# create the plot
fig, axs = plt.subplots(1, 2, figsize=(16, 10))
axs = axs.ravel()
axs[0].axis('off')
axs[0].set_title('image source')
axs[0].imshow(img)
axs[1].axis('off')
