(stylization_module, smoothing_module, content_image_path, style_image_path, content_seg_path,
style_seg_path, output_image_path,
cuda, save_intermediate, no_post, label_remapping, output_visualization=False)
| 122 | |
| 123 | |
| 124 | def stylization(stylization_module, smoothing_module, content_image_path, style_image_path, content_seg_path, |
| 125 | style_seg_path, output_image_path, |
| 126 | cuda, save_intermediate, no_post, label_remapping, output_visualization=False): |
| 127 | # Load image |
| 128 | with torch.no_grad(): |
| 129 | cont_img = Image.open(content_image_path).convert('RGB') |
| 130 | styl_img = Image.open(style_image_path).convert('RGB') |
| 131 | |
| 132 | new_cw, new_ch = memory_limit_image_resize(cont_img) |
| 133 | new_sw, new_sh = memory_limit_image_resize(styl_img) |
| 134 | cont_pilimg = cont_img.copy() |
| 135 | styl_pilimg = styl_img.copy() |
| 136 | cw = cont_pilimg.width |
| 137 | ch = cont_pilimg.height |
| 138 | try: |
| 139 | cont_seg = Image.open(content_seg_path) |
| 140 | styl_seg = Image.open(style_seg_path) |
| 141 | cont_seg.resize((new_cw, new_ch), Image.NEAREST) |
| 142 | styl_seg.resize((new_sw, new_sh), Image.NEAREST) |
| 143 | |
| 144 | except: |
| 145 | cont_seg = [] |
| 146 | styl_seg = [] |
| 147 | |
| 148 | cont_img = transforms.ToTensor()(cont_img).unsqueeze(0) |
| 149 | styl_img = transforms.ToTensor()(styl_img).unsqueeze(0) |
| 150 | |
| 151 | if cuda: |
| 152 | cont_img = cont_img.cuda(0) |
| 153 | styl_img = styl_img.cuda(0) |
| 154 | stylization_module.cuda(0) |
| 155 | |
| 156 | # cont_img = Variable(cont_img, volatile=True) |
| 157 | # styl_img = Variable(styl_img, volatile=True) |
| 158 | |
| 159 | cont_seg = np.asarray(cont_seg) |
| 160 | styl_seg = np.asarray(styl_seg) |
| 161 | |
| 162 | cont_seg = label_remapping.self_remapping(cont_seg) |
| 163 | styl_seg = label_remapping.self_remapping(styl_seg) |
| 164 | cont_seg, styl_seg = label_remapping.cross_remapping(cont_seg, styl_seg) |
| 165 | |
| 166 | if output_visualization: |
| 167 | import cv2 |
| 168 | cont_seg_vis = visualize_result(cont_seg) |
| 169 | styl_seg_vis = visualize_result(styl_seg) |
| 170 | cont_seg_vis = overlay(cv2.imread(content_image_path), cont_seg_vis) |
| 171 | styl_seg_vis = overlay(cv2.imread(style_image_path), styl_seg_vis) |
| 172 | cv2.imwrite(content_seg_path + '.visualization.jpg', cont_seg_vis) |
| 173 | cv2.imwrite(style_seg_path + '.visualization.jpg', styl_seg_vis) |
| 174 | |
| 175 | if save_intermediate: |
| 176 | with Timer("Elapsed time in stylization: %f"): |
| 177 | stylized_img = stylization_module.transform(cont_img, styl_img, cont_seg, styl_seg) |
| 178 | if ch != new_ch or cw != new_cw: |
| 179 | print("De-resize image: (%d,%d)->(%d,%d)" % (new_cw, new_ch, cw, ch)) |
| 180 | stylized_img = nn.functional.upsample(stylized_img, size=(ch, cw), mode='bilinear') |
| 181 | utils.save_image(stylized_img.data.cpu().float(), output_image_path, nrow=1, padding=0) |
nothing calls this directly
no test coverage detected