import re
from pathlib import Path
import glob
import logging
import numpy as np
import torch
import cv2
import os
import math
from adamp import AdamP
import random
import torch.nn as nn
_logger = None
def increment_path(path):
# Increment path, i.e. runs/exp1 --> runs/exp{sep}1, runs/exp{sep}2 etc.
res = re.search("\d+", path)
if res is None:
print("Set initial exp number!")
exit(1)
if not Path(path).exists():
return str(path)
else:
path = path[:res.start()]
dirs = glob.glob(f"{path}*") # similar paths
matches = [re.search(rf"%s(\d+)" % Path(path).stem, d) for d in dirs]
i = [int(m.groups()[0]) for m in matches if m] # indices
n = max(i) + 1 # increment number
return f"{path}{n}" # update path
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, fmt=':f'):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def create_logger(log_file, level=logging.INFO):
global _logger
_logger = logging.getLogger()
formatter = logging.Formatter(
'[%(asctime)s][%(filename)15s][line:%(lineno)4d][%(levelname)8s] %(message)s')
fh = logging.FileHandler(log_file)
fh.setFormatter(formatter)
sh = logging.StreamHandler()
sh.setFormatter(formatter)
_logger.setLevel(level)
_logger.addHandler(fh)
_logger.addHandler(sh)
return _logger
def get_mgrid(sidelen, dim=2):
'''Generates a flattened grid of (x,y,...) coordinates in a range of -1 to 1.'''
if isinstance(sidelen, int):
sidelen = dim * (sidelen,)
if dim == 2:
pixel_coords = np.stack(np.mgrid[:sidelen[0], :sidelen[1]], axis=-1)[None, ...].astype(np.float32)
pixel_coords[0, :, :, 0] = pixel_coords[0, :, :, 0] / (sidelen[0] - 1)
pixel_coords[0, :, :, 1] = pixel_coords[0, :, :, 1] / (sidelen[1] - 1)
elif dim == 3:
pixel_coords = np.stack(np.mgrid[:sidelen[0], :sidelen[1], :sidelen[2]], axis=-1)[None, ...].astype(np.float32)
pixel_coords[..., 0] = pixel_coords[..., 0] / max(sidelen[0] - 1, 1)
pixel_coords[..., 1] = pixel_coords[..., 1] / (sidelen[1] - 1)
pixel_coords[..., 2] = pixel_coords[..., 2] / (sidelen[2] - 1)
else:
raise NotImplementedError('Not implemented for dim=%d' % dim)
pixel_coords -= 0.5
pixel_coords *= 2.
pixel_coords = torch.Tensor(pixel_coords).view(-1, dim)
return pixel_coords
def lin2img(tensor, image_resolution=None):
batch_size, num_samples, channels = tensor.shape
if image_resolution is None:
width = np.sqrt(num_samples).astype(int)
height = width
else:
if isinstance(image_resolution, int):
image_resolution = (image_resolution, image_resolution)
height = image_resolution[0]
width = image_resolution[1]
return tensor.permute(0, 2, 1).contiguous().view(batch_size, channels, height, width)
def normalize(x, opt, mode='normal'):
device = x.device
mean = torch.tensor(np.array(opt.transform_mean), dtype=x.dtype)[np.newaxis, :, np.newaxis, np.newaxis].to(device)
var = torch.tensor(np.array(opt.transform_var), dtype=x.dtype)[np.newaxis, :, np.newaxis, np.newaxis].to(device)
if mode == 'normal':
return (x - mean) / var
elif mode == 'inv':
return x * var + mean
def prepare_cooridinate_input(mask, dim=2):
'''Generates a flattened grid of (x,y,...) coordinates in a range of -1 to 1.'''
if mask.shape[0] == mask.shape[1]:
sidelen = mask.shape[0]
else:
sidelen = mask.shape[:2]
if isinstance(sidelen, int):
sidelen = dim * (sidelen,)
if dim == 2:
pixel_coords = np.stack(np.mgrid[:sidelen[0], :sidelen[1]], axis=-1)[None, ...].astype(np.float32)
pixel_coords[0, :, :, 0] = pixel_coords[0, :, :, 0] / (sidelen[0] - 1)
pixel_coords[0, :, :, 1] = pixel_coords[0, :, :, 1] / (sidelen[1] - 1)
elif dim == 3:
pixel_coords = np.stack(np.mgrid[:sidelen[0], :sidelen[1], :sidelen[2]], axis=-1)[None, ...].astype(np.float32)
pixel_coords[..., 0] = pixel_coords[..., 0] / max(sidelen[0] - 1, 1)
pixel_coords[..., 1] = pixel_coords[..., 1] / (sidelen[1] - 1)
pixel_coords[..., 2] = pixel_coords[..., 2] / (sidelen[2] - 1)
else:
raise NotImplementedError('Not implemented for dim=%d' % dim)
pixel_coords -= 0.5
pixel_coords *= 2.
return pixel_coords.squeeze(0).transpose(2, 0, 1)
def visualize(real, composite, mask, pred_fg, pred_harmonized, lut_transform_image, opt, epoch,
show=False, wandb=True, isAll=False, step=None):
save_path = os.path.join(opt.save_path, "figs", str(epoch))
os.makedirs(save_path, exist_ok=True)
if isAll:
final_index = 1
"""
Uncomment the following code if you want to save all the results, otherwise will only save the first image
of each batch
"""
# final_index = len(real)
else:
final_index = 1
for id in range(final_index):
if show:
cv2.imshow("pred_fg", normalize(pred_fg, opt, 'inv')[id].permute(1, 2, 0).cpu().numpy())
cv2.imshow("real", normalize(real, opt, 'inv')[id].permute(1, 2, 0).cpu().numpy())
cv2.imshow("lut_transform", normalize(lut_transform_image, opt, 'inv')[id].permute(1, 2, 0).cpu().numpy())
cv2.imshow("composite", normalize(composite, opt, 'inv')[id].permute(1, 2, 0).cpu().numpy())
cv2.imshow("mask", mask[id].permute(1, 2, 0).cpu().numpy())
cv2.imshow("pred_harmonized_image",
normalize(pred_harmonized, opt, 'inv')[id].permute(1, 2, 0).cpu().numpy())
cv2.waitKey()
if not opt.INRDecode:
real_tmp = cv2.cvtColor(
normalize(real, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(
np.uint8),
cv2.COLOR_RGB2BGR)
composite_tmp = cv2.cvtColor(
normalize(composite, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(
np.uint8), cv2.COLOR_RGB2BGR)
mask_tmp = mask[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(np.uint8)
lut_transform_image_tmp = cv2.cvtColor(
normalize(lut_transform_image, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(
0., 255.).numpy().astype(np.uint8), cv2.COLOR_RGB2BGR)
else:
pred_fg_tmp = cv2.cvtColor(
normalize(pred_fg, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(
np.uint8), cv2.COLOR_RGB2BGR)
real_tmp = cv2.cvtColor(
normalize(real, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(
np.uint8),
cv2.COLOR_RGB2BGR)
composite_tmp = cv2.cvtColor(
normalize(composite, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(
np.uint8), cv2.COLOR_RGB2BGR)
lut_transform_image_tmp = cv2.cvtColor(
normalize(lut_transform_image, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(
0., 255.).numpy().astype(np.uint8), cv2.COLOR_RGB2BGR)
mask_tmp = mask[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(0., 255.).numpy().astype(np.uint8)
pred_harmonized_tmp = cv2.cvtColor(
normalize(pred_harmonized, opt, 'inv')[id].permute(1, 2, 0).cpu().mul_(255.).clamp_(
0., 255.).numpy().astype(np.uint8), cv2.COLOR_RGB2BGR)
if isAll:
cv2.imwrite(os.path.join(save_path, f"{step}_{id}_composite.jpg"), composite_tmp)
cv2.imwrite(os.path.join(save_path, f"{step}_{id}_real.jpg"), real_tmp)
if opt.INRDecode:
cv2.imwrite(os.path.join(save_path, f"{step}_{id}_pred_harmonized_image.jpg"), pred_harmonized_tmp)
cv2.imwrite(os.path.join(save_path, f"{step}_{id}_lut_transform_image.jpg"), lut_transform_image_tmp)
cv2.imwrite(os.path.join(save_path, f"{step}_{id}_mask.jpg"), mask_tmp)
else:
if not opt.INRDecode:
cv2.imwrite(os.path.join(save_path, f"real_{step}_{id}.jpg"), real_tmp)
cv2.imwrite(os.path.join(save_path, f"composite_{step}_{id}.jpg"), composite_tmp)
cv2.imwrite(os.path.join(save_path, f"mask_{step}_{id}.jpg"), mask_tmp)
cv2.imwrite(os.path.join(save_path, f"lut_transform_image_{step}_{id}.jpg"), lut_transform_image_tmp)
else:
cv2.imwrite(os.path.join(save_path, f"pred_fg_{step}_{id}.jpg"), pred_fg_tmp)
cv2.imwrite(os.path.join(save_path, f"real_{step}_{id}.jpg"), real_tmp)
cv2.imwrite(os.path.join(save_path, f"composite_{step}_{id}.jpg"), composite_tmp)
cv2.imwrite(os.path.join(save_path, f"mask_{step}_{id}.jpg"), mask_tmp)
cv2.imwrite(os.path.join(save_path, f"pred_harmonized_image_{step}_{id}.jpg"), pred_harmonized_tmp)
cv2.imwrite(os.path.join(save_path, f"lut_transform_image_{step}_{id}.jpg"), lut_transform_image_tmp)
"Only upload images of the first batch of the first epoch to save storage."
if wandb and id == 0 and step == 0:
import wandb
real_tmp = wandb.Image(real_tmp, caption=epoch)
composite_tmp = wandb.Image(composite_tmp, caption=epoch)
if opt.INRDecode:
pred_fg_tmp = wandb.Image(pred_fg_tmp, caption=epoch)
pred_harmonized_tmp = wandb.Image(pred_harmonized_tmp, caption=epoch)
lut_transform_image_tmp = wandb.Image(lut_transform_image_tmp, caption=epoch)
mask_tmp = wandb.Image(mask_tmp, caption=epoch)
if not opt.INRDecode:
wandb.log(
{"pic/real": real_tmp, "pic/composite": composite_tmp,
"pic/mask": mask_tmp,
"pic/lut_trans": lut_transform_image_tmp,
"pic/epoch": epoch})
else:
wandb.log(
{"pic/pred_fg": pred_fg_tmp, "pic/real": real_tmp, "pic/composite": composite_tmp,
"pic/mask": mask_tmp,
"pic/lut_trans": lut_transform_image_tmp,
"pic/pred_harmonized": pred_harmonized_tmp,
"pic/epoch": epoch})
wandb.log({})
def get_optimizer(model, opt_name, opt_kwargs):
params = []
base_lr = opt_kwargs['lr']
for name, param in model.named_parameters():
param_group = {'params': [param]}
if not param.requires_grad:
params.append(param_group)
continue
if not math.isclose(getattr(param, 'lr_mult', 1.0), 1.0):
# print(f'Applied lr_mult={param.lr_mult} to "{name}" parameter.')
param_group['lr'] = param_group.get('lr', base_lr) * param.lr_mult
params.append(param_group)
optimizer = {
'sgd': torch.optim.SGD,
'adam': torch.optim.Adam,
'adamw': torch.optim.AdamW,
'adamp': AdamP
}[opt_name.lower()](params, **opt_kwargs)
return optimizer
def improved_efficient_matmul(a, c, index, batch=256):
"""
Reduce the unneed memory cost, but the speed is very slow.
:param a: N * I * J
:param b: N * J * K
:return: N * I * K
"""
"The first can only support when a is not requires_grad_, and have high speed. While the second one supports "
"whatever situations, but speed is quite slow. More Details in "
"https://discuss.pytorch.org/t/many-weird-phenomena-about-torch-matmul-operation/158208"
# out = torch.cat(
# [torch.matmul(a[i * batch:i * batch + batch, :, :], c[index[i * batch:i * batch + batch], :, :]) for i in
# range(a.shape[0] // batch)], dim=0)
batch = 1
out = torch.cat(
[torch.matmul(a[i * batch:i * batch + batch, :, :], c[index[i * batch], :, :]) for i in
range(a.shape[0] // batch)], dim=0)
return out
class LRMult(object):
def __init__(self, lr_mult=1.):
self.lr_mult = lr_mult
def __call__(self, m):
if getattr(m, 'weight', None) is not None:
m.weight.lr_mult = self.lr_mult
if getattr(m, 'bias', None) is not None:
m.bias.lr_mult = self.lr_mult
def customRandomCrop(objects, crop_height, crop_width, h_start=None, w_start=None):
if h_start is None:
h_start = random.random()
if w_start is None:
w_start = random.random()
if isinstance(objects, list):
out = []
for obj in objects:
out.append(random_crop(obj, crop_height, crop_width, h_start, w_start))
else:
out = random_crop(objects, crop_height, crop_width, h_start, w_start)
return out, h_start, w_start
def get_random_crop_coords(height: int, width: int, crop_height: int, crop_width: int, h_start: float,
w_start: float):
y1 = int((height - crop_height) * h_start)
y2 = y1 + crop_height
x1 = int((width - crop_width) * w_start)
x2 = x1 + crop_width
return x1, y1, x2, y2
def random_crop(img: np.ndarray, crop_height: int, crop_width: int, h_start: float, w_start: float):
height, width = img.shape[:2]
if height < crop_height or width < crop_width:
raise ValueError(
"Requested crop size ({crop_height}, {crop_width}) is "
"larger than the image size ({height}, {width})".format(
crop_height=crop_height, crop_width=crop_width, height=height, width=width
)
)
x1, y1, x2, y2 = get_random_crop_coords(height, width, crop_height, crop_width, h_start, w_start)
img = img[y1:y2, x1:x2]
return img
class PadToDivisor:
def __init__(self, divisor):
super().__init__()
self.divisor = divisor
def transform(self, images):
self._pads = (*self._get_dim_padding(images[0].shape[-1]), *self._get_dim_padding(images[0].shape[-2]))
self.pad_operation = nn.ZeroPad2d(padding=self._pads)
out = []
for im in images:
out.append(self.pad_operation(im))
return out
def inv_transform(self, image):
assert self._pads is not None,\
'Something went wrong, inv_transform(...) should be called after transform(...)'
return self._remove_padding(image)
def _get_dim_padding(self, dim_size):
pad = (self.divisor - dim_size % self.divisor) % self.divisor
pad_upper = pad // 2
pad_lower = pad - pad_upper
return pad_upper, pad_lower
def _remove_padding(self, tensors):
tensor_h, tensor_w = tensors[0].shape[-2:]
out = []
for t in tensors:
out.append(t[..., self._pads[2]:tensor_h - self._pads[3], self._pads[0]:tensor_w - self._pads[1]])
return out