Initial commit

app.py

@@ -0,0 +1,176 @@
+import gradio as gr
+import os
+import cv2
+import torch
+import numpy as np
+import argparse
+import torch.nn as nn
+import torch.nn.functional as F
+
+from baseline.DRL.actor import *
+from baseline.Renderer.stroke_gen import *
+from baseline.Renderer.model import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+width = 128
+
+
+actor_path = 'ckpts/actor.pkl'
+renderer_path = 'ckpts/renderer.pkl'
+# 
+divide = 4
+canvas_cnt = divide * divide
+
+Decoder = FCN()
+Decoder.load_state_dict(torch.load(renderer_path))
+
+def decode(x, canvas): # b * (10 + 3)
+    x = x.view(-1, 10 + 3)
+    stroke = 1 - Decoder(x[:, :10])
+    stroke = stroke.view(-1, width, width, 1)
+    color_stroke = stroke * x[:, -3:].view(-1, 1, 1, 3)
+    stroke = stroke.permute(0, 3, 1, 2)
+    color_stroke = color_stroke.permute(0, 3, 1, 2)
+    stroke = stroke.view(-1, 5, 1, width, width)
+    color_stroke = color_stroke.view(-1, 5, 3, width, width)
+    res = []
+    for i in range(5):
+        canvas = canvas * (1 - stroke[:, i]) + color_stroke[:, i]
+        res.append(canvas)
+    return canvas, res
+
+def small2large(x):
+    # (d * d, width, width) -> (d * width, d * width)    
+    x = x.reshape(divide, divide, width, width, -1)
+    x = np.transpose(x, (0, 2, 1, 3, 4))
+    x = x.reshape(divide * width, divide * width, -1)
+    return x
+
+def large2small(x):
+    # (d * width, d * width) -> (d * d, width, width)
+    x = x.reshape(divide, width, divide, width, 3)
+    x = np.transpose(x, (0, 2, 1, 3, 4))
+    x = x.reshape(canvas_cnt, width, width, 3)
+    return x
+
+def smooth(img):
+    def smooth_pix(img, tx, ty):
+        if tx == divide * width - 1 or ty == divide * width - 1 or tx == 0 or ty == 0: 
+            return img
+        img[tx, ty] = (img[tx, ty] + img[tx + 1, ty] + img[tx, ty + 1] + img[tx - 1, ty] + img[tx, ty - 1] + img[tx + 1, ty - 1] + img[tx - 1, ty + 1] + img[tx - 1, ty - 1] + img[tx + 1, ty + 1]) / 9
+        return img
+
+    for p in range(divide):
+        for q in range(divide):
+            x = p * width
+            y = q * width
+            for k in range(width):
+                img = smooth_pix(img, x + k, y + width - 1)
+                if q != divide - 1:
+                    img = smooth_pix(img, x + k, y + width)
+            for k in range(width):
+                img = smooth_pix(img, x + width - 1, y + k)
+                if p != divide - 1:
+                    img = smooth_pix(img, x + width, y + k)
+    return img
+
+def save_img(res, imgid, origin_shape, output_name, divide=False):
+    output = res.detach().cpu().numpy() # d * d, 3, width, width    
+    output = np.transpose(output, (0, 2, 3, 1))
+    if divide:
+        output = small2large(output)
+        output = smooth(output)
+    else:
+        output = output[0]
+    output = (output * 255).astype('uint8')
+    output = cv2.resize(output, origin_shape)
+    cv2.imwrite(output_name +"/" + str(imgid) + '.jpg', output)
+
+actor = ResNet(9, 18, 65) # action_bundle = 5, 65 = 5 * 13
+actor.load_state_dict(torch.load(actor_path))
+actor = actor.to(device).eval()
+Decoder = Decoder.to(device).eval()
+
+
+
+def paint_img(img):
+    max_step = 40
+    # imgid = 0
+    # output_name = os.path.join('output', str(len(os.listdir('output'))) if os.path.exists('output') else '0')
+    # os.makedirs(output_name, exist_ok= True)
+    # img = cv2.imread(args.img, cv2.IMREAD_COLOR)
+    origin_shape = (img.shape[1], img.shape[0])
+    patch_img = cv2.resize(img, (width * divide, width * divide))
+    patch_img = large2small(patch_img)
+    patch_img = np.transpose(patch_img, (0, 3, 1, 2))
+    patch_img = torch.tensor(patch_img).to(device).float() / 255.
+
+    img = cv2.resize(img, (width, width))
+    img = img.reshape(1, width, width, 3)
+    img = np.transpose(img, (0, 3, 1, 2))
+    img = torch.tensor(img).to(device).float() / 255.
+
+    T = torch.ones([1, 1, width, width], dtype=torch.float32).to(device)
+    coord = torch.zeros([1, 2, width, width])
+    for i in range(width):
+        for j in range(width):
+            coord[0, 0, i, j] = i / (width - 1.)
+            coord[0, 1, i, j] = j / (width - 1.)
+    coord = coord.to(device) # Coordconv
+    canvas = torch.zeros([1, 3, width, width]).to(device)
+
+    with torch.no_grad():
+        if divide != 1:
+            max_step = max_step // 2
+        for i in range(max_step):
+            stepnum = T * i / max_step
+            actions = actor(torch.cat([canvas, img, stepnum, coord], 1))
+            canvas, res = decode(actions, canvas)
+            for j in range(5):
+                # save_img(res[j], imgid)
+                # imgid += 1
+                output = res[j].detach().cpu().numpy() # d * d, 3, width, width    
+                output = np.transpose(output, (0, 2, 3, 1))
+                output = output[0]
+                output = (output * 255).astype('uint8')
+                output = cv2.resize(output, origin_shape)
+                yield output
+        if divide != 1:
+            canvas = canvas[0].detach().cpu().numpy()
+            canvas = np.transpose(canvas, (1, 2, 0))    
+            canvas = cv2.resize(canvas, (width * divide, width * divide))
+            canvas = large2small(canvas)
+            canvas = np.transpose(canvas, (0, 3, 1, 2))
+            canvas = torch.tensor(canvas).to(device).float()
+            coord = coord.expand(canvas_cnt, 2, width, width)
+            T = T.expand(canvas_cnt, 1, width, width)
+            for i in range(max_step):
+                stepnum = T * i / max_step
+                actions = actor(torch.cat([canvas, patch_img, stepnum, coord], 1))
+                canvas, res = decode(actions, canvas)
+                # print('divided canvas step {}, L2Loss = {}'.format(i, ((canvas - patch_img) ** 2).mean()))
+                for j in range(5):
+                    # save_img(res[j], imgid, True)
+                    # imgid += 1
+                    output = res[j].detach().cpu().numpy() # d * d, 3, width, width    
+                    output = np.transpose(output, (0, 2, 3, 1))
+                    output = small2large(output)
+                    output = smooth(output)
+                    output = (output * 255).astype('uint8')
+                    output = cv2.resize(output, origin_shape)
+                    yield output
+
+        return output
+
+examples = [
+    ["image\chaoyue.png"],
+    ["image\degang.png"],
+    ["image\JayChou.png"],
+    ["image\Leslie.png"],
+    ["image\mayun.png"],
+
+]
+
+demo = gr.Interface(fn=paint_img, inputs=gr.Image(), outputs="image", examples = examples)
+demo.queue()
+demo.launch(server_name="0.0.0.0")

