ui_model.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347
  1. import torch
  2. from torch.autograd import Variable
  3. from collections import OrderedDict
  4. import numpy as np
  5. import os
  6. from PIL import Image
  7. import util.util as util
  8. from .base_model import BaseModel
  9. from . import networks
  10. class UIModel(BaseModel):
  11. def name(self):
  12. return 'UIModel'
  13. def initialize(self, opt):
  14. assert(not opt.isTrain)
  15. BaseModel.initialize(self, opt)
  16. self.use_features = opt.instance_feat or opt.label_feat
  17. netG_input_nc = opt.label_nc
  18. if not opt.no_instance:
  19. netG_input_nc += 1
  20. if self.use_features:
  21. netG_input_nc += opt.feat_num
  22. self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG,
  23. opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers,
  24. opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)
  25. self.load_network(self.netG, 'G', opt.which_epoch)
  26. print('---------- Networks initialized -------------')
  27. def toTensor(self, img, normalize=False):
  28. tensor = torch.from_numpy(np.array(img, np.int32, copy=False))
  29. tensor = tensor.view(1, img.size[1], img.size[0], len(img.mode))
  30. tensor = tensor.transpose(1, 2).transpose(1, 3).contiguous()
  31. if normalize:
  32. return (tensor.float()/255.0 - 0.5) / 0.5
  33. return tensor.float()
  34. def load_image(self, label_path, inst_path, feat_path):
  35. opt = self.opt
  36. # read label map
  37. label_img = Image.open(label_path)
  38. if label_path.find('face') != -1:
  39. label_img = label_img.convert('L')
  40. ow, oh = label_img.size
  41. w = opt.loadSize
  42. h = int(w * oh / ow)
  43. label_img = label_img.resize((w, h), Image.NEAREST)
  44. label_map = self.toTensor(label_img)
  45. # onehot vector input for label map
  46. self.label_map = label_map.cuda()
  47. oneHot_size = (1, opt.label_nc, h, w)
  48. input_label = self.Tensor(torch.Size(oneHot_size)).zero_()
  49. self.input_label = input_label.scatter_(1, label_map.long().cuda(), 1.0)
  50. # read instance map
  51. if not opt.no_instance:
  52. inst_img = Image.open(inst_path)
  53. inst_img = inst_img.resize((w, h), Image.NEAREST)
  54. self.inst_map = self.toTensor(inst_img).cuda()
  55. self.edge_map = self.get_edges(self.inst_map)
  56. self.net_input = Variable(torch.cat((self.input_label, self.edge_map), dim=1), volatile=True)
  57. else:
  58. self.net_input = Variable(self.input_label, volatile=True)
  59. self.features_clustered = np.load(feat_path).item()
  60. self.object_map = self.inst_map if opt.instance_feat else self.label_map
  61. object_np = self.object_map.cpu().numpy().astype(int)
  62. self.feat_map = self.Tensor(1, opt.feat_num, h, w).zero_()
  63. self.cluster_indices = np.zeros(self.opt.label_nc, np.uint8)
  64. for i in np.unique(object_np):
  65. label = i if i < 1000 else i//1000
  66. if label in self.features_clustered:
  67. feat = self.features_clustered[label]
  68. np.random.seed(i+1)
  69. cluster_idx = np.random.randint(0, feat.shape[0])
  70. self.cluster_indices[label] = cluster_idx
  71. idx = (self.object_map == i).nonzero()
  72. self.set_features(idx, feat, cluster_idx)
  73. self.net_input_original = self.net_input.clone()
  74. self.label_map_original = self.label_map.clone()
  75. self.feat_map_original = self.feat_map.clone()
  76. if not opt.no_instance:
  77. self.inst_map_original = self.inst_map.clone()
  78. def reset(self):
  79. self.net_input = self.net_input_prev = self.net_input_original.clone()
  80. self.label_map = self.label_map_prev = self.label_map_original.clone()
  81. self.feat_map = self.feat_map_prev = self.feat_map_original.clone()
  82. if not self.opt.no_instance:
  83. self.inst_map = self.inst_map_prev = self.inst_map_original.clone()
  84. self.object_map = self.inst_map if self.opt.instance_feat else self.label_map
  85. def undo(self):
  86. self.net_input = self.net_input_prev
  87. self.label_map = self.label_map_prev
  88. self.feat_map = self.feat_map_prev
  89. if not self.opt.no_instance:
  90. self.inst_map = self.inst_map_prev
  91. self.object_map = self.inst_map if self.opt.instance_feat else self.label_map
  92. # get boundary map from instance map
  93. def get_edges(self, t):
  94. edge = torch.cuda.ByteTensor(t.size()).zero_()
  95. edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
  96. edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
  97. edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
  98. edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
  99. return edge.float()
  100. # change the label at the source position to the label at the target position
  101. def change_labels(self, click_src, click_tgt):
  102. y_src, x_src = click_src[0], click_src[1]
  103. y_tgt, x_tgt = click_tgt[0], click_tgt[1]
  104. label_src = int(self.label_map[0, 0, y_src, x_src])
  105. inst_src = self.inst_map[0, 0, y_src, x_src]
  106. label_tgt = int(self.label_map[0, 0, y_tgt, x_tgt])
  107. inst_tgt = self.inst_map[0, 0, y_tgt, x_tgt]
  108. idx_src = (self.inst_map == inst_src).nonzero()
  109. # need to change 3 things: label map, instance map, and feature map
  110. if idx_src.shape:
  111. # backup current maps
  112. self.backup_current_state()
  113. # change both the label map and the network input
  114. self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
  115. self.net_input[idx_src[:,0], idx_src[:,1] + label_src, idx_src[:,2], idx_src[:,3]] = 0
  116. self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
  117. # update the instance map (and the network input)
  118. if inst_tgt > 1000:
  119. # if different instances have different ids, give the new object a new id
  120. tgt_indices = (self.inst_map > label_tgt * 1000) & (self.inst_map < (label_tgt+1) * 1000)
  121. inst_tgt = self.inst_map[tgt_indices].max() + 1
  122. self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = inst_tgt
  123. self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
  124. # also copy the source features to the target position
  125. idx_tgt = (self.inst_map == inst_tgt).nonzero()
  126. if idx_tgt.shape:
  127. self.copy_features(idx_src, idx_tgt[0,:])
  128. self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
  129. # add strokes of target label in the image
  130. def add_strokes(self, click_src, label_tgt, bw, save):
  131. # get the region of the new strokes (bw is the brush width)
  132. size = self.net_input.size()
  133. h, w = size[2], size[3]
  134. idx_src = torch.LongTensor(bw**2, 4).fill_(0)
  135. for i in range(bw):
  136. idx_src[i*bw:(i+1)*bw, 2] = min(h-1, max(0, click_src[0]-bw//2 + i))
  137. for j in range(bw):
  138. idx_src[i*bw+j, 3] = min(w-1, max(0, click_src[1]-bw//2 + j))
  139. idx_src = idx_src.cuda()
  140. # again, need to update 3 things
  141. if idx_src.shape:
  142. # backup current maps
  143. if save:
  144. self.backup_current_state()
  145. # update the label map (and the network input) in the stroke region
  146. self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
  147. for k in range(self.opt.label_nc):
  148. self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
  149. self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
  150. # update the instance map (and the network input)
  151. self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
  152. self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
  153. # also update the features if available
  154. if self.opt.instance_feat:
  155. feat = self.features_clustered[label_tgt]
  156. #np.random.seed(label_tgt+1)
  157. #cluster_idx = np.random.randint(0, feat.shape[0])
  158. cluster_idx = self.cluster_indices[label_tgt]
  159. self.set_features(idx_src, feat, cluster_idx)
  160. self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
  161. # add an object to the clicked position with selected style
  162. def add_objects(self, click_src, label_tgt, mask, style_id=0):
  163. y, x = click_src[0], click_src[1]
  164. mask = np.transpose(mask, (2, 0, 1))[np.newaxis,...]
  165. idx_src = torch.from_numpy(mask).cuda().nonzero()
  166. idx_src[:,2] += y
  167. idx_src[:,3] += x
  168. # backup current maps
  169. self.backup_current_state()
  170. # update label map
  171. self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
  172. for k in range(self.opt.label_nc):
  173. self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
  174. self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
  175. # update instance map
  176. self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
  177. self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
  178. # update feature map
  179. self.set_features(idx_src, self.feat, style_id)
  180. self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
  181. def single_forward(self, net_input, feat_map):
  182. net_input = torch.cat((net_input, feat_map), dim=1)
  183. fake_image = self.netG.forward(net_input)
  184. if fake_image.size()[0] == 1:
  185. return fake_image.data[0]
  186. return fake_image.data
  187. # generate all outputs for different styles
  188. def style_forward(self, click_pt, style_id=-1):
  189. if click_pt is None:
  190. self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
  191. self.crop = None
  192. self.mask = None
  193. else:
  194. instToChange = int(self.object_map[0, 0, click_pt[0], click_pt[1]])
  195. self.instToChange = instToChange
  196. label = instToChange if instToChange < 1000 else instToChange//1000
  197. self.feat = self.features_clustered[label]
  198. self.fake_image = []
  199. self.mask = self.object_map == instToChange
  200. idx = self.mask.nonzero()
  201. self.get_crop_region(idx)
  202. if idx.size():
  203. if style_id == -1:
  204. (min_y, min_x, max_y, max_x) = self.crop
  205. ### original
  206. for cluster_idx in range(self.opt.multiple_output):
  207. self.set_features(idx, self.feat, cluster_idx)
  208. fake_image = self.single_forward(self.net_input, self.feat_map)
  209. fake_image = util.tensor2im(fake_image[:,min_y:max_y,min_x:max_x])
  210. self.fake_image.append(fake_image)
  211. """### To speed up previewing different style results, either crop or downsample the label maps
  212. if instToChange > 1000:
  213. (min_y, min_x, max_y, max_x) = self.crop
  214. ### crop
  215. _, _, h, w = self.net_input.size()
  216. offset = 512
  217. y_start, x_start = max(0, min_y-offset), max(0, min_x-offset)
  218. y_end, x_end = min(h, (max_y + offset)), min(w, (max_x + offset))
  219. y_region = slice(y_start, y_start+(y_end-y_start)//16*16)
  220. x_region = slice(x_start, x_start+(x_end-x_start)//16*16)
  221. net_input = self.net_input[:,:,y_region,x_region]
  222. for cluster_idx in range(self.opt.multiple_output):
  223. self.set_features(idx, self.feat, cluster_idx)
  224. fake_image = self.single_forward(net_input, self.feat_map[:,:,y_region,x_region])
  225. fake_image = util.tensor2im(fake_image[:,min_y-y_start:max_y-y_start,min_x-x_start:max_x-x_start])
  226. self.fake_image.append(fake_image)
  227. else:
  228. ### downsample
  229. (min_y, min_x, max_y, max_x) = [crop//2 for crop in self.crop]
  230. net_input = self.net_input[:,:,::2,::2]
  231. size = net_input.size()
  232. net_input_batch = net_input.expand(self.opt.multiple_output, size[1], size[2], size[3])
  233. for cluster_idx in range(self.opt.multiple_output):
  234. self.set_features(idx, self.feat, cluster_idx)
  235. feat_map = self.feat_map[:,:,::2,::2]
  236. if cluster_idx == 0:
  237. feat_map_batch = feat_map
  238. else:
  239. feat_map_batch = torch.cat((feat_map_batch, feat_map), dim=0)
  240. fake_image_batch = self.single_forward(net_input_batch, feat_map_batch)
  241. for i in range(self.opt.multiple_output):
  242. self.fake_image.append(util.tensor2im(fake_image_batch[i,:,min_y:max_y,min_x:max_x]))"""
  243. else:
  244. self.set_features(idx, self.feat, style_id)
  245. self.cluster_indices[label] = style_id
  246. self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
  247. def backup_current_state(self):
  248. self.net_input_prev = self.net_input.clone()
  249. self.label_map_prev = self.label_map.clone()
  250. self.inst_map_prev = self.inst_map.clone()
  251. self.feat_map_prev = self.feat_map.clone()
  252. # crop the ROI and get the mask of the object
  253. def get_crop_region(self, idx):
  254. size = self.net_input.size()
  255. h, w = size[2], size[3]
  256. min_y, min_x = idx[:,2].min(), idx[:,3].min()
  257. max_y, max_x = idx[:,2].max(), idx[:,3].max()
  258. crop_min = 128
  259. if max_y - min_y < crop_min:
  260. min_y = max(0, (max_y + min_y) // 2 - crop_min // 2)
  261. max_y = min(h-1, min_y + crop_min)
  262. if max_x - min_x < crop_min:
  263. min_x = max(0, (max_x + min_x) // 2 - crop_min // 2)
  264. max_x = min(w-1, min_x + crop_min)
  265. self.crop = (min_y, min_x, max_y, max_x)
  266. self.mask = self.mask[:,:, min_y:max_y, min_x:max_x]
  267. # update the feature map once a new object is added or the label is changed
  268. def update_features(self, cluster_idx, mask=None, click_pt=None):
  269. self.feat_map_prev = self.feat_map.clone()
  270. # adding a new object
  271. if mask is not None:
  272. y, x = click_pt[0], click_pt[1]
  273. mask = np.transpose(mask, (2,0,1))[np.newaxis,...]
  274. idx = torch.from_numpy(mask).cuda().nonzero()
  275. idx[:,2] += y
  276. idx[:,3] += x
  277. # changing the label of an existing object
  278. else:
  279. idx = (self.object_map == self.instToChange).nonzero()
  280. # update feature map
  281. self.set_features(idx, self.feat, cluster_idx)
  282. # set the class features to the target feature
  283. def set_features(self, idx, feat, cluster_idx):
  284. for k in range(self.opt.feat_num):
  285. self.feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k]
  286. # copy the features at the target position to the source position
  287. def copy_features(self, idx_src, idx_tgt):
  288. for k in range(self.opt.feat_num):
  289. val = self.feat_map[idx_tgt[0], idx_tgt[1] + k, idx_tgt[2], idx_tgt[3]]
  290. self.feat_map[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = val
  291. def get_current_visuals(self, getLabel=False):
  292. mask = self.mask
  293. if self.mask is not None:
  294. mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8)
  295. dict_list = [('fake_image', self.fake_image), ('mask', mask)]
  296. if getLabel: # only output label map if needed to save bandwidth
  297. label = util.tensor2label(self.net_input.data[0], self.opt.label_nc)
  298. dict_list += [('label', label)]
  299. return OrderedDict(dict_list)