# -*- coding: utf-8 -*-
# Copyright (c) XiMing Xing. All rights reserved.
# Author: XiMing Xing
# Description:
import math
import copy
import random
import pathlib
from typing import Dict
from shapely.geometry.polygon import Polygon
import omegaconf
import cv2
import numpy as np
import pydiffvg
import torch
from torch.optim.lr_scheduler import LambdaLR
from pytorch_svgrender.diffvg_warp import DiffVGState
from pytorch_svgrender.libs.solver.optim import get_optimizer
class Painter(DiffVGState):
def __init__(
self,
diffvg_cfg: omegaconf.DictConfig,
style: str,
num_segments: int,
segment_init: str,
radius: int = 20,
canvas_size: int = 600,
n_grid: int = 32,
trainable_bg: bool = False,
stroke_width: int = 3,
path_svg=None,
device=None,
):
super().__init__(device, print_timing=diffvg_cfg.print_timing,
canvas_width=canvas_size, canvas_height=canvas_size)
self.style = style
self.num_segments = num_segments
self.segment_init = segment_init
self.radius = radius
"""pixelart params"""
self.n_grid = n_grid # divide the canvas into n grids
self.pixel_per_grid = self.canvas_width // self.n_grid
"""sketch params"""
self.stroke_width = stroke_width
"""iconography params"""
self.color_ref = None
self.path_svg = path_svg
self.optimize_flag = []
self.strokes_counter = 0 # counts the number of calls to "get_path"
# Background color
self.para_bg = torch.tensor([1., 1., 1.], requires_grad=trainable_bg, device=self.device)
self.target_img = None
self.pos_init_method = None
def component_wise_path_init(self, gt, pred, init_type: str = 'sparse'):
# set target image
self.target_img = gt
if init_type == 'random':
self.pos_init_method = RandomCoordInit(self.canvas_height, self.canvas_width)
elif init_type == 'sparse':
# when initialized for the first time, the render result is None
if pred is None:
pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width)
# then pred is the render result
self.pos_init_method = SparseCoordInit(pred, gt)
elif init_type == 'naive':
if pred is None:
pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width)
self.pos_init_method = NaiveCoordInit(pred, gt)
else:
raise NotImplementedError(f"'{init_type}' is not support.")
def init_image(self, stage=0, num_paths=0):
self.cur_shapes, self.cur_shape_groups = [], []
# or init svg by pydiffvg
if self.style in ['pixelart', 'low-poly']: # update path definition
num_paths = self.n_grid
if stage > 0:
# Noting: if multi stages training than add new strokes on existing ones
# don't optimize on previous strokes
self.optimize_flag = [False for i in range(len(self.shapes))]
for i in range(num_paths):
if self.style == 'iconography':
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
fill_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
fill_color_init[-1] = 1.0
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([self.strokes_counter - 1]),
fill_color=fill_color_init,
stroke_color=None
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
self.optimize_flag.append(True)
elif self.style in ['pixelart', 'low-poly']:
for j in range(num_paths):
path = self.get_path(coord=[i, j])
self.shapes.append(path)
self.cur_shapes.append(path)
fill_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
fill_color_init[-1] = 1.0
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.LongTensor([i * num_paths + j]),
fill_color=fill_color_init,
stroke_color=None,
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
self.optimize_flag.append(True)
elif self.style in ['ink', 'sketch']:
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
stroke_color_init = [0.0, 0.0, 0.0] + [random.random()]
stroke_color_init = torch.FloatTensor(stroke_color_init)
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(self.shapes) - 1]),
fill_color=None,
stroke_color=stroke_color_init
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
elif self.style == 'painting':
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
wref, href = self.color_ref
wref = max(0, min(int(wref), self.canvas_width - 1))
href = max(0, min(int(href), self.canvas_height - 1))
stroke_color_init = list(self.target_img[0, :, href, wref]) + [1.]
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(self.shapes) - 1]),
fill_color=None,
stroke_color=torch.FloatTensor(stroke_color_init)
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
else:
num_paths_exists = 0
if self.path_svg is not None and pathlib.Path(self.path_svg).exists():
print(f"-> init svg from `{self.path_svg}` ...")
self.canvas_width, self.canvas_height, self.shapes, self.shape_groups = self.load_svg(self.path_svg)
# if you want to add more strokes to existing ones and optimize on all of them
num_paths_exists = len(self.shapes)
self.cur_shapes = self.shapes
self.cur_shape_groups = self.shape_groups
for i in range(num_paths_exists, num_paths):
if self.style == 'iconography':
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
wref, href = self.color_ref
wref = max(0, min(int(wref), self.canvas_width - 1))
href = max(0, min(int(href), self.canvas_height - 1))
fill_color_init = list(self.target_img[0, :, href, wref]) + [1.]
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([self.strokes_counter - 1]),
fill_color=torch.FloatTensor(fill_color_init),
stroke_color=None
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
elif self.style in ['pixelart', 'low-poly']:
for j in range(num_paths):
path = self.get_path(coord=[i, j])
self.shapes.append(path)
self.cur_shapes.append(path)
fill_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
fill_color_init[-1] = 1.0
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.LongTensor([i * num_paths + j]),
fill_color=fill_color_init,
stroke_color=None,
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
elif self.style in ['sketch', 'ink']:
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
stroke_color_init = [0.0, 0.0, 0.0] + [random.random()]
stroke_color_init = torch.FloatTensor(stroke_color_init)
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(self.shapes) - 1]),
fill_color=None,
stroke_color=stroke_color_init
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
elif self.style in ['painting']:
path = self.get_path()
self.shapes.append(path)
self.cur_shapes.append(path)
if self.color_ref is None:
stroke_color_val = np.random.uniform(size=[4])
stroke_color_val[-1] = 1.0
stroke_color_init = torch.FloatTensor(stroke_color_val)
else:
wref, href = self.color_ref
wref = max(0, min(int(wref), self.canvas_width - 1))
href = max(0, min(int(href), self.canvas_height - 1))
stroke_color_init = list(self.target_img[0, :, href, wref]) + [1.]
stroke_color_init = torch.FloatTensor(stroke_color_init)
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(self.shapes) - 1]),
fill_color=None,
stroke_color=stroke_color_init
)
self.shape_groups.append(path_group)
self.cur_shape_groups.append(path_group)
self.optimize_flag = [True for i in range(len(self.shapes))]
img = self.get_image()
return img
def get_image(self, step: int = 0):
img = self.render_warp(step)
img = img[:, :, 3:4] * img[:, :, :3] + self.para_bg * (1 - img[:, :, 3:4])
img = img.unsqueeze(0) # convert img from HWC to NCHW
img = img.permute(0, 3, 1, 2).to(self.device) # NHWC -> NCHW
return img
def get_path(self, coord=None):
num_segments = self.num_segments
points = []
if self.style == 'iconography':
# init segment
if self.segment_init == 'circle':
num_control_points = [2] * num_segments
radius = self.radius if self.radius is not None else np.random.uniform(0.5, 1)
if self.pos_init_method is not None:
center = self.pos_init_method()
else:
center = (random.random(), random.random())
bias = center
self.color_ref = copy.deepcopy(bias)
avg_degree = 360 / (num_segments * 3)
for i in range(0, num_segments * 3):
point = (
np.cos(np.deg2rad(i * avg_degree)), np.sin(np.deg2rad(i * avg_degree))
)
points.append(point)
points = torch.FloatTensor(points) * radius + torch.FloatTensor(bias).unsqueeze(dim=0)
elif self.segment_init == 'random':
num_control_points = [2] * num_segments
p0 = self.pos_init_method()
self.color_ref = copy.deepcopy(p0)
points.append(p0)
for j in range(num_segments):
radius = self.radius
p1 = (p0[0] + radius * np.random.uniform(-0.5, 0.5),
p0[1] + radius * np.random.uniform(-0.5, 0.5))
p2 = (p1[0] + radius * np.random.uniform(-0.5, 0.5),
p1[1] + radius * np.random.uniform(-0.5, 0.5))
p3 = (p2[0] + radius * np.random.uniform(-0.5, 0.5),
p2[1] + radius * np.random.uniform(-0.5, 0.5))
points.append(p1)
points.append(p2)
if j < num_segments - 1:
points.append(p3)
p0 = p3
points = torch.FloatTensor(points)
else:
raise NotImplementedError(f"{self.segment_init} is not exists.")
path = pydiffvg.Path(
num_control_points=torch.LongTensor(num_control_points),
points=points,
stroke_width=torch.tensor(0.0),
is_closed=True
)
elif self.style in ['sketch', 'painting', 'ink']:
num_control_points = torch.zeros(num_segments, dtype=torch.long) + 2
points = []
p0 = [random.random(), random.random()]
points.append(p0)
# select color by first point coordinate
color_ref = copy.deepcopy(p0)
color_ref[0] *= self.canvas_width
color_ref[1] *= self.canvas_height
self.color_ref = color_ref
for j in range(num_segments):
radius = 0.1
p1 = (p0[0] + radius * (random.random() - 0.5), p0[1] + radius * (random.random() - 0.5))
p2 = (p1[0] + radius * (random.random() - 0.5), p1[1] + radius * (random.random() - 0.5))
p3 = (p2[0] + radius * (random.random() - 0.5), p2[1] + radius * (random.random() - 0.5))
points.append(p1)
points.append(p2)
points.append(p3)
p0 = p3
points = torch.tensor(points).to(self.device)
points[:, 0] *= self.canvas_width
points[:, 1] *= self.canvas_height
path = pydiffvg.Path(num_control_points=torch.LongTensor(num_control_points),
points=points,
stroke_width=torch.tensor(float(self.stroke_width)),
is_closed=False)
elif self.style in ['pixelart', 'low-poly']:
x = coord[0] * self.pixel_per_grid
y = coord[1] * self.pixel_per_grid
points = torch.FloatTensor([
[x, y],
[x + self.pixel_per_grid, y],
[x + self.pixel_per_grid, y + self.pixel_per_grid],
[x, y + self.pixel_per_grid]
]).to(self.device)
path = pydiffvg.Polygon(points=points,
stroke_width=torch.tensor(0.0),
is_closed=True)
self.strokes_counter += 1
return path
def clip_curve_shape(self):
if self.style in ['sketch', 'ink']:
for group in self.shape_groups:
group.stroke_color.data[:3].clamp_(0., 0.) # to force black stroke
group.stroke_color.data[-1].clamp_(0., 1.) # clip alpha
else:
for group in self.shape_groups:
if group.stroke_color is not None:
group.stroke_color.data.clamp_(0.0, 1.0) # clip rgba
if group.fill_color is not None:
group.fill_color.data.clamp_(0.0, 1.0) # clip rgba
def reinitialize_paths(self,
reinit_path: bool = False,
opacity_threshold: float = None,
area_threshold: float = None,
fpath: pathlib.Path = None):
"""
reinitialize paths, also known as 'Reinitializing paths' in VectorFusion paper.
Args:
reinit_path: whether to reinitialize paths or not.
opacity_threshold: Threshold of opacity.
area_threshold: Threshold of the closed polygon area.
fpath: The path to save the reinitialized SVG.
"""
if not reinit_path:
return
if self.style not in ['iconography', 'low-poly', 'painting']:
return
def get_keys_below_threshold(my_dict, threshold):
keys_below_threshold = [key for key, value in my_dict.items() if value < threshold]
return keys_below_threshold
select_path_ids_by_opc = []
select_path_ids_by_area = []
if self.style in ['iconography', 'low-poly']:
# re-init by opacity_threshold
if opacity_threshold != 0 and opacity_threshold is not None:
opacity_record_ = {group.shape_ids.item(): group.fill_color[-1].item()
for group in self.cur_shape_groups}
# print("-> opacity_record: ", opacity_record_)
print("-> opacity_record: ", [f"{k}: {v:.3f}" for k, v in opacity_record_.items()])
select_path_ids_by_opc = get_keys_below_threshold(opacity_record_, opacity_threshold)
print("select_path_ids_by_opc: ", select_path_ids_by_opc)
# remove path by area_threshold
if area_threshold != 0 and area_threshold is not None:
area_records = [Polygon(shape.points.detach().cpu().numpy()).area for shape in self.cur_shapes]
# print("-> area_records: ", area_records)
print("-> area_records: ", ['%.2f' % i for i in area_records])
for i, shape in enumerate(self.cur_shapes):
points_ = shape.points.detach().cpu().numpy()
if Polygon(points_).area < area_threshold:
select_path_ids_by_area.append(shape.id)
print("select_path_ids_by_area: ", select_path_ids_by_area)
elif self.style in ['painting']:
# re-init by opacity_threshold
if opacity_threshold != 0 and opacity_threshold is not None:
opacity_record_ = {group.shape_ids.item(): group.stroke_color[-1].item()
for group in self.cur_shape_groups}
# print("-> opacity_record: ", opacity_record_)
print("-> opacity_record: ", [f"{k}: {v:.3f}" for k, v in opacity_record_.items()])
select_path_ids_by_opc = get_keys_below_threshold(opacity_record_, opacity_threshold)
print("select_path_ids_by_opc: ", select_path_ids_by_opc)
# re-init paths
reinit_union = list(set(select_path_ids_by_opc + select_path_ids_by_area))
if len(reinit_union) > 0:
for i, path in enumerate(self.cur_shapes):
if path.id in reinit_union:
coord = [i, i] if self.style == 'low-poly' else None
self.cur_shapes[i] = self.get_path(coord=coord)
for i, group in enumerate(self.cur_shape_groups):
shp_ids = group.shape_ids.cpu().numpy().tolist()
if set(shp_ids).issubset(reinit_union):
if self.style in ['iconography', 'low-poly']:
fill_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
fill_color_init[-1] = 1.0
self.cur_shape_groups[i] = pydiffvg.ShapeGroup(
shape_ids=torch.tensor(list(shp_ids)),
fill_color=fill_color_init,
stroke_color=None)
elif self.style in ['painting']:
stroke_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
stroke_color_init[-1] = 1.0
self.cur_shape_groups[i] = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(self.shapes) - 1]),
fill_color=None,
stroke_color=stroke_color_init
)
# save reinit svg
self.pretty_save_svg(fpath)
print("-" * 40)
def calc_distance_weight(self, loss_weight_keep):
shapes_forsdf = copy.deepcopy(self.cur_shapes)
shape_groups_forsdf = copy.deepcopy(self.cur_shape_groups)
for si in shapes_forsdf:
si.stroke_width = torch.FloatTensor([0]).to(self.device)
for sg_idx, sgi in enumerate(shape_groups_forsdf):
sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(self.device)
sgi.shape_ids = torch.LongTensor([sg_idx]).to(self.device)
sargs_forsdf = pydiffvg.RenderFunction.serialize_scene(
self.canvas_width, self.canvas_height, shapes_forsdf, shape_groups_forsdf
)
_render = pydiffvg.RenderFunction.apply
with torch.no_grad():
im_forsdf = _render(self.canvas_width, # width
self.canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None,
*sargs_forsdf)
# use alpha channel is a trick to get 0-1 image
im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy()
loss_weight = get_sdf(im_forsdf, normalize='to1')
loss_weight += loss_weight_keep
loss_weight = np.clip(loss_weight, 0, 1)
loss_weight = torch.FloatTensor(loss_weight).to(self.device)
return loss_weight
def set_point_parameters(self, id_delta=0):
self.point_vars = []
for i, path in enumerate(self.cur_shapes):
path.id = i + id_delta # set point id
path.points.requires_grad = True
self.point_vars.append(path.points)
def get_point_parameters(self):
return self.point_vars
def set_color_parameters(self):
self.color_vars = []
for i, group in enumerate(self.cur_shape_groups):
if group.fill_color is not None:
group.fill_color.requires_grad = True
self.color_vars.append(group.fill_color)
if group.stroke_color is not None:
group.stroke_color.requires_grad = True
self.color_vars.append(group.stroke_color)
def get_color_parameters(self):
return self.color_vars
def set_width_parameters(self):
# stroke`s width optimization
self.width_vars = []
for i, path in enumerate(self.shapes):
path.stroke_width.requires_grad = True
self.width_vars.append(path.stroke_width)
def get_width_parameters(self):
return self.width_vars
def pretty_save_svg(self, filename, width=None, height=None, shapes=None, shape_groups=None):
width = self.canvas_width if width is None else width
height = self.canvas_height if height is None else height
shapes = self.shapes if shapes is None else shapes
shape_groups = self.shape_groups if shape_groups is None else shape_groups
self.save_svg(filename, width, height, shapes, shape_groups, use_gamma=False, background=None)
def load_svg(self, path_svg):
canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(path_svg)
return canvas_width, canvas_height, shapes, shape_groups
def get_sdf(phi, **kwargs):
import skfmm # local import
phi = (phi - 0.5) * 2
if (phi.max() <= 0) or (phi.min() >= 0):
return np.zeros(phi.shape).astype(np.float32)
sd = skfmm.distance(phi, dx=1)
flip_negative = kwargs.get('flip_negative', True)
if flip_negative:
sd = np.abs(sd)
truncate = kwargs.get('truncate', 10)
sd = np.clip(sd, -truncate, truncate)
# print(f"max sd value is: {sd.max()}")
zero2max = kwargs.get('zero2max', True)
if zero2max and flip_negative:
sd = sd.max() - sd
elif zero2max:
raise ValueError
normalize = kwargs.get('normalize', 'sum')
if normalize == 'sum':
sd /= sd.sum()
elif normalize == 'to1':
sd /= sd.max()
return sd
class SparseCoordInit:
def __init__(self, pred, gt, format='[bs x c x 2D]', quantile_interval=200, nodiff_thres=0.1):
if torch.is_tensor(pred):
pred = pred.detach().cpu().numpy()
if torch.is_tensor(gt):
gt = gt.detach().cpu().numpy()
if format == '[bs x c x 2D]':
self.map = ((pred[0] - gt[0]) ** 2).sum(0)
self.reference_gt = copy.deepcopy(np.transpose(gt[0], (1, 2, 0)))
elif format == ['[2D x c]']:
self.map = (np.abs(pred - gt)).sum(-1)
self.reference_gt = copy.deepcopy(gt[0])
else:
raise ValueError
# OptionA: Zero too small errors to avoid the error too small deadloop
self.map[self.map < nodiff_thres] = 0
quantile_interval = np.linspace(0., 1., quantile_interval)
quantized_interval = np.quantile(self.map, quantile_interval)
# remove redundant
quantized_interval = np.unique(quantized_interval)
quantized_interval = sorted(quantized_interval[1:-1])
self.map = np.digitize(self.map, quantized_interval, right=False)
self.map = np.clip(self.map, 0, 255).astype(np.uint8)
self.idcnt = {}
for idi in sorted(np.unique(self.map)):
self.idcnt[idi] = (self.map == idi).sum()
# remove smallest one to remove the correct region
self.idcnt.pop(min(self.idcnt.keys()))
def __call__(self):
if len(self.idcnt) == 0:
h, w = self.map.shape
return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h]
target_id = max(self.idcnt, key=self.idcnt.get)
_, component, cstats, ccenter = cv2.connectedComponentsWithStats(
(self.map == target_id).astype(np.uint8),
connectivity=4
)
# remove cid = 0, it is the invalid area
csize = [ci[-1] for ci in cstats[1:]]
target_cid = csize.index(max(csize)) + 1
center = ccenter[target_cid][::-1]
coord = np.stack(np.where(component == target_cid)).T
dist = np.linalg.norm(coord - center, axis=1)
target_coord_id = np.argmin(dist)
coord_h, coord_w = coord[target_coord_id]
# replace_sampling
self.idcnt[target_id] -= max(csize)
if self.idcnt[target_id] == 0:
self.idcnt.pop(target_id)
self.map[component == target_cid] = 0
return [coord_w, coord_h]
class RandomCoordInit:
def __init__(self, canvas_width, canvas_height):
self.canvas_width, self.canvas_height = canvas_width, canvas_height
def __call__(self):
w, h = self.canvas_width, self.canvas_height
return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h]
class NaiveCoordInit:
def __init__(self, pred, gt, format='[bs x c x 2D]', replace_sampling=True):
if isinstance(pred, torch.Tensor):
pred = pred.detach().cpu().numpy()
if isinstance(gt, torch.Tensor):
gt = gt.detach().cpu().numpy()
if format == '[bs x c x 2D]':
self.map = ((pred[0] - gt[0]) ** 2).sum(0)
elif format == ['[2D x c]']:
self.map = ((pred - gt) ** 2).sum(-1)
else:
raise ValueError
self.replace_sampling = replace_sampling
def __call__(self):
coord = np.where(self.map == self.map.max())
coord_h, coord_w = coord[0][0], coord[1][0]
if self.replace_sampling:
self.map[coord_h, coord_w] = -1
return [coord_w, coord_h]
class PainterOptimizer:
def __init__(self,
renderer: Painter,
style: str,
num_iter: int,
lr_config: omegaconf.DictConfig,
trainable_bg: bool = False):
self.renderer = renderer
self.num_iter = num_iter
self.trainable_bg = trainable_bg
self.lr_config = lr_config
# set optimized params via style
self.optim_point, self.optim_color, self.optim_width = {
"iconography": (True, True, False),
"pixelart": (False, True, False),
"low-poly": (True, True, False),
"sketch": (True, True, False),
"ink": (True, True, True),
"painting": (True, True, True)
}.get(style, (False, False, False))
self.optim_bg = trainable_bg
# set lr schedule
schedule_cfg = lr_config.schedule
if schedule_cfg.name == 'linear':
self.lr_lambda = LinearDecayWithKeepLRLambda(init_lr=lr_config.point,
keep_ratio=schedule_cfg.keep_ratio,
decay_every=self.num_iter,
decay_ratio=schedule_cfg.decay_ratio)
elif schedule_cfg.name == 'cosine':
self.lr_lambda = CosineWithWarmupLRLambda(num_steps=self.num_iter,
warmup_steps=schedule_cfg.warmup_steps,
warmup_start_lr=schedule_cfg.warmup_start_lr,
warmup_end_lr=schedule_cfg.warmup_end_lr,
cosine_end_lr=schedule_cfg.cosine_end_lr)
else:
print(f"{schedule_cfg.name} is not support.")
self.lr_lambda = None
self.point_optimizer = None
self.color_optimizer = None
self.width_optimizer = None
self.bg_optimizer = None
self.point_scheduler = None
def init_optimizers(self, pid_delta: int = 0):
# optimizer
optim_cfg = self.lr_config.optim
optim_name = optim_cfg.name
params = {}
if self.optim_point:
self.renderer.set_point_parameters(pid_delta)
params['point'] = self.renderer.get_point_parameters()
self.point_optimizer = get_optimizer(optim_name, params['point'], self.lr_config.point, optim_cfg)
if self.optim_color:
self.renderer.set_color_parameters()
params['color'] = self.renderer.get_color_parameters()
self.color_optimizer = get_optimizer(optim_name, params['color'], self.lr_config.color, optim_cfg)
if self.optim_width:
self.renderer.set_width_parameters()
params['width'] = self.renderer.get_width_parameters()
if len(params['width']) > 0:
self.width_optimizer = get_optimizer(optim_name, params['width'], self.lr_config.width, optim_cfg)
if self.optim_bg:
self.renderer.para_bg.requires_grad = True
self.bg_optimizer = get_optimizer(optim_name, self.renderer.para_bg, self.lr_config.bg, optim_cfg)
# lr schedule
if self.lr_lambda is not None and self.optim_point:
self.point_scheduler = LambdaLR(self.point_optimizer, lr_lambda=self.lr_lambda, last_epoch=-1)
def update_lr(self):
if self.point_scheduler is not None:
self.point_scheduler.step()
def zero_grad_(self):
if self.point_optimizer is not None:
self.point_optimizer.zero_grad()
if self.color_optimizer is not None:
self.color_optimizer.zero_grad()
if self.width_optimizer is not None:
self.width_optimizer.zero_grad()
if self.bg_optimizer is not None:
self.bg_optimizer.zero_grad()
def step_(self):
if self.point_optimizer is not None:
self.point_optimizer.step()
if self.color_optimizer is not None:
self.color_optimizer.step()
if self.width_optimizer is not None:
self.width_optimizer.step()
if self.bg_optimizer is not None:
self.bg_optimizer.step()
def get_lr(self) -> Dict:
lr = {}
if self.point_optimizer is not None:
lr['pnt'] = self.point_optimizer.param_groups[0]['lr']
if self.color_optimizer is not None:
lr['clr'] = self.color_optimizer.param_groups[0]['lr']
if self.width_optimizer is not None:
lr['wd'] = self.width_optimizer.param_groups[0]['lr']
if self.bg_optimizer is not None:
lr['bg'] = self.bg_optimizer.param_groups[0]['lr']
return lr
class LinearDecayWithKeepLRLambda:
"""apply in LIVE stage"""
def __init__(self, init_lr, keep_ratio, decay_every, decay_ratio):
self.init_lr = init_lr
self.keep_ratio = keep_ratio
self.decay_every = decay_every
self.decay_ratio = decay_ratio
def __call__(self, n):
if n < self.keep_ratio * self.decay_every:
return self.init_lr
decay_time = n // self.decay_every
decay_step = n % self.decay_every
lr_s = self.decay_ratio ** decay_time
lr_e = self.decay_ratio ** (decay_time + 1)
r = decay_step / self.decay_every
lr = lr_s * (1 - r) + lr_e * r
return lr
class CosineWithWarmupLRLambda:
"""apply in fine-tuning stage"""
def __init__(self, num_steps, warmup_steps, warmup_start_lr, warmup_end_lr, cosine_end_lr):
self.n_steps = num_steps
self.n_warmup = warmup_steps
self.warmup_start_lr = warmup_start_lr
self.warmup_end_lr = warmup_end_lr
self.cosine_end_lr = cosine_end_lr
def __call__(self, n):
if n < self.n_warmup:
# linearly warmup
return self.warmup_start_lr + (n / self.n_warmup) * (self.warmup_end_lr - self.warmup_start_lr)
else:
# cosine decayed schedule
return self.cosine_end_lr + 0.5 * (self.warmup_end_lr - self.cosine_end_lr) * (
1 + math.cos(math.pi * (n - self.n_warmup) / (self.n_steps - self.n_warmup)))