baseline/DRL/actor.py

@@ -0,0 +1,114 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.weight_norm as weightNorm
+
+from torch.autograd import Variable
+import sys
+
+def conv3x3(in_planes, out_planes, stride=1):
+    return (nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False))
+
+def cfg(depth):
+    depth_lst = [18, 34, 50, 101, 152]
+    assert (depth in depth_lst), "Error : Resnet depth should be either 18, 34, 50, 101, 152"
+    cf_dict = {
+        '18': (BasicBlock, [2,2,2,2]),
+        '34': (BasicBlock, [3,4,6,3]),
+        '50': (Bottleneck, [3,4,6,3]),
+        '101':(Bottleneck, [3,4,23,3]),
+        '152':(Bottleneck, [3,8,36,3]),
+    }
+
+    return cf_dict[str(depth)]
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                (nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = (nn.Conv2d(in_planes, planes, kernel_size=1, bias=False))
+        self.conv2 = (nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False))
+        self.conv3 = (nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False))
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                (nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+    def __init__(self, num_inputs, depth, num_outputs):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        block, num_blocks = cfg(depth)
+
+        self.conv1 = conv3x3(num_inputs, 64, 2)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.fc = nn.Linear(512 * block.expansion, num_outputs)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = F.avg_pool2d(x, 4)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        x = torch.sigmoid(x)
+        return x

baseline/DRL/critic.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.weight_norm as weightNorm
+
+from torch.autograd import Variable
+import sys
+
+def conv3x3(in_planes, out_planes, stride=1):
+    return weightNorm(nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=True))
+
+class TReLU(nn.Module):
+    def __init__(self):
+        super(TReLU, self).__init__()
+        self.alpha = nn.Parameter(torch.FloatTensor(1), requires_grad=True)
+        self.alpha.data.fill_(0)
+        
+    def forward(self, x):
+        x = F.relu(x - self.alpha) + self.alpha
+        return x
+
+def cfg(depth):
+    depth_lst = [18, 34, 50, 101, 152]
+    assert (depth in depth_lst), "Error : Resnet depth should be either 18, 34, 50, 101, 152"
+    cf_dict = {
+        '18': (BasicBlock, [2,2,2,2]),
+        '34': (BasicBlock, [3,4,6,3]),
+        '50': (Bottleneck, [3,4,6,3]),
+        '101':(Bottleneck, [3,4,23,3]),
+        '152':(Bottleneck, [3,8,36,3]),
+    }
+
+    return cf_dict[str(depth)]
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.conv2 = conv3x3(planes, planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                weightNorm(nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=True)),
+            )
+        self.relu_1 = TReLU()
+        self.relu_2 = TReLU()
+
+    def forward(self, x):
+        out = self.relu_1(self.conv1(x))
+        out = self.conv2(out)
+        out += self.shortcut(x)
+        out = self.relu_2(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = weightNorm(nn.Conv2d(in_planes, planes, kernel_size=1, bias=True))
+        self.conv2 = weightNorm(nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=True))
+        self.conv3 = weightNorm(nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=True))
+        self.relu_1 = TReLU()
+        self.relu_2 = TReLU()
+        self.relu_3 = TReLU()
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                weightNorm(nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=True)),
+            )
+
+    def forward(self, x):
+        out = self.relu_1(self.conv1(x))
+        out = self.relu_2(self.conv2(out))
+        out = self.conv3(out)
+        out += self.shortcut(x)
+        out = self.relu_3(out)
+
+        return out
+
+class ResNet_wobn(nn.Module):
+    def __init__(self, num_inputs, depth, num_outputs):
+        super(ResNet_wobn, self).__init__()
+        self.in_planes = 64
+
+        block, num_blocks = cfg(depth)
+
+        self.conv1 = conv3x3(num_inputs, 64, 2)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.fc = nn.Linear(512 * block.expansion, num_outputs)
+        self.relu_1 = TReLU()
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.relu_1(self.conv1(x))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = F.avg_pool2d(x, 4)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x

