(src_video, src_mask, src_ref_images, num_frames, image_size,
device)
| 272 | |
| 273 | |
| 274 | def prepare_source(src_video, src_mask, src_ref_images, num_frames, image_size, |
| 275 | device): |
| 276 | for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)): |
| 277 | if sub_src_video is None and sub_src_mask is None: |
| 278 | src_video[i] = torch.zeros( |
| 279 | (3, num_frames, image_size[0], image_size[1]), device=device) |
| 280 | src_mask[i] = torch.ones( |
| 281 | (1, num_frames, image_size[0], image_size[1]), device=device) |
| 282 | for i, ref_images in enumerate(src_ref_images): |
| 283 | if ref_images is not None: |
| 284 | for j, ref_img in enumerate(ref_images): |
| 285 | if ref_img is not None and ref_img.shape[-2:] != image_size: |
| 286 | canvas_height, canvas_width = image_size |
| 287 | ref_height, ref_width = ref_img.shape[-2:] |
| 288 | white_canvas = torch.ones( |
| 289 | (3, 1, canvas_height, canvas_width), |
| 290 | device=device) # [-1, 1] |
| 291 | scale = min(canvas_height / ref_height, |
| 292 | canvas_width / ref_width) |
| 293 | new_height = int(ref_height * scale) |
| 294 | new_width = int(ref_width * scale) |
| 295 | resized_image = F.interpolate( |
| 296 | ref_img.squeeze(1).unsqueeze(0), |
| 297 | size=(new_height, new_width), |
| 298 | mode='bilinear', |
| 299 | align_corners=False).squeeze(0).unsqueeze(1) |
| 300 | top = (canvas_height - new_height) // 2 |
| 301 | left = (canvas_width - new_width) // 2 |
| 302 | white_canvas[:, :, top:top + new_height, |
| 303 | left:left + new_width] = resized_image |
| 304 | src_ref_images[i][j] = white_canvas |
| 305 | return src_video, src_mask, src_ref_images |
nothing calls this directly
no outgoing calls
no test coverage detected