| 299 | return blends, masks, blends_colormap |
| 300 | |
| 301 | def blend_SPSG(scene_images, marker, matching=None, source=None, blend_type='D', use_colormap=True): |
| 302 | print('blend with SPSG') |
| 303 | blends = [] |
| 304 | blends_colormap = [] |
| 305 | masks = [] |
| 306 | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| 307 | for scene in scene_images: |
| 308 | marker = cv2.resize(marker, (scene.shape[1], scene.shape[0])) |
| 309 | marker_gray = cv2.cvtColor(marker, cv2.COLOR_BGR2GRAY) |
| 310 | scene_gray = cv2.cvtColor(scene, cv2.COLOR_BGR2GRAY) |
| 311 | inp0 = frame2tensor(marker_gray, device=device) |
| 312 | inp1 = frame2tensor(scene_gray, device=device) |
| 313 | pred = matching({'image0': inp0, 'image1': inp1}) |
| 314 | pred = {k: v[0].cpu().numpy() for k, v in pred.items()} |
| 315 | kpts0, kpts1 = pred['keypoints0'], pred['keypoints1'] |
| 316 | matches, conf = pred['matches0'], pred['matching_scores0'] |
| 317 | |
| 318 | valid = matches > -1 |
| 319 | mkpts0 = kpts0[valid] |
| 320 | mkpts1 = kpts1[matches[valid]] |
| 321 | mconf = conf[valid] |
| 322 | |
| 323 | valid = mconf > 0.5 |
| 324 | mkpts0 = mkpts0[valid] |
| 325 | mkpts1 = mkpts1[valid] |
| 326 | |
| 327 | if np.count_nonzero(valid) < 4: |
| 328 | blends.append(None) |
| 329 | masks.append(None) |
| 330 | blends_colormap.append(None) |
| 331 | continue |
| 332 | |
| 333 | src_pts = np.float32(mkpts0).reshape(-1,1,2) |
| 334 | dst_pts = np.float32(mkpts1).reshape(-1,1,2) |
| 335 | M, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) |
| 336 | if M is None: |
| 337 | blends.append(None) |
| 338 | masks.append(None) |
| 339 | blends_colormap.append(None) |
| 340 | continue |
| 341 | out = cv2.warpPerspective(marker, M, (scene.shape[1], scene.shape[0])) |
| 342 | |
| 343 | blend_i, mask_i = blend(out, source, scene, blend_type) |
| 344 | blends.append(blend_i) |
| 345 | masks.append(mask_i) |
| 346 | |
| 347 | if use_colormap: |
| 348 | colormap = np.asarray(Image.open('./colormap.jpg')) |
| 349 | colormap = cv2.resize(colormap[:,:,::-1], (scene.shape[1], scene.shape[0])) |
| 350 | blend_colormap = cv2.warpPerspective(colormap, M, (scene.shape[1], scene.shape[0])) |
| 351 | blend_colormap, _ = blend(blend_colormap, source, scene, blend_type) |
| 352 | blends_colormap.append(blend_colormap) |
| 353 | return blends, masks, blends_colormap |
| 354 | |
| 355 | def eval(args, id, scene_images, marker, source=None, |
| 356 | estimator=None, |