baseline/DRL/ddpg.py

@@ -0,0 +1,220 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import Adam, SGD
+from Renderer.model import *
+from DRL.rpm import rpm
+from DRL.actor import *
+from DRL.critic import *
+from DRL.wgan import *
+from utils.util import *
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+coord = torch.zeros([1, 2, 128, 128])
+for i in range(128):
+    for j in range(128):
+        coord[0, 0, i, j] = i / 127.
+        coord[0, 1, i, j] = j / 127.
+coord = coord.to(device)
+
+criterion = nn.MSELoss()
+
+Decoder = FCN()
+Decoder.load_state_dict(torch.load('../renderer.pkl'))
+
+def decode(x, canvas): # b * (10 + 3)
+    x = x.view(-1, 10 + 3)
+    stroke = 1 - Decoder(x[:, :10])
+    stroke = stroke.view(-1, 128, 128, 1)
+    color_stroke = stroke * x[:, -3:].view(-1, 1, 1, 3)
+    stroke = stroke.permute(0, 3, 1, 2)
+    color_stroke = color_stroke.permute(0, 3, 1, 2)
+    stroke = stroke.view(-1, 5, 1, 128, 128)
+    color_stroke = color_stroke.view(-1, 5, 3, 128, 128)
+    for i in range(5):
+        canvas = canvas * (1 - stroke[:, i]) + color_stroke[:, i]
+    return canvas
+
+def cal_trans(s, t):
+    return (s.transpose(0, 3) * t).transpose(0, 3)
+    
+class DDPG(object):
+    def __init__(self, batch_size=64, env_batch=1, max_step=40, \
+                 tau=0.001, discount=0.9, rmsize=800, \
+                 writer=None, resume=None, output_path=None):
+
+        self.max_step = max_step
+        self.env_batch = env_batch
+        self.batch_size = batch_size        
+
+        self.actor = ResNet(9, 18, 65) # target, canvas, stepnum, coordconv 3 + 3 + 1 + 2
+        self.actor_target = ResNet(9, 18, 65)
+        self.critic = ResNet_wobn(3 + 9, 18, 1) # add the last canvas for better prediction
+        self.critic_target = ResNet_wobn(3 + 9, 18, 1) 
+
+        self.actor_optim  = Adam(self.actor.parameters(), lr=1e-2)
+        self.critic_optim  = Adam(self.critic.parameters(), lr=1e-2)
+
+        if (resume != None):
+            self.load_weights(resume)
+
+        hard_update(self.actor_target, self.actor)
+        hard_update(self.critic_target, self.critic)
+        
+        # Create replay buffer
+        self.memory = rpm(rmsize * max_step)
+
+        # Hyper-parameters
+        self.tau = tau
+        self.discount = discount
+
+        # Tensorboard
+        self.writer = writer
+        self.log = 0
+        
+        self.state = [None] * self.env_batch # Most recent state
+        self.action = [None] * self.env_batch # Most recent action
+        self.choose_device()        
+
+    def play(self, state, target=False):
+        state = torch.cat((state[:, :6].float() / 255, state[:, 6:7].float() / self.max_step, coord.expand(state.shape[0], 2, 128, 128)), 1)
+        if target:
+            return self.actor_target(state)
+        else:
+            return self.actor(state)
+
+    def update_gan(self, state):
+        canvas = state[:, :3]
+        gt = state[:, 3 : 6]
+        fake, real, penal = update(canvas.float() / 255, gt.float() / 255)
+        if self.log % 20 == 0:
+            self.writer.add_scalar('train/gan_fake', fake, self.log)
+            self.writer.add_scalar('train/gan_real', real, self.log)
+            self.writer.add_scalar('train/gan_penal', penal, self.log)       
+        
+    def evaluate(self, state, action, target=False):
+        T = state[:, 6 : 7]
+        gt = state[:, 3 : 6].float() / 255
+        canvas0 = state[:, :3].float() / 255
+        canvas1 = decode(action, canvas0)
+        gan_reward = cal_reward(canvas1, gt) - cal_reward(canvas0, gt)
+        # L2_reward = ((canvas0 - gt) ** 2).mean(1).mean(1).mean(1) - ((canvas1 - gt) ** 2).mean(1).mean(1).mean(1)        
+        coord_ = coord.expand(state.shape[0], 2, 128, 128)
+        merged_state = torch.cat([canvas0, canvas1, gt, (T + 1).float() / self.max_step, coord_], 1)
+        # canvas0 is not necessarily added
+        if target:
+            Q = self.critic_target(merged_state)
+            return (Q + gan_reward), gan_reward
+        else:
+            Q = self.critic(merged_state)
+            if self.log % 20 == 0:
+                self.writer.add_scalar('train/expect_reward', Q.mean(), self.log)
+                self.writer.add_scalar('train/gan_reward', gan_reward.mean(), self.log)
+            return (Q + gan_reward), gan_reward
+    
+    def update_policy(self, lr):
+        self.log += 1
+        
+        for param_group in self.critic_optim.param_groups:
+            param_group['lr'] = lr[0]
+        for param_group in self.actor_optim.param_groups:
+            param_group['lr'] = lr[1]
+            
+        # Sample batch
+        state, action, reward, \
+            next_state, terminal = self.memory.sample_batch(self.batch_size, device)
+
+        self.update_gan(next_state)
+        
+        with torch.no_grad():
+            next_action = self.play(next_state, True)
+            target_q, _ = self.evaluate(next_state, next_action, True)
+            target_q = self.discount * ((1 - terminal.float()).view(-1, 1)) * target_q
+                
+        cur_q, step_reward = self.evaluate(state, action)
+        target_q += step_reward.detach()
+        
+        value_loss = criterion(cur_q, target_q)
+        self.critic.zero_grad()
+        value_loss.backward(retain_graph=True)
+        self.critic_optim.step()
+
+        action = self.play(state)
+        pre_q, _ = self.evaluate(state.detach(), action)
+        policy_loss = -pre_q.mean()
+        self.actor.zero_grad()
+        policy_loss.backward(retain_graph=True)
+        self.actor_optim.step()
+        
+        # Target update
+        soft_update(self.actor_target, self.actor, self.tau)
+        soft_update(self.critic_target, self.critic, self.tau)
+
+        return -policy_loss, value_loss
+
+    def observe(self, reward, state, done, step):
+        s0 = torch.tensor(self.state, device='cpu')
+        a = to_tensor(self.action, "cpu")
+        r = to_tensor(reward, "cpu")
+        s1 = torch.tensor(state, device='cpu')
+        d = to_tensor(done.astype('float32'), "cpu")
+        for i in range(self.env_batch):
+            self.memory.append([s0[i], a[i], r[i], s1[i], d[i]])
+        self.state = state
+
+    def noise_action(self, noise_factor, state, action):
+        noise = np.zeros(action.shape)
+        for i in range(self.env_batch):
+            action[i] = action[i] + np.random.normal(0, self.noise_level[i], action.shape[1:]).astype('float32')
+        return np.clip(action.astype('float32'), 0, 1)
+    
+    def select_action(self, state, return_fix=False, noise_factor=0):
+        self.eval()
+        with torch.no_grad():
+            action = self.play(state)
+            action = to_numpy(action)
+        if noise_factor > 0:        
+            action = self.noise_action(noise_factor, state, action)
+        self.train()
+        self.action = action
+        if return_fix:
+            return action
+        return self.action
+
+    def reset(self, obs, factor):
+        self.state = obs
+        self.noise_level = np.random.uniform(0, factor, self.env_batch)
+
+    def load_weights(self, path):
+        if path is None: return
+        self.actor.load_state_dict(torch.load('{}/actor.pkl'.format(path)))
+        self.critic.load_state_dict(torch.load('{}/critic.pkl'.format(path)))
+        load_gan(path)
+        
+    def save_model(self, path):
+        self.actor.cpu()
+        self.critic.cpu()
+        torch.save(self.actor.state_dict(),'{}/actor.pkl'.format(path))
+        torch.save(self.critic.state_dict(),'{}/critic.pkl'.format(path))
+        save_gan(path)
+        self.choose_device()
+
+    def eval(self):
+        self.actor.eval()
+        self.actor_target.eval()
+        self.critic.eval()
+        self.critic_target.eval()
+    
+    def train(self):
+        self.actor.train()
+        self.actor_target.train()
+        self.critic.train()
+        self.critic_target.train()
+    
+    def choose_device(self):
+        Decoder.to(device)
+        self.actor.to(device)
+        self.actor_target.to(device)
+        self.critic.to(device)
+        self.critic_target.to(device)

