Exemple #1
0
    def __init__(self, args, device=DEVICE):
        self.args = args
        self.device = device

        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, 128, 128]).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)

        if os.path.exists(args.vector_file):
            npzfile = np.load(args.vector_file)
            print(npzfile['x_ctt'].shape)
            print(npzfile['x_color'].shape)
            print(npzfile['x_alpha'].shape)
            self.x_ctt = np.copy(npzfile['x_ctt'][:, ::-1, :])
            self.x_color = np.copy(npzfile['x_color'][:, ::-1, :])
            self.x_alpha = np.copy(npzfile['x_alpha'][:, ::-1, :])
            self.m_grid = int(np.sqrt(self.x_ctt.shape[0]))
            self.anchor_id = self.x_ctt.shape[1] - 1
Exemple #2
0
    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)