(stylization_module, smoothing_module, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path,
cuda, save_intermediate, no_post, cont_seg_remapping=None, styl_seg_remapping=None)
| 58 | |
| 59 | |
| 60 | def stylization(stylization_module, smoothing_module, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path, |
| 61 | cuda, save_intermediate, no_post, cont_seg_remapping=None, styl_seg_remapping=None): |
| 62 | # Load image |
| 63 | with torch.no_grad(): |
| 64 | cont_img = Image.open(content_image_path).convert('RGB') |
| 65 | styl_img = Image.open(style_image_path).convert('RGB') |
| 66 | |
| 67 | new_cw, new_ch = memory_limit_image_resize(cont_img) |
| 68 | new_sw, new_sh = memory_limit_image_resize(styl_img) |
| 69 | cont_pilimg = cont_img.copy() |
| 70 | cw = cont_pilimg.width |
| 71 | ch = cont_pilimg.height |
| 72 | try: |
| 73 | cont_seg = Image.open(content_seg_path) |
| 74 | styl_seg = Image.open(style_seg_path) |
| 75 | cont_seg.resize((new_cw,new_ch),Image.NEAREST) |
| 76 | styl_seg.resize((new_sw,new_sh),Image.NEAREST) |
| 77 | |
| 78 | except: |
| 79 | cont_seg = [] |
| 80 | styl_seg = [] |
| 81 | |
| 82 | cont_img = transforms.ToTensor()(cont_img).unsqueeze(0) |
| 83 | styl_img = transforms.ToTensor()(styl_img).unsqueeze(0) |
| 84 | |
| 85 | if cuda: |
| 86 | cont_img = cont_img.cuda(0) |
| 87 | styl_img = styl_img.cuda(0) |
| 88 | stylization_module.cuda(0) |
| 89 | |
| 90 | # cont_img = Variable(cont_img, volatile=True) |
| 91 | # styl_img = Variable(styl_img, volatile=True) |
| 92 | |
| 93 | cont_seg = np.asarray(cont_seg) |
| 94 | styl_seg = np.asarray(styl_seg) |
| 95 | if cont_seg_remapping is not None: |
| 96 | cont_seg = cont_seg_remapping.process(cont_seg) |
| 97 | if styl_seg_remapping is not None: |
| 98 | styl_seg = styl_seg_remapping.process(styl_seg) |
| 99 | |
| 100 | if save_intermediate: |
| 101 | with Timer("Elapsed time in stylization: %f"): |
| 102 | stylized_img = stylization_module.transform(cont_img, styl_img, cont_seg, styl_seg) |
| 103 | if ch != new_ch or cw != new_cw: |
| 104 | print("De-resize image: (%d,%d)->(%d,%d)" %(new_cw,new_ch,cw,ch)) |
| 105 | stylized_img = nn.functional.upsample(stylized_img, size=(ch,cw), mode='bilinear') |
| 106 | utils.save_image(stylized_img.data.cpu().float(), output_image_path, nrow=1, padding=0) |
| 107 | |
| 108 | with Timer("Elapsed time in propagation: %f"): |
| 109 | out_img = smoothing_module.process(output_image_path, content_image_path) |
| 110 | out_img.save(output_image_path) |
| 111 | |
| 112 | if not cuda: |
| 113 | print("NotImplemented: The CPU version of smooth filter has not been implemented currently.") |
| 114 | return |
| 115 | |
| 116 | if no_post is False: |
| 117 | with Timer("Elapsed time in post processing: %f"): |
nothing calls this directly
no test coverage detected