text: import os
  import cv2
  import random
  
  import utils
  import loss
  from networks import *
  import morphology
  
  import renderer
  
  import torch
  
  # Decide which device we want to run on
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  
  class PainterBase():
      def __init__(self, args):
          self.args = args
  
          self.rderr = renderer.Renderer(renderer=args.renderer,
                                         CANVAS_WIDTH=args.canvas_size, canvas_color=args.canvas_color)
  
          # define G
          self.net_G = define_G(rdrr=self.rderr, netG=args.net_G).to(device)
  
          # define some other vars to record the training states
          self.x_ctt = None
          self.x_color = None
          self.x_alpha = None
  
          self.G_pred_foreground = None
          self.G_pred_alpha = None
          self.G_final_pred_canvas = torch.zeros(
              [1, 3, self.net_G.out_size, self.net_G.out_size]).to(device)
  
          self.G_loss = torch.tensor(0.0)
          self.step_id = 0
          self.anchor_id = 0
          self.renderer_checkpoint_dir = args.renderer_checkpoint_dir
          self.output_dir = args.output_dir
          self.lr = args.lr
  
          # define the loss functions
          self._pxl_loss = loss.PixelLoss(p=1)
          self._sinkhorn_loss = loss.SinkhornLoss(epsilon=0.01, niter=5, normalize=False)
  
          # some other vars to be initialized in child classes
          self.input_aspect_ratio = None
          self.img_path = None
          self.img_batch = None
          self.img_ = None
          self.final_rendered_images = None
          self.m_grid = None
          self.m_strokes_per_block = None
  
          if os.path.exists(self.output_dir) is False:
              os.mkdir(self.output_dir)
  
      def _load_checkpoint(self):
  
          # load renderer G
          if os.path.exists((os.path.join(
                  self.renderer_checkpoint_dir, 'last_ckpt.pt'))):
              print('loading renderer from pre-trained checkpoint...')
              # load the entire checkpoint
              checkpoint = torch.load(os.path.join(self.renderer_checkpoint_dir, 'last_ckpt.pt'),
                                  map_location=None if torch.cuda.is_available() else device)
              # update net_G states
              self.net_G.load_state_dict(checkpoint['model_G_state_dict'])
              self.net_G.to(device)
              self.net_G.eval()
          else:
              print('pre-trained renderer does not exist...')
              exit()
  
  
      def _compute_acc(self):
  
          target = self.img_batch.detach()
          canvas = self.G_pred_canvas.detach()
          psnr = utils.cpt_batch_psnr(canvas, target, PIXEL_MAX=1.0)
  
          return psnr
  
      def _save_stroke_params(self, v):
  
          d_shape = self.rderr.d_shape
          d_color = self.rderr.d_color
          d_alpha = self.rderr.d_alpha
  
          x_ctt = v[:, :, 0:d_shape]
          x_color = v[:, :, d_shape:d_shape+d_color]
          x_alpha = v[:, :, d_shape+d_color:d_shape+d_color+d_alpha]
          print('saving stroke parameters...')
          file_name = os.path.join(
              self.output_dir, self.img_path.split('/')[-1][:-4])
          np.savez(file_name + '_strokes.npz', x_ctt=x_ctt,
                   x_color=x_color, x_alpha=x_alpha)
  
      def _shuffle_strokes_and_reshape(self, v):
  
          grid_idx = list(range(self.m_grid ** 2))
          random.shuffle(grid_idx)
          v = v[grid_idx, :, :]
          v = np.reshape(np.transpose(v, [1,0,2]), [-1, self.rderr.d])
          v = np.expand_dims(v, axis=0)
  
          return v
  
      def _render(self, v, save_jpgs=True, save_video=True):
  
          v = v[0,:,:]
          if self.args.keep_aspect_ratio:
              if self.input_aspect_ratio < 1:
                  out_h = int(self.args.canvas_size * self.input_aspect_ratio)
                  out_w = self.args.canvas_size
              else:
                  out_h = self.args.canvas_size
                  out_w = int(self.args.canvas_size / self.input_aspect_ratio)
          else:
              out_h = self.args.canvas_size
              out_w = self.args.canvas_size
  
          file_name = os.path.join(
              self.output_dir, self.img_path.split('/')[-1][:-4])
  
          if save_video:
              video_writer = cv2.VideoWriter(
                  file_name + '_animated.mp4', cv2.VideoWriter_fourcc(*'MP4V'), 40,
                  (out_w, out_h))
  
          print('rendering canvas...')
          self.rderr.create_empty_canvas()
          for i in range(v.shape[0]):  # for each stroke
              self.rderr.stroke_params = v[i, :]
              if self.rderr.check_stroke():
                  self.rderr.draw_stroke()
              this_frame = self.rderr.canvas
              this_frame = cv2.resize(this_frame, (out_w, out_h), cv2.INTER_AREA)
              if save_jpgs:
                  plt.imsave(file_name + '_rendered_stroke_' + str((i+1)).zfill(4) +
                             '.png', this_frame)
              if save_video:
                  video_writer.write((this_frame[:,:,::-1] * 255.).astype(np.uint8))
  
          if save_jpgs:
              print('saving input photo...')
              out_img = cv2.resize(self.img_, (out_w, out_h), cv2.INTER_AREA)
              plt.imsave(file_name + '_input.png', out_img)
  
          final_rendered_image = np.copy(this_frame)
          if save_jpgs:
              print('saving final rendered result...')
              plt.imsave(file_name + '_final.png', final_rendered_image)
  
          return final_rendered_image
  
  
  
  
      def _normalize_strokes(self, v):
  
          v = np.array(v.detach().cpu())
  
          if self.rderr.renderer in ['watercolor', 'markerpen']:
              # x0, y0, x1, y1, x2, y2, radius0, radius2, ...
              xs = np.array([0, 4])
              ys = np.array([1, 5])
              rs = np.array([6, 7])
          elif self.rderr.renderer in ['oilpaintbrush', 'rectangle']:
              # xc, yc, w, h, theta ...
              xs = np.array([0])
              ys = np.array([1])
              rs = np.array([2, 3])
          else:
              raise NotImplementedError('renderer [%s] is not implemented' % self.rderr.renderer)
  
          for y_id in range(self.m_grid):
              for x_id in range(self.m_grid):
                  y_bias = y_id / self.m_grid
                  x_bias = x_id / self.m_grid
                  v[y_id * self.m_grid + x_id, :, ys] = \
                      y_bias + v[y_id * self.m_grid + x_id, :, ys] / self.m_grid
                  v[y_id * self.m_grid + x_id, :, xs] = \
                      x_bias + v[y_id * self.m_grid + x_id, :, xs] / self.m_grid
                  v[y_id * self.m_grid + x_id, :, rs] /= self.m_grid
  
          return v
  
  
      def initialize_params(self):
  
          self.x_ctt = np.random.rand(
              self.m_grid*self.m_grid, self.m_strokes_per_block,
              self.rderr.d_shape).astype(np.float32)
          self.x_ctt = torch.tensor(self.x_ctt).to(device)
  
          self.x_color = np.random.rand(
              self.m_grid*self.m_grid, self.m_strokes_per_block,
              self.rderr.d_color).astype(np.float32)
          self.x_color = torch.tensor(self.x_color).to(device)
  
          self.x_alpha = np.random.rand(
              self.m_grid*self.m_grid, self.m_strokes_per_block,
              self.rderr.d_alpha).astype(np.float32)
          self.x_alpha = torch.tensor(self.x_alpha).to(device)
  
  
      def stroke_sampler(self, anchor_id):
  
          if anchor_id == self.m_strokes_per_block:
              return
  
          err_maps = torch.sum(
              torch.abs(self.img_batch - self.G_final_pred_canvas),
              dim=1, keepdim=True).detach()
  
          for i in range(self.m_grid*self.m_grid):
              this_err_map = err_maps[i,0,:,:].cpu().numpy()
              ks = int(this_err_map.shape[0] / 8)
              this_err_map = cv2.blur(this_err_map, (ks, ks))
              this_err_map = this_err_map ** 4
              this_img = self.img_batch[i, :, :, :].detach().permute([1, 2, 0]).cpu().numpy()
  
              self.rderr.random_stroke_params_sampler(
                  err_map=this_err_map, img=this_img)
  
              self.x_ctt.data[i, anchor_id, :] = torch.tensor(
                  self.rderr.stroke_params[0:self.rderr.d_shape])
              self.x_color.data[i, anchor_id, :] = torch.tensor(
                  self.rderr.stroke_params[self.rderr.d_shape:self.rderr.d_shape+self.rderr.d_color])
              self.x_alpha.data[i, anchor_id, :] = torch.tensor(self.rderr.stroke_params[-1])
  
  
      def _backward_x(self):
  
          self.G_loss = 0
          self.G_loss += self.args.beta_L1 * self._pxl_loss(
              canvas=self.G_final_pred_canvas, gt=self.img_batch)
          if self.args.with_ot_loss:
              self.G_loss += self.args.beta_ot * self._sinkhorn_loss(
                  self.G_final_pred_canvas, self.img_batch)
          self.G_loss.backward()
  
  
      def _forward_pass(self):
  
          self.x = torch.cat([self.x_ctt, self.x_color, self.x_alpha], dim=-1)
  
          v = torch.reshape(self.x[:, 0:self.anchor_id+1, :],
                            [self.m_grid*self.m_grid*(self.anchor_id+1), -1, 1, 1])
          self.G_pred_foregrounds, self.G_pred_alphas = self.net_G(v)
  
          self.G_pred_foregrounds = morphology.Dilation2d(m=1)(self.G_pred_foregrounds)
          self.G_pred_alphas = morphology.Erosion2d(m=1)(self.G_pred_alphas)
  
          self.G_pred_foregrounds = torch.reshape(
              self.G_pred_foregrounds, [self.m_grid*self.m_grid, self.anchor_id+1, 3,
                                        self.net_G.out_size, self.net_G.out_size])
          self.G_pred_alphas = torch.reshape(
              self.G_pred_alphas, [self.m_grid*self.m_grid, self.anchor_id+1, 3,
                                   self.net_G.out_size, self.net_G.out_size])
  
          for i in range(self.anchor_id+1):
              G_pred_foreground = self.G_pred_foregrounds[:, i]
              G_pred_alpha = self.G_pred_alphas[:, i]
              self.G_pred_canvas = G_pred_foreground * G_pred_alpha \
                                   + self.G_pred_canvas * (1 - G_pred_alpha)
  
          self.G_final_pred_canvas = self.G_pred_canvas
  
  
  
  
  class Painter(PainterBase):
  
      def __init__(self, args):
          super(Painter, self).__init__(args=args)
  
          self.m_grid = args.m_grid
  
          self.max_m_strokes = args.max_m_strokes
  
          self.img_path = args.img_path
          self.img_ = cv2.imread(args.img_path, cv2.IMREAD_COLOR)
          self.img_ = cv2.cvtColor(self.img_, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
          self.input_aspect_ratio = self.img_.shape[0] / self.img_.shape[1]
          self.img_ = cv2.resize(self.img_, (self.net_G.out_size * args.m_grid,
                                             self.net_G.out_size * args.m_grid), cv2.INTER_AREA)
  
          self.m_strokes_per_block = int(args.max_m_strokes / (args.m_grid * args.m_grid))
  
          self.img_batch = utils.img2patches(self.img_, args.m_grid, self.net_G.out_size).to(device)
  
          self.final_rendered_images = None
  
  
      def _drawing_step_states(self):
          acc = self._compute_acc().item()
          print('iteration step %d, G_loss: %.5f, step_psnr: %.5f, strokes: %d / %d'
                % (self.step_id, self.G_loss.item(), acc,
                   (self.anchor_id+1)*self.m_grid*self.m_grid,
                   self.max_m_strokes))
          vis2 = utils.patches2img(self.G_final_pred_canvas, self.m_grid).clip(min=0, max=1)
          if self.args.disable_preview:
              pass
          else:
              cv2.namedWindow('G_pred', cv2.WINDOW_NORMAL)
              cv2.namedWindow('input', cv2.WINDOW_NORMAL)
              cv2.imshow('G_pred', vis2[:,:,::-1])
              cv2.imshow('input', self.img_[:, :, ::-1])
              cv2.waitKey(1)
  
  
  
  
  
  class ProgressivePainter(PainterBase):
  
      def __init__(self, args):
          super(ProgressivePainter, self).__init__(args=args)
  
          self.max_divide = args.max_divide
  
          self.max_m_strokes = args.max_m_strokes
  
          self.m_strokes_per_block = self.stroke_parser()
  
          self.m_grid = 1
  
          self.img_path = args.img_path
          self.img_ = cv2.imread(args.img_path, cv2.IMREAD_COLOR)
          self.img_ = cv2.cvtColor(self.img_, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
          self.input_aspect_ratio = self.img_.shape[0] / self.img_.shape[1]
          self.img_ = cv2.resize(self.img_, (self.net_G.out_size * args.max_divide,
                                             self.net_G.out_size * args.max_divide), cv2.INTER_AREA)
  
  
      def stroke_parser(self):
  
          total_blocks = 0
          for i in range(0, self.max_divide + 1):
              total_blocks += i ** 2
  
          return int(self.max_m_strokes / total_blocks)
  
  
      def _drawing_step_states(self):
          acc = self._compute_acc().item()
          # print('iteration step %d, G_loss: %.5f, step_acc: %.5f, grid_scale: %d / %d, strokes: %d / %d'
          #       % (self.step_id, self.G_loss.item(), acc,
          #          self.m_grid, self.max_divide,
          #          self.anchor_id + 1, self.m_strokes_per_block))
          vis2 = utils.patches2img(self.G_final_pred_canvas, self.m_grid).clip(min=0, max=1)
          if self.args.disable_preview:
              pass
          else:
              cv2.namedWindow('G_pred', cv2.WINDOW_NORMAL)
              cv2.namedWindow('input', cv2.WINDOW_NORMAL)
              cv2.imshow('G_pred', vis2[:,:,::-1])
              cv2.imshow('input', self.img_[:, :, ::-1])
              cv2.waitKey(1)
  
  
  
  class NeuralStyleTransfer(PainterBase):
  
      def __init__(self, args):
          super(NeuralStyleTransfer, self).__init__(args=args)
  
          self.args = args
  
          self._style_loss = loss.VGGStyleLoss(transfer_mode=args.transfer_mode, resize=True)
  
          print('loading pre-generated vector file...')
          if os.path.exists(args.vector_file) is False:
              exit('vector file does not exist, pls check --vector_file, or run demo.py fist')
          else:
              npzfile = np.load(args.vector_file)
  
          self.x_ctt = torch.tensor(npzfile['x_ctt']).to(device)
          self.x_color = torch.tensor(npzfile['x_color']).to(device)
          self.x_alpha = torch.tensor(npzfile['x_alpha']).to(device)
          self.m_grid = int(np.sqrt(self.x_ctt.shape[0]))
  
          self.anchor_id = self.x_ctt.shape[1] - 1
  
          img_ = cv2.imread(args.content_img_path, cv2.IMREAD_COLOR)
          img_ = cv2.cvtColor(img_, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
          self.input_aspect_ratio = img_.shape[0] / img_.shape[1]
          self.img_ = cv2.resize(img_, (self.net_G.out_size*self.m_grid,
                                        self.net_G.out_size*self.m_grid), cv2.INTER_AREA)
          self.img_batch = utils.img2patches(self.img_, self.m_grid, self.net_G.out_size).to(device)
  
          style_img = cv2.imread(args.style_img_path, cv2.IMREAD_COLOR)
          self.style_img_ = cv2.cvtColor(style_img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
          self.style_img = cv2.blur(cv2.resize(self.style_img_, (128, 128)), (2, 2))
          self.style_img = torch.tensor(self.style_img).permute([2, 0, 1]).unsqueeze(0).to(device)
  
          self.content_img_path = args.content_img_path
          self.img_path = args.content_img_path
          self.style_img_path = args.style_img_path
  
  
      def _style_transfer_step_states(self):
          acc = self._compute_acc().item()
          print('running style transfer... iteration step %d, G_loss: %.5f, step_psnr: %.5f'
                % (self.step_id, self.G_loss.item(), acc))
          vis2 = utils.patches2img(self.G_final_pred_canvas, self.m_grid).clip(min=0, max=1)
          if self.args.disable_preview:
              pass
          else:
              cv2.namedWindow('G_pred', cv2.WINDOW_NORMAL)
              cv2.namedWindow('input', cv2.WINDOW_NORMAL)
              cv2.namedWindow('style_img', cv2.WINDOW_NORMAL)
              cv2.imshow('G_pred', vis2[:,:,::-1])
              cv2.imshow('input', self.img_[:, :, ::-1])
              cv2.imshow('style_img', self.style_img_[:, :, ::-1])
              cv2.waitKey(1)
  
  
      def _backward_x_sty(self):
          canvas = utils.patches2img(
              self.G_final_pred_canvas, self.m_grid, to_numpy=False).to(device)
          self.G_loss = self.args.beta_L1 * self._pxl_loss(
              canvas=self.G_final_pred_canvas, gt=self.img_batch, ignore_color=True)
          self.G_loss += self.args.beta_sty * self._style_loss(canvas, self.style_img)
          self.G_loss.backward()
  
  
      def _render_on_grids(self, v):
  
          rendered_imgs = []
  
          self.rderr.create_empty_canvas()
  
          grid_idx = list(range(self.m_grid ** 2))
          random.shuffle(grid_idx)
          for j in range(v.shape[1]):  # for each group of stroke
              for i in range(len(grid_idx)):  # for each random patch
                  self.rderr.stroke_params = v[grid_idx[i], j, :]
                  if self.rderr.check_stroke():
                      self.rderr.draw_stroke()
                  rendered_imgs.append(self.rderr.canvas)
  
          return rendered_imgs
  
  
      def _save_style_transfer_images(self, final_rendered_image):
  
          if self.args.keep_aspect_ratio:
              if self.input_aspect_ratio < 1:
                  out_h = int(self.args.canvas_size * self.input_aspect_ratio)
                  out_w = self.args.canvas_size
              else:
                  out_h = self.args.canvas_size
                  out_w = int(self.args.canvas_size / self.input_aspect_ratio)
          else:
              out_h = self.args.canvas_size
              out_w = self.args.canvas_size
  
          print('saving style transfer results...')
  
          file_dir = os.path.join(
              self.output_dir, self.content_img_path.split('/')[-1][:-4])
  
          out_img = cv2.resize(self.style_img_, (out_w, out_h), cv2.INTER_AREA)
          plt.imsave(file_dir + '_style_img_' +
                     self.style_img_path.split('/')[-1][:-4] + '.png', out_img)
  
          out_img = cv2.resize(final_rendered_image, (out_w, out_h), cv2.INTER_AREA)
          plt.imsave(file_dir + '_style_transfer_' +
                     self.style_img_path.split('/')[-1][:-4] + '.png', out_img)
  