baseline/DRL/evaluator.py

@@ -0,0 +1,31 @@
+import numpy as np
+from utils.util import *
+
+class Evaluator(object):
+
+    def __init__(self, args, writer):    
+        self.validate_episodes = args.validate_episodes
+        self.max_step = args.max_step
+        self.env_batch = args.env_batch
+        self.writer = writer
+        self.log = 0
+
+    def __call__(self, env, policy, debug=False):        
+        observation = None
+        for episode in range(self.validate_episodes):
+            # reset at the start of episode
+            observation = env.reset(test=True, episode=episode)
+            episode_steps = 0
+            episode_reward = 0.     
+            assert observation is not None            
+            # start episode
+            episode_reward = np.zeros(self.env_batch)
+            while (episode_steps < self.max_step or not self.max_step):
+                action = policy(observation)
+                observation, reward, done, (step_num) = env.step(action)
+                episode_reward += reward
+                episode_steps += 1
+                env.save_image(self.log, episode_steps)
+            dist = env.get_dist()
+            self.log += 1
+        return episode_reward, dist

baseline/DRL/multi.py

@@ -0,0 +1,53 @@
+import cv2
+import torch
+import numpy as np
+from env import Paint
+from utils.util import *
+from DRL.ddpg import decode
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class fastenv():
+    def __init__(self, 
+                 max_episode_length=10, env_batch=64, \
+                 writer=None):
+        self.max_episode_length = max_episode_length
+        self.env_batch = env_batch
+        self.env = Paint(self.env_batch, self.max_episode_length)
+        self.env.load_data()
+        self.observation_space = self.env.observation_space
+        self.action_space = self.env.action_space
+        self.writer = writer
+        self.test = False
+        self.log = 0
+
+    def save_image(self, log, step):
+        for i in range(self.env_batch):
+            if self.env.imgid[i] <= 10:
+                canvas = cv2.cvtColor((to_numpy(self.env.canvas[i].permute(1, 2, 0))), cv2.COLOR_BGR2RGB)
+                self.writer.add_image('{}/canvas_{}.png'.format(str(self.env.imgid[i]), str(step)), canvas, log)
+        if step == self.max_episode_length:
+            for i in range(self.env_batch):
+                if self.env.imgid[i] < 50:
+                    gt = cv2.cvtColor((to_numpy(self.env.gt[i].permute(1, 2, 0))), cv2.COLOR_BGR2RGB)
+                    canvas = cv2.cvtColor((to_numpy(self.env.canvas[i].permute(1, 2, 0))), cv2.COLOR_BGR2RGB)
+                    self.writer.add_image(str(self.env.imgid[i]) + '/_target.png', gt, log)
+                    self.writer.add_image(str(self.env.imgid[i]) + '/_canvas.png', canvas, log)
+    
+    def step(self, action):
+        with torch.no_grad():
+            ob, r, d, _ = self.env.step(torch.tensor(action).to(device))
+        if d[0]:
+            if not self.test:
+                self.dist = self.get_dist()
+                for i in range(self.env_batch):
+                    self.writer.add_scalar('train/dist', self.dist[i], self.log)
+                    self.log += 1
+        return ob, r, d, _
+
+    def get_dist(self):
+        return to_numpy((((self.env.gt.float() - self.env.canvas.float()) / 255) ** 2).mean(1).mean(1).mean(1))
+        
+    def reset(self, test=False, episode=0):
+        self.test = test
+        ob = self.env.reset(self.test, episode * self.env_batch)
+        return ob

