Reevee/ohayo_face_style
0
1import torch2from torch.autograd import Variable3from collections import OrderedDict4import numpy as np5import os6from PIL import Image7import util.util as util8from .base_model import BaseModel9from . import networks10 11class UIModel(BaseModel):12 def name(self):13 return 'UIModel'14 15 def initialize(self, opt):16 assert(not opt.isTrain)17 BaseModel.initialize(self, opt)18 self.use_features = opt.instance_feat or opt.label_feat19 20 netG_input_nc = opt.label_nc21 if not opt.no_instance:22 netG_input_nc += 1 23 if self.use_features: 24 netG_input_nc += opt.feat_num 25 26 self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG, 27 opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 28 opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids) 29 self.load_network(self.netG, 'G', opt.which_epoch)30 31 print('---------- Networks initialized -------------')32 33 def toTensor(self, img, normalize=False):34 tensor = torch.from_numpy(np.array(img, np.int32, copy=False))35 tensor = tensor.view(1, img.size[1], img.size[0], len(img.mode)) 36 tensor = tensor.transpose(1, 2).transpose(1, 3).contiguous()37 if normalize:38 return (tensor.float()/255.0 - 0.5) / 0.5 39 return tensor.float()40 41 def load_image(self, label_path, inst_path, feat_path):42 opt = self.opt43 # read label map44 label_img = Image.open(label_path) 45 if label_path.find('face') != -1:46 label_img = label_img.convert('L')47 ow, oh = label_img.size 48 w = opt.loadSize49 h = int(w * oh / ow) 50 label_img = label_img.resize((w, h), Image.NEAREST)51 label_map = self.toTensor(label_img) 52 53 # onehot vector input for label map54 self.label_map = label_map.cuda()55 oneHot_size = (1, opt.label_nc, h, w)56 input_label = self.Tensor(torch.Size(oneHot_size)).zero_()57 self.input_label = input_label.scatter_(1, label_map.long().cuda(), 1.0)58 59 # read instance map60 if not opt.no_instance:61 inst_img = Image.open(inst_path) 62 inst_img = inst_img.resize((w, h), Image.NEAREST) 63 self.inst_map = self.toTensor(inst_img).cuda()64 self.edge_map = self.get_edges(self.inst_map) 65 self.net_input = Variable(torch.cat((self.input_label, self.edge_map), dim=1), volatile=True)66 else:67 self.net_input = Variable(self.input_label, volatile=True) 68 69 self.features_clustered = np.load(feat_path).item()70 self.object_map = self.inst_map if opt.instance_feat else self.label_map 71 72 object_np = self.object_map.cpu().numpy().astype(int) 73 self.feat_map = self.Tensor(1, opt.feat_num, h, w).zero_() 74 self.cluster_indices = np.zeros(self.opt.label_nc, np.uint8)75 for i in np.unique(object_np): 76 label = i if i < 1000 else i//100077 if label in self.features_clustered:78 feat = self.features_clustered[label]79 np.random.seed(i+1)80 cluster_idx = np.random.randint(0, feat.shape[0])81 self.cluster_indices[label] = cluster_idx82 idx = (self.object_map == i).nonzero() 83 self.set_features(idx, feat, cluster_idx)84 85 self.net_input_original = self.net_input.clone() 86 self.label_map_original = self.label_map.clone()87 self.feat_map_original = self.feat_map.clone()88 if not opt.no_instance:89 self.inst_map_original = self.inst_map.clone() 90 91 def reset(self):92 self.net_input = self.net_input_prev = self.net_input_original.clone() 93 self.label_map = self.label_map_prev = self.label_map_original.clone()94 self.feat_map = self.feat_map_prev = self.feat_map_original.clone()95 if not self.opt.no_instance:96 self.inst_map = self.inst_map_prev = self.inst_map_original.clone()97 self.object_map = self.inst_map if self.opt.instance_feat else self.label_map 98 99 def undo(self): 100 self.net_input = self.net_input_prev101 self.label_map = self.label_map_prev102 self.feat_map = self.feat_map_prev103 if not self.opt.no_instance:104 self.inst_map = self.inst_map_prev105 self.object_map = self.inst_map if self.opt.instance_feat else self.label_map 106 107 # get boundary map from instance map108 def get_edges(self, t):109 edge = torch.cuda.ByteTensor(t.size()).zero_()110 edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])111 edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])112 edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])113 edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])114 return edge.float()115 116 # change the label at the source position to the label at the target position117 def change_labels(self, click_src, click_tgt): 118 y_src, x_src = click_src[0], click_src[1]119 y_tgt, x_tgt = click_tgt[0], click_tgt[1]120 label_src = int(self.label_map[0, 0, y_src, x_src])121 inst_src = self.inst_map[0, 0, y_src, x_src]122 label_tgt = int(self.label_map[0, 0, y_tgt, x_tgt])123 inst_tgt = self.inst_map[0, 0, y_tgt, x_tgt]124 125 idx_src = (self.inst_map == inst_src).nonzero() 126 # need to change 3 things: label map, instance map, and feature map127 if idx_src.shape:128 # backup current maps129 self.backup_current_state() 130 131 # change both the label map and the network input132 self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt133 self.net_input[idx_src[:,0], idx_src[:,1] + label_src, idx_src[:,2], idx_src[:,3]] = 0134 self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1 135 136 # update the instance map (and the network input)137 if inst_tgt > 1000:138 # if different instances have different ids, give the new object a new id139 tgt_indices = (self.inst_map > label_tgt * 1000) & (self.inst_map < (label_tgt+1) * 1000)140 inst_tgt = self.inst_map[tgt_indices].max() + 1141 self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = inst_tgt142 self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)143 144 # also copy the source features to the target position 145 idx_tgt = (self.inst_map == inst_tgt).nonzero() 146 if idx_tgt.shape:147 self.copy_features(idx_src, idx_tgt[0,:])148 149 self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))150 151 # add strokes of target label in the image152 def add_strokes(self, click_src, label_tgt, bw, save):153 # get the region of the new strokes (bw is the brush width) 154 size = self.net_input.size()155 h, w = size[2], size[3]156 idx_src = torch.LongTensor(bw**2, 4).fill_(0)157 for i in range(bw):158 idx_src[i*bw:(i+1)*bw, 2] = min(h-1, max(0, click_src[0]-bw//2 + i))159 for j in range(bw):160 idx_src[i*bw+j, 3] = min(w-1, max(0, click_src[1]-bw//2 + j))161 idx_src = idx_src.cuda()162 163 # again, need to update 3 things164 if idx_src.shape:165 # backup current maps166 if save:167 self.backup_current_state()168 169 # update the label map (and the network input) in the stroke region 170 self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt171 for k in range(self.opt.label_nc):172 self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0173 self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1 174 175 # update the instance map (and the network input)176 self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt177 self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)178 179 # also update the features if available180 if self.opt.instance_feat: 181 feat = self.features_clustered[label_tgt]182 #np.random.seed(label_tgt+1) 183 #cluster_idx = np.random.randint(0, feat.shape[0])184 cluster_idx = self.cluster_indices[label_tgt]185 self.set_features(idx_src, feat, cluster_idx) 186 187 self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))188 189 # add an object to the clicked position with selected style190 def add_objects(self, click_src, label_tgt, mask, style_id=0):191 y, x = click_src[0], click_src[1]192 mask = np.transpose(mask, (2, 0, 1))[np.newaxis,...] 193 idx_src = torch.from_numpy(mask).cuda().nonzero() 194 idx_src[:,2] += y195 idx_src[:,3] += x196 197 # backup current maps198 self.backup_current_state()199 200 # update label map201 self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt 202 for k in range(self.opt.label_nc):203 self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0204 self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1 205 206 # update instance map207 self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt208 self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)209 210 # update feature map211 self.set_features(idx_src, self.feat, style_id) 212 213 self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))214 215 def single_forward(self, net_input, feat_map):216 net_input = torch.cat((net_input, feat_map), dim=1)217 fake_image = self.netG.forward(net_input)218 219 if fake_image.size()[0] == 1:220 return fake_image.data[0] 221 return fake_image.data222 223 224 # generate all outputs for different styles225 def style_forward(self, click_pt, style_id=-1): 226 if click_pt is None: 227 self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))228 self.crop = None229 self.mask = None 230 else: 231 instToChange = int(self.object_map[0, 0, click_pt[0], click_pt[1]])232 self.instToChange = instToChange233 label = instToChange if instToChange < 1000 else instToChange//1000 234 self.feat = self.features_clustered[label]235 self.fake_image = []236 self.mask = self.object_map == instToChange237 idx = self.mask.nonzero()238 self.get_crop_region(idx) 239 if idx.size(): 240 if style_id == -1:241 (min_y, min_x, max_y, max_x) = self.crop242 ### original243 for cluster_idx in range(self.opt.multiple_output):244 self.set_features(idx, self.feat, cluster_idx)245 fake_image = self.single_forward(self.net_input, self.feat_map)246 fake_image = util.tensor2im(fake_image[:,min_y:max_y,min_x:max_x])247 self.fake_image.append(fake_image) 248 """### To speed up previewing different style results, either crop or downsample the label maps249 if instToChange > 1000:250 (min_y, min_x, max_y, max_x) = self.crop 251 ### crop 252 _, _, h, w = self.net_input.size()253 offset = 512254 y_start, x_start = max(0, min_y-offset), max(0, min_x-offset)255 y_end, x_end = min(h, (max_y + offset)), min(w, (max_x + offset))256 y_region = slice(y_start, y_start+(y_end-y_start)//16*16)257 x_region = slice(x_start, x_start+(x_end-x_start)//16*16)258 net_input = self.net_input[:,:,y_region,x_region] 259 for cluster_idx in range(self.opt.multiple_output): 260 self.set_features(idx, self.feat, cluster_idx)261 fake_image = self.single_forward(net_input, self.feat_map[:,:,y_region,x_region]) 262 fake_image = util.tensor2im(fake_image[:,min_y-y_start:max_y-y_start,min_x-x_start:max_x-x_start])263 self.fake_image.append(fake_image)264 else:265 ### downsample266 (min_y, min_x, max_y, max_x) = [crop//2 for crop in self.crop] 267 net_input = self.net_input[:,:,::2,::2] 268 size = net_input.size()269 net_input_batch = net_input.expand(self.opt.multiple_output, size[1], size[2], size[3]) 270 for cluster_idx in range(self.opt.multiple_output): 271 self.set_features(idx, self.feat, cluster_idx)272 feat_map = self.feat_map[:,:,::2,::2]273 if cluster_idx == 0:274 feat_map_batch = feat_map275 else:276 feat_map_batch = torch.cat((feat_map_batch, feat_map), dim=0)277 fake_image_batch = self.single_forward(net_input_batch, feat_map_batch)278 for i in range(self.opt.multiple_output):279 self.fake_image.append(util.tensor2im(fake_image_batch[i,:,min_y:max_y,min_x:max_x]))"""280 281 else:282 self.set_features(idx, self.feat, style_id)283 self.cluster_indices[label] = style_id284 self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map)) 285 286 def backup_current_state(self):287 self.net_input_prev = self.net_input.clone()288 self.label_map_prev = self.label_map.clone() 289 self.inst_map_prev = self.inst_map.clone() 290 self.feat_map_prev = self.feat_map.clone() 291 292 # crop the ROI and get the mask of the object293 def get_crop_region(self, idx):294 size = self.net_input.size()295 h, w = size[2], size[3]296 min_y, min_x = idx[:,2].min(), idx[:,3].min()297 max_y, max_x = idx[:,2].max(), idx[:,3].max() 298 crop_min = 128299 if max_y - min_y < crop_min:300 min_y = max(0, (max_y + min_y) // 2 - crop_min // 2)301 max_y = min(h-1, min_y + crop_min)302 if max_x - min_x < crop_min:303 min_x = max(0, (max_x + min_x) // 2 - crop_min // 2)304 max_x = min(w-1, min_x + crop_min)305 self.crop = (min_y, min_x, max_y, max_x) 306 self.mask = self.mask[:,:, min_y:max_y, min_x:max_x]307 308 # update the feature map once a new object is added or the label is changed309 def update_features(self, cluster_idx, mask=None, click_pt=None): 310 self.feat_map_prev = self.feat_map.clone()311 # adding a new object312 if mask is not None:313 y, x = click_pt[0], click_pt[1]314 mask = np.transpose(mask, (2,0,1))[np.newaxis,...] 315 idx = torch.from_numpy(mask).cuda().nonzero() 316 idx[:,2] += y317 idx[:,3] += x 318 # changing the label of an existing object 319 else: 320 idx = (self.object_map == self.instToChange).nonzero() 321 322 # update feature map323 self.set_features(idx, self.feat, cluster_idx) 324 325 # set the class features to the target feature326 def set_features(self, idx, feat, cluster_idx): 327 for k in range(self.opt.feat_num):328 self.feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k] 329 330 # copy the features at the target position to the source position331 def copy_features(self, idx_src, idx_tgt): 332 for k in range(self.opt.feat_num):333 val = self.feat_map[idx_tgt[0], idx_tgt[1] + k, idx_tgt[2], idx_tgt[3]]334 self.feat_map[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = val 335 336 def get_current_visuals(self, getLabel=False): 337 mask = self.mask 338 if self.mask is not None:339 mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8) 340 341 dict_list = [('fake_image', self.fake_image), ('mask', mask)]342 343 if getLabel: # only output label map if needed to save bandwidth344 label = util.tensor2label(self.net_input.data[0], self.opt.label_nc) 345 dict_list += [('label', label)]346 347 return OrderedDict(dict_list)