pytorch_svgrender/painter/vectorfusion/painter_params.py

# -*- 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)))