baseline/DRL/rpm.py

@@ -0,0 +1,43 @@
+# from collections import deque
+import numpy as np
+import random
+import torch
+import pickle as pickle
+
+class rpm(object):
+    # replay memory
+    def __init__(self, buffer_size):
+        self.buffer_size = buffer_size
+        self.buffer = []
+        self.index = 0
+        
+    def append(self, obj):
+        if self.size() > self.buffer_size:
+            print('buffer size larger than set value, trimming...')
+            self.buffer = self.buffer[(self.size() - self.buffer_size):]
+        elif self.size() == self.buffer_size:
+            self.buffer[self.index] = obj
+            self.index += 1
+            self.index %= self.buffer_size
+        else:
+            self.buffer.append(obj)
+
+    def size(self):
+        return len(self.buffer)
+
+    def sample_batch(self, batch_size, device, only_state=False):
+        if self.size() < batch_size:
+            batch = random.sample(self.buffer, self.size())
+        else:
+            batch = random.sample(self.buffer, batch_size)
+
+        if only_state:
+            res = torch.stack(tuple(item[3] for item in batch), dim=0)            
+            return res.to(device)
+        else:
+            item_count = 5
+            res = []
+            for i in range(5):
+                k = torch.stack(tuple(item[i] for item in batch), dim=0)
+                res.append(k.to(device))
+            return res[0], res[1], res[2], res[3], res[4]

