Cartoonify an image.
(input_image, style_id=0, load_size=450)
| 41 | |
| 42 | |
| 43 | def cartoonify(input_image, style_id=0, load_size=450): |
| 44 | """Cartoonify an image.""" |
| 45 | style = styles[int(style_id)] |
| 46 | model = models[style] |
| 47 | load_size = int(load_size) |
| 48 | |
| 49 | # resize image, keep aspect ratio |
| 50 | h = input_image.size[0] |
| 51 | w = input_image.size[1] |
| 52 | ratio = h * 1.0 / w |
| 53 | if ratio > 1: |
| 54 | h = load_size |
| 55 | w = int(h * 1.0 / ratio) |
| 56 | else: |
| 57 | w = load_size |
| 58 | h = int(w * ratio) |
| 59 | input_image = input_image.resize((h, w), Image.BICUBIC) |
| 60 | |
| 61 | input_image = np.asarray(input_image) |
| 62 | # RGB -> BGR |
| 63 | input_image = input_image[:, :, [2, 1, 0]] |
| 64 | input_image = transforms.ToTensor()(input_image).unsqueeze(0) |
| 65 | |
| 66 | # preprocess, (-1, 1) |
| 67 | input_image = -1 + 2 * input_image |
| 68 | if gpu: |
| 69 | input_image = Variable(input_image).cuda() |
| 70 | else: |
| 71 | input_image = Variable(input_image).float() |
| 72 | |
| 73 | # forward |
| 74 | with torch.no_grad(): |
| 75 | output_image = model(input_image) |
| 76 | output_image = output_image[0] |
| 77 | |
| 78 | # BGR -> RGB |
| 79 | output_image = output_image[[2, 1, 0], :, :] |
| 80 | |
| 81 | # deprocess, (0, 1) |
| 82 | output_image = output_image.data.cpu().float() * 0.5 + 0.5 |
| 83 | |
| 84 | output_image = output_image.numpy() |
| 85 | output_image = np.uint8(output_image.transpose(1, 2, 0) * 255) |
| 86 | output_image = Image.fromarray(output_image) |
| 87 | |
| 88 | return output_image |
| 89 |
no outgoing calls