| 177 | return blends, masks, blends_colormap |
| 178 | |
| 179 | def blend_RANSAC(scene_images, marker, coarseModel=None, network=None, source=None, blend_type='D', use_colormap=False): |
| 180 | print('blend with ransac-flow') |
| 181 | blends = [] |
| 182 | masks = [] |
| 183 | blends_colormap = [] |
| 184 | for idx, scene in tqdm(enumerate(scene_images), total=len(scene_images)): |
| 185 | scene_ = Image.fromarray(scene) |
| 186 | marker_ = Image.fromarray(marker) |
| 187 | marker_ = marker_.resize(scene_.size) |
| 188 | coarseModel.setSource(marker_) |
| 189 | coarseModel.setTarget(scene_) |
| 190 | |
| 191 | I2w, I2h = coarseModel.It.size |
| 192 | featt = F.normalize(network['netFeatCoarse'](coarseModel.ItTensor)) |
| 193 | |
| 194 | #### -- grid |
| 195 | gridY = torch.linspace(-1, 1, steps = I2h).view(1, -1, 1, 1).expand(1, I2h, I2w, 1) |
| 196 | gridX = torch.linspace(-1, 1, steps = I2w).view(1, 1, -1, 1).expand(1, I2h, I2w, 1) |
| 197 | grid = torch.cat((gridX, gridY), dim=3).cuda() |
| 198 | warper = tgm.HomographyWarper(I2h, I2w) |
| 199 | |
| 200 | bestPara, InlierMask = coarseModel.getCoarse(np.zeros((I2h, I2w))) |
| 201 | bestPara = torch.from_numpy(bestPara).unsqueeze(0).cuda() |
| 202 | |
| 203 | flowCoarse = warper.warp_grid(bestPara) |
| 204 | I1_coarse = F.grid_sample(coarseModel.IsTensor, flowCoarse) |
| 205 | |
| 206 | featsSample = F.normalize(network['netFeatCoarse'](I1_coarse.cuda())) |
| 207 | |
| 208 | corr12 = network['netCorr'](featt, featsSample) |
| 209 | flowDown8 = network['netFlowCoarse'](corr12, False) ## output is with dimension B, 2, W, H |
| 210 | |
| 211 | flowUp = F.interpolate(flowDown8, size=(grid.size()[1], grid.size()[2]), mode='bilinear') |
| 212 | flowUp = flowUp.permute(0, 2, 3, 1) |
| 213 | |
| 214 | flowUp = flowUp + grid |
| 215 | |
| 216 | flow12 = F.grid_sample(flowCoarse.permute(0, 3, 1, 2), flowUp).permute(0, 2, 3, 1).contiguous() |
| 217 | |
| 218 | I1_fine = F.grid_sample(coarseModel.IsTensor, flow12) |
| 219 | I1_fine_pil = transforms.ToPILImage()(I1_fine.cpu().squeeze()) |
| 220 | |
| 221 | if blend_type == "mix": |
| 222 | source_id = 3 * int(idx / 3) |
| 223 | source = scene_images[source_id] |
| 224 | blend_i, mask_i = blend(np.array(I1_fine_pil), source, scene, blend_type) |
| 225 | blends.append(blend_i) |
| 226 | masks.append(mask_i) |
| 227 | |
| 228 | if use_colormap: |
| 229 | colormap = Image.open('./colormap.jpg') |
| 230 | coarseModel.setSource(colormap) |
| 231 | colormap_fine = F.grid_sample(coarseModel.IsTensor, flow12) |
| 232 | out_colormap = transforms.ToPILImage()(colormap_fine.cpu().squeeze()) |
| 233 | blend_colormap, _ = blend(np.array(out_colormap), source, scene, blend_type, use_colormap=use_colormap)#, mask=mask) |
| 234 | blends_colormap.append(blend_colormap) |
| 235 | |
| 236 | return blends, masks, blends_colormap |