baseline/DRL/wgan.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim import Adam, SGD
+from torch import autograd
+from torch.autograd import Variable
+import torch.nn.functional as F
+from torch.autograd import grad as torch_grad
+import torch.nn.utils.weight_norm as weightNorm
+from utils.util import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dim = 128
+LAMBDA = 10 # Gradient penalty lambda hyperparameter
+
+class TReLU(nn.Module):
+    def __init__(self):
+            super(TReLU, self).__init__()
+            self.alpha = nn.Parameter(torch.FloatTensor(1), requires_grad=True)
+            self.alpha.data.fill_(0)
+
+    def forward(self, x):
+        x = F.relu(x - self.alpha) + self.alpha
+        return x
+
+class Discriminator(nn.Module):
+        def __init__(self):
+            super(Discriminator, self).__init__()
+
+            self.conv0 = weightNorm(nn.Conv2d(6, 16, 5, 2, 2))
+            self.conv1 = weightNorm(nn.Conv2d(16, 32, 5, 2, 2))
+            self.conv2 = weightNorm(nn.Conv2d(32, 64, 5, 2, 2))
+            self.conv3 = weightNorm(nn.Conv2d(64, 128, 5, 2, 2))
+            self.conv4 = weightNorm(nn.Conv2d(128, 1, 5, 2, 2))
+            self.relu0 = TReLU()
+            self.relu1 = TReLU()
+            self.relu2 = TReLU()
+            self.relu3 = TReLU()
+
+        def forward(self, x):
+            x = self.conv0(x)
+            x = self.relu0(x)
+            x = self.conv1(x)
+            x = self.relu1(x)
+            x = self.conv2(x)
+            x = self.relu2(x)
+            x = self.conv3(x)
+            x = self.relu3(x)
+            x = self.conv4(x)
+            x = F.avg_pool2d(x, 4)
+            x = x.view(-1, 1)
+            return x
+
+netD = Discriminator()
+target_netD = Discriminator()
+netD = netD.to(device)
+target_netD = target_netD.to(device)
+hard_update(target_netD, netD)
+
+optimizerD = Adam(netD.parameters(), lr=3e-4, betas=(0.5, 0.999))
+def cal_gradient_penalty(netD, real_data, fake_data, batch_size):
+    alpha = torch.rand(batch_size, 1)
+    alpha = alpha.expand(batch_size, int(real_data.nelement()/batch_size)).contiguous()
+    alpha = alpha.view(batch_size, 6, dim, dim)
+    alpha = alpha.to(device)
+    fake_data = fake_data.view(batch_size, 6, dim, dim)
+    interpolates = Variable(alpha * real_data.data + ((1 - alpha) * fake_data.data), requires_grad=True)
+    disc_interpolates = netD(interpolates)
+    gradients = autograd.grad(disc_interpolates, interpolates,
+                              grad_outputs=torch.ones(disc_interpolates.size()).to(device),
+                              create_graph=True, retain_graph=True)[0]
+    gradients = gradients.view(gradients.size(0), -1)
+    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * LAMBDA
+    return gradient_penalty
+
+def cal_reward(fake_data, real_data):
+    return target_netD(torch.cat([real_data, fake_data], 1))
+
+def save_gan(path):
+    netD.cpu()
+    torch.save(netD.state_dict(),'{}/wgan.pkl'.format(path))
+    netD.to(device)
+
+def load_gan(path):
+    netD.load_state_dict(torch.load('{}/wgan.pkl'.format(path)))
+
+def update(fake_data, real_data):
+    fake_data = fake_data.detach()
+    real_data = real_data.detach()
+    fake = torch.cat([real_data, fake_data], 1)
+    real = torch.cat([real_data, real_data], 1)
+    D_real = netD(real)
+    D_fake = netD(fake)
+    gradient_penalty = cal_gradient_penalty(netD, real, fake, real.shape[0])
+    optimizerD.zero_grad()
+    D_cost = D_fake.mean() - D_real.mean() + gradient_penalty
+    D_cost.backward()
+    optimizerD.step()
+    soft_update(target_netD, netD, 0.001)
+    return D_fake.mean(), D_real.mean(), gradient_penalty

baseline/Renderer/model.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.weight_norm as weightNorm
+
+class FCN(nn.Module):
+    def __init__(self):
+        super(FCN, self).__init__()
+        self.fc1 = (nn.Linear(10, 512))
+        self.fc2 = (nn.Linear(512, 1024))
+        self.fc3 = (nn.Linear(1024, 2048))
+        self.fc4 = (nn.Linear(2048, 4096))
+        self.conv1 = (nn.Conv2d(16, 32, 3, 1, 1))
+        self.conv2 = (nn.Conv2d(32, 32, 3, 1, 1))
+        self.conv3 = (nn.Conv2d(8, 16, 3, 1, 1))
+        self.conv4 = (nn.Conv2d(16, 16, 3, 1, 1))
+        self.conv5 = (nn.Conv2d(4, 8, 3, 1, 1))
+        self.conv6 = (nn.Conv2d(8, 4, 3, 1, 1))
+        self.pixel_shuffle = nn.PixelShuffle(2)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc3(x))
+        x = F.relu(self.fc4(x))
+        x = x.view(-1, 16, 16, 16)
+        x = F.relu(self.conv1(x))
+        x = self.pixel_shuffle(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = self.pixel_shuffle(self.conv4(x))
+        x = F.relu(self.conv5(x))
+        x = self.pixel_shuffle(self.conv6(x))
+        x = torch.sigmoid(x)
+        return 1 - x.view(-1, 128, 128)

baseline/Renderer/stroke_gen.py

@@ -0,0 +1,28 @@
+import cv2
+import numpy as np
+
+def normal(x, width):
+    return (int)(x * (width - 1) + 0.5)
+
+def draw(f, width=128):
+    x0, y0, x1, y1, x2, y2, z0, z2, w0, w2 = f
+    x1 = x0 + (x2 - x0) * x1
+    y1 = y0 + (y2 - y0) * y1
+    x0 = normal(x0, width * 2)
+    x1 = normal(x1, width * 2)
+    x2 = normal(x2, width * 2)
+    y0 = normal(y0, width * 2)
+    y1 = normal(y1, width * 2)
+    y2 = normal(y2, width * 2)
+    z0 = (int)(1 + z0 * width // 2)
+    z2 = (int)(1 + z2 * width // 2)
+    canvas = np.zeros([width * 2, width * 2]).astype('float32')
+    tmp = 1. / 100
+    for i in range(100):
+        t = i * tmp
+        x = (int)((1-t) * (1-t) * x0 + 2 * t * (1-t) * x1 + t * t * x2)
+        y = (int)((1-t) * (1-t) * y0 + 2 * t * (1-t) * y1 + t * t * y2)
+        z = (int)((1-t) * z0 + t * z2)
+        w = (1-t) * w0 + t * w2
+        cv2.circle(canvas, (y, x), z, w, -1)
+    return 1 - cv2.resize(canvas, dsize=(width, width))

baseline/env.py

@@ -0,0 +1,107 @@
+import sys
+import json
+import torch
+import numpy as np
+import argparse
+import torchvision.transforms as transforms
+import cv2
+from DRL.ddpg import decode
+from utils.util import *
+from PIL import Image
+from torchvision import transforms, utils
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+aug = transforms.Compose(
+            [transforms.ToPILImage(),
+             transforms.RandomHorizontalFlip(),
+             ])
+
+width = 128
+convas_area = width * width
+
+img_train = []
+img_test = []
+train_num = 0
+test_num = 0
+
+class Paint:
+    def __init__(self, batch_size, max_step):
+        self.batch_size = batch_size
+        self.max_step = max_step
+        self.action_space = (13)
+        self.observation_space = (self.batch_size, width, width, 7)
+        self.test = False
+        
+    def load_data(self):
+        # CelebA
+        global train_num, test_num
+        for i in range(200000):
+            img_id = '%06d' % (i + 1)
+            try:
+                img = cv2.imread('./data/img_align_celeba/' + img_id + '.jpg', cv2.IMREAD_UNCHANGED)
+                img = cv2.resize(img, (width, width))
+                if i > 2000:                
+                    train_num += 1
+                    img_train.append(img)
+                else:
+                    test_num += 1
+                    img_test.append(img)
+            finally:
+                if (i + 1) % 10000 == 0:                    
+                    print('loaded {} images'.format(i + 1))
+        print('finish loading data, {} training images, {} testing images'.format(str(train_num), str(test_num)))
+        
+    def pre_data(self, id, test):
+        if test:
+            img = img_test[id]
+        else:
+            img = img_train[id]
+        if not test:
+            img = aug(img)
+        img = np.asarray(img)
+        return np.transpose(img, (2, 0, 1))
+    
+    def reset(self, test=False, begin_num=False):
+        self.test = test
+        self.imgid = [0] * self.batch_size
+        self.gt = torch.zeros([self.batch_size, 3, width, width], dtype=torch.uint8).to(device)
+        for i in range(self.batch_size):
+            if test:
+                id = (i + begin_num)  % test_num
+            else:
+                id = np.random.randint(train_num)
+            self.imgid[i] = id
+            self.gt[i] = torch.tensor(self.pre_data(id, test))
+        self.tot_reward = ((self.gt.float() / 255) ** 2).mean(1).mean(1).mean(1)
+        self.stepnum = 0
+        self.canvas = torch.zeros([self.batch_size, 3, width, width], dtype=torch.uint8).to(device)
+        self.lastdis = self.ini_dis = self.cal_dis()
+        return self.observation()
+    
+    def observation(self):
+        # canvas B * 3 * width * width
+        # gt B * 3 * width * width
+        # T B * 1 * width * width
+        ob = []
+        T = torch.ones([self.batch_size, 1, width, width], dtype=torch.uint8) * self.stepnum
+        return torch.cat((self.canvas, self.gt, T.to(device)), 1) # canvas, img, T
+
+    def cal_trans(self, s, t):
+        return (s.transpose(0, 3) * t).transpose(0, 3)
+    
+    def step(self, action):
+        self.canvas = (decode(action, self.canvas.float() / 255) * 255).byte()
+        self.stepnum += 1
+        ob = self.observation()
+        done = (self.stepnum == self.max_step)
+        reward = self.cal_reward() # np.array([0.] * self.batch_size)
+        return ob.detach(), reward, np.array([done] * self.batch_size), None
+
+    def cal_dis(self):
+        return (((self.canvas.float() - self.gt.float()) / 255) ** 2).mean(1).mean(1).mean(1)
+    
+    def cal_reward(self):
+        dis = self.cal_dis()
+        reward = (self.lastdis - dis) / (self.ini_dis + 1e-8)
+        self.lastdis = dis
+        return to_numpy(reward)

baseline/utils/tensorboard.py

@@ -0,0 +1,29 @@
+import PIL
+import scipy.misc
+from io import BytesIO
+import tensorboardX as tb
+from tensorboardX.summary import Summary
+
+class TensorBoard(object):
+    def __init__(self, model_dir):
+        self.summary_writer = tb.FileWriter(model_dir)
+
+    def add_image(self, tag, img, step):
+        summary = Summary()
+        bio = BytesIO()
+
+        if type(img) == str:
+            img = PIL.Image.open(img)
+        elif type(img) == PIL.Image.Image:
+            pass
+        else:
+            img = scipy.misc.toimage(img)
+
+        img.save(bio, format="png")
+        image_summary = Summary.Image(encoded_image_string=bio.getvalue())
+        summary.value.add(tag=tag, image=image_summary)
+        self.summary_writer.add_summary(summary, global_step=step)
+
+    def add_scalar(self, tag, value, step):
+        summary = Summary(value=[Summary.Value(tag=tag, simple_value=value)])
+        self.summary_writer.add_summary(summary, global_step=step)

baseline/utils/util.py

@@ -0,0 +1,69 @@
+import os
+import torch
+from torch.autograd import Variable
+
+USE_CUDA = torch.cuda.is_available()
+
+def prRed(prt): print("\033[91m {}\033[00m" .format(prt))
+def prGreen(prt): print("\033[92m {}\033[00m" .format(prt))
+def prYellow(prt): print("\033[93m {}\033[00m" .format(prt))
+def prLightPurple(prt): print("\033[94m {}\033[00m" .format(prt))
+def prPurple(prt): print("\033[95m {}\033[00m" .format(prt))
+def prCyan(prt): print("\033[96m {}\033[00m" .format(prt))
+def prLightGray(prt): print("\033[97m {}\033[00m" .format(prt))
+def prBlack(prt): print("\033[98m {}\033[00m" .format(prt))
+
+def to_numpy(var):
+    return var.cpu().data.numpy() if USE_CUDA else var.data.numpy()
+
+def to_tensor(ndarray, device):
+    return torch.tensor(ndarray, dtype=torch.float, device=device)
+
+def soft_update(target, source, tau):
+    for target_param, param in zip(target.parameters(), source.parameters()):
+        target_param.data.copy_(
+            target_param.data * (1.0 - tau) + param.data * tau
+        )
+
+def hard_update(target, source):
+    for m1, m2 in zip(target.modules(), source.modules()):
+        m1._buffers = m2._buffers.copy()
+    for target_param, param in zip(target.parameters(), source.parameters()):
+            target_param.data.copy_(param.data)
+
+def get_output_folder(parent_dir, env_name):
+    """Return save folder.
+
+    Assumes folders in the parent_dir have suffix -run{run
+    number}. Finds the highest run number and sets the output folder
+    to that number + 1. This is just convenient so that if you run the
+    same script multiple times tensorboard can plot all of the results
+    on the same plots with different names.
+
+    Parameters
+    ----------
+    parent_dir: str
+      Path of the directory containing all experiment runs.
+
+    Returns
+    -------
+    parent_dir/run_dir
+      Path to this run's save directory.
+    """
+    os.makedirs(parent_dir, exist_ok=True)
+    experiment_id = 0
+    for folder_name in os.listdir(parent_dir):
+        if not os.path.isdir(os.path.join(parent_dir, folder_name)):
+            continue
+        try:
+            folder_name = int(folder_name.split('-run')[-1])
+            if folder_name > experiment_id:
+                experiment_id = folder_name
+        except:
+            pass
+    experiment_id += 1
+
+    parent_dir = os.path.join(parent_dir, env_name)
+    parent_dir = parent_dir + '-run{}'.format(experiment_id)
+    os.makedirs(parent_dir, exist_ok=True)
+    return parent_dir

ckpts/actor.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75e908a42d4c90ad092892a9183bd467453696d8e2202b9640f9a5b4a488eab9
+size 44898539

ckpts/renderer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34639d4ab4f30807fb056bf4b71de0c3f434d6cbe8600cf97b6af3a855a5ca8e
+size 44165821