myshop-capsule/image-super-resolution
6
1# install thing, just like in segment anything2 3 4from typing import Dict, List, Any5from PIL import Image6from io import BytesIO7from transformers import AutoModelForSemanticSegmentation, AutoFeatureExtractor8import base649import torch10from torch import nn11 12 13# import subprocess14# result = subprocess.run(["pip", "install", "git+https://github.com/sberbank-ai/Real-ESRGAN.git"], check=True)15# print(f"git+https://github.com/sberbank-ai/Real-ESRGAN.git = {result}")16# from RealESRGAN import RealESRGAN17 18# no need to install, just take in all of the necessary files from the notebook19import math20import torch21from torch import nn as nn22from torch.nn import functional as F23from torch.nn import init as init24from torch.nn.modules.batchnorm import _BatchNorm25 26@torch.no_grad()27def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):28 """Initialize network weights.29 30 Args:31 module_list (list[nn.Module] | nn.Module): Modules to be initialized.32 scale (float): Scale initialized weights, especially for residual33 blocks. Default: 1.34 bias_fill (float): The value to fill bias. Default: 035 kwargs (dict): Other arguments for initialization function.36 """37 if not isinstance(module_list, list):38 module_list = [module_list]39 for module in module_list:40 for m in module.modules():41 if isinstance(m, nn.Conv2d):42 init.kaiming_normal_(m.weight, **kwargs)43 m.weight.data *= scale44 if m.bias is not None:45 m.bias.data.fill_(bias_fill)46 elif isinstance(m, nn.Linear):47 init.kaiming_normal_(m.weight, **kwargs)48 m.weight.data *= scale49 if m.bias is not None:50 m.bias.data.fill_(bias_fill)51 elif isinstance(m, _BatchNorm):52 init.constant_(m.weight, 1)53 if m.bias is not None:54 m.bias.data.fill_(bias_fill)55 56 57def make_layer(basic_block, num_basic_block, **kwarg):58 """Make layers by stacking the same blocks.59 60 Args:61 basic_block (nn.module): nn.module class for basic block.62 num_basic_block (int): number of blocks.63 64 Returns:65 nn.Sequential: Stacked blocks in nn.Sequential.66 """67 layers = []68 for _ in range(num_basic_block):69 layers.append(basic_block(**kwarg))70 return nn.Sequential(*layers)71 72 73 74class ResidualBlockNoBN(nn.Module):75 """Residual block without BN.76 77 It has a style of:78 ---Conv-ReLU-Conv-+-79 |________________|80 81 Args:82 num_feat (int): Channel number of intermediate features.83 Default: 64.84 res_scale (float): Residual scale. Default: 1.85 pytorch_init (bool): If set to True, use pytorch default init,86 otherwise, use default_init_weights. Default: False.87 """88 89 def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):90 super(ResidualBlockNoBN, self).__init__()91 self.res_scale = res_scale92 self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)93 self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)94 self.relu = nn.ReLU(inplace=True)95 96 if not pytorch_init:97 default_init_weights([self.conv1, self.conv2], 0.1)98 99 def forward(self, x):100 identity = x101 out = self.conv2(self.relu(self.conv1(x)))102 return identity + out * self.res_scale103 104 105 106class Upsample(nn.Sequential):107 """Upsample module.108 109 Args:110 scale (int): Scale factor. Supported scales: 2^n and 3.111 num_feat (int): Channel number of intermediate features.112 """113 114 def __init__(self, scale, num_feat):115 m = []116 if (scale & (scale - 1)) == 0: # scale = 2^n117 for _ in range(int(math.log(scale, 2))):118 m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))119 m.append(nn.PixelShuffle(2))120 elif scale == 3:121 m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))122 m.append(nn.PixelShuffle(3))123 else:124 raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')125 super(Upsample, self).__init__(*m)126 127 128def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):129 """Warp an image or feature map with optical flow.130 131 Args:132 x (Tensor): Tensor with size (n, c, h, w).133 flow (Tensor): Tensor with size (n, h, w, 2), normal value.134 interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.135 padding_mode (str): 'zeros' or 'border' or 'reflection'.136 Default: 'zeros'.137 align_corners (bool): Before pytorch 1.3, the default value is138 align_corners=True. After pytorch 1.3, the default value is139 align_corners=False. Here, we use the True as default.140 141 Returns:142 Tensor: Warped image or feature map.143 """144 assert x.size()[-2:] == flow.size()[1:3]145 _, _, h, w = x.size()146 # create mesh grid147 grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))148 grid = torch.stack((grid_x, grid_y), 2).float() # W(x), H(y), 2149 grid.requires_grad = False150 151 vgrid = grid + flow152 # scale grid to [-1,1]153 vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0154 vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0155 vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)156 output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)157 158 # TODO, what if align_corners=False159 return output160 161 162def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):163 """Resize a flow according to ratio or shape.164 165 Args:166 flow (Tensor): Precomputed flow. shape [N, 2, H, W].167 size_type (str): 'ratio' or 'shape'.168 sizes (list[int | float]): the ratio for resizing or the final output169 shape.170 1) The order of ratio should be [ratio_h, ratio_w]. For171 downsampling, the ratio should be smaller than 1.0 (i.e., ratio172 < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,173 ratio > 1.0).174 2) The order of output_size should be [out_h, out_w].175 interp_mode (str): The mode of interpolation for resizing.176 Default: 'bilinear'.177 align_corners (bool): Whether align corners. Default: False.178 179 Returns:180 Tensor: Resized flow.181 """182 _, _, flow_h, flow_w = flow.size()183 if size_type == 'ratio':184 output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])185 elif size_type == 'shape':186 output_h, output_w = sizes[0], sizes[1]187 else:188 raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')189 190 input_flow = flow.clone()191 ratio_h = output_h / flow_h192 ratio_w = output_w / flow_w193 input_flow[:, 0, :, :] *= ratio_w194 input_flow[:, 1, :, :] *= ratio_h195 resized_flow = F.interpolate(196 input=input_flow, size=(output_h, output_w), mode=interp_mode, align_corners=align_corners)197 return resized_flow198 199 200# TODO: may write a cpp file201def pixel_unshuffle(x, scale):202 """ Pixel unshuffle.203 204 Args:205 x (Tensor): Input feature with shape (b, c, hh, hw).206 scale (int): Downsample ratio.207 208 Returns:209 Tensor: the pixel unshuffled feature.210 """211 print('PIXEL UNSHUFFLE X SIZE', x.size())212 output = []213 # new batch size for it here214 b, c, hh, hw = x.size()215 216 # okay ugh, what is this all doing ...217 # i mean you could concat each of those in a llok218 out_channel = c * (scale**2)219 assert hh % scale == 0 and hw % scale == 0220 h = hh // scale221 w = hw // scale222 x_view = x.view(b, c, h, scale, w, scale)223 x_view = x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)224 225 # output = torch.stack(output)226 # print('output shape', x_view.shape)227 # 1/0228 return x_view229 230 231import os232import torch233from torch.nn import functional as F234from PIL import Image235import numpy as np236from huggingface_hub import hf_hub_url, cached_download237 238 239HF_MODELS = {240 2: dict(241 repo_id='sberbank-ai/Real-ESRGAN',242 filename='RealESRGAN_x2.pth',243 ),244 4: dict(245 repo_id='sberbank-ai/Real-ESRGAN',246 filename='RealESRGAN_x4.pth',247 ),248 8: dict(249 repo_id='sberbank-ai/Real-ESRGAN',250 filename='RealESRGAN_x8.pth',251 ),252}253 254 255class RealESRGAN:256 def __init__(self, device, scale=4):257 self.device = device258 self.scale = scale259 self.model = RRDBNet(260 num_in_ch=3, num_out_ch=3, num_feat=64,261 num_block=23, num_grow_ch=32, scale=scale262 )263 264 def load_weights(self, model_path, download=True):265 if not os.path.exists(model_path) and download:266 assert self.scale in [2,4,8], 'You can download models only with scales: 2, 4, 8'267 config = HF_MODELS[self.scale]268 cache_dir = os.path.dirname(model_path)269 local_filename = os.path.basename(model_path)270 config_file_url = hf_hub_url(repo_id=config['repo_id'], filename=config['filename'])271 cached_download(config_file_url, cache_dir=cache_dir, force_filename=local_filename)272 print('Weights downloaded to:', os.path.join(cache_dir, local_filename))273 274 loadnet = torch.load(model_path)275 if 'params' in loadnet:276 self.model.load_state_dict(loadnet['params'], strict=True)277 elif 'params_ema' in loadnet:278 self.model.load_state_dict(loadnet['params_ema'], strict=True)279 else:280 self.model.load_state_dict(loadnet, strict=True)281 self.model.eval()282 self.model.to(self.device)283 284 @torch.cuda.amp.autocast()285 def predict(self, numpy_images, batch_size=4, patches_size=192,286 padding=24, pad_size=15):287 import time288 start = time.time()289 # okay i think that's good with variability for now ... 290 # ***IMPORTANT VARIABLE***291 batch_size = len(numpy_images) * 4292 scale = self.scale293 device = self.device294 295 list_of_inputs = []296 for lr_image in numpy_images:297 lr_image = np.array(lr_image)298 lr_image = pad_reflect(lr_image, pad_size)299 300 patches, p_shape = split_image_into_overlapping_patches(301 lr_image, patch_size=patches_size, padding_size=padding302 )303 304 # print('patches.shape', patches.shape)305 # print('p_shape', p_shape)306 307 img = torch.FloatTensor(patches/255).permute((0,3,1,2)).to(device).detach()308 list_of_inputs.append(img)309 310 311 input_batch = torch.concat(list_of_inputs)312 313 # print('input_batch.shape', input_batch.shape)314 315 start2 = time.time()316 with torch.no_grad():317 # res = self.model(input_batch[0:batch_size])318 319 # okay what does the input size really need to be?320 321 # print('input_batch.shape', input_batch.shape)322 # print('input_batch[0:batch_size].shape', input_batch[0:batch_size].shape)323 # 1/0324 res = self.model(input_batch[0:batch_size])325 326 # print('res.shape 1', res.shape)327 # print('batch_size', batch_size)328 # 1/0329 for i in range(batch_size, img.shape[0], batch_size):330 print('i is', i)331 res = torch.cat((res, self.model(img[i:i+batch_size])), 0)332 # print('res.shape 2', res.shape)333 print('inference alone takes', time.time() - start2)334 # print('res.shape 3', res.shape)335 336 # 1/0337 338 sr_image = res.permute((0,2,3,1)).clamp_(0, 1).cpu()339 np_sr_image_batch = sr_image.numpy()340 341 # print('np_sr_image_batch.shape', np_sr_image_batch.shape)342 # print('np_sr_image_batch[0].shape', np_sr_image_batch[0].shape)343 # 1/0344 345 padded_size_scaled = tuple(np.multiply(p_shape[0:2], scale)) + (3,)346 347 output_images = []348 for i in range(0,batch_size,4):349 # get first time from original input image size350 scaled_image_shape = tuple(np.multiply(lr_image.shape[0:2], scale)) + (3,)351 # print('scaled_image_shape', scaled_image_shape)352 # print('padded_size_scaled', padded_size_scaled)353 # print("padding * scale", padding * scale)354 np_sr_image = stich_together(355 np_sr_image_batch[i:i+4], padded_image_shape=padded_size_scaled,356 target_shape=scaled_image_shape, padding_size=padding * scale357 )358 sr_img = (np_sr_image*255).astype(np.uint8)359 # print('sr_img.shape', sr_img.shape)360 sr_img = unpad_image(sr_img, pad_size*scale)361 sr_img = Image.fromarray(sr_img)362 output_images.append(sr_img)363 364 print('len of output_images', len(output_images))365 366 # for debugging367 # for idx, image in enumerate(output_images):368 # image.save(f'output_image_{idx}.png')369 370 371 print("EVERYTHING TOOK", time.time() - start)372 373 return output_images374 375 376import torch377from torch import nn as nn378from torch.nn import functional as F379 380 381class ResidualDenseBlock(nn.Module):382 """Residual Dense Block.383 384 Used in RRDB block in ESRGAN.385 386 Args:387 num_feat (int): Channel number of intermediate features.388 num_grow_ch (int): Channels for each growth.389 """390 391 def __init__(self, num_feat=64, num_grow_ch=32):392 super(ResidualDenseBlock, self).__init__()393 self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)394 self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)395 self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)396 self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)397 self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)398 399 self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)400 401 # initialization402 default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)403 404 def forward(self, x):405 x1 = self.lrelu(self.conv1(x))406 x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))407 x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))408 x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))409 x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))410 # Emperically, we use 0.2 to scale the residual for better performance411 return x5 * 0.2 + x412 413 414class RRDB(nn.Module):415 """Residual in Residual Dense Block.416 417 Used in RRDB-Net in ESRGAN.418 419 Args:420 num_feat (int): Channel number of intermediate features.421 num_grow_ch (int): Channels for each growth.422 """423 424 def __init__(self, num_feat, num_grow_ch=32):425 super(RRDB, self).__init__()426 self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)427 self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)428 self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)429 430 def forward(self, x):431 # this part happens 23 times per pass 432 out = self.rdb1(x)433 out = self.rdb2(out)434 out = self.rdb3(out)435 # Emperically, we use 0.2 to scale the residual for better performance436 return out * 0.2 + x437 438 439class RRDBNet(nn.Module):440 """Networks consisting of Residual in Residual Dense Block, which is used441 in ESRGAN.442 443 ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.444 445 We extend ESRGAN for scale x2 and scale x1.446 Note: This is one option for scale 1, scale 2 in RRDBNet.447 We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size448 and enlarge the channel size before feeding inputs into the main ESRGAN architecture.449 450 Args:451 num_in_ch (int): Channel number of inputs.452 num_out_ch (int): Channel number of outputs.453 num_feat (int): Channel number of intermediate features.454 Default: 64455 num_block (int): Block number in the trunk network. Defaults: 23456 num_grow_ch (int): Channels for each growth. Default: 32.457 """458 459 def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):460 super(RRDBNet, self).__init__()461 462 self.scale = scale463 if scale == 2:464 num_in_ch = num_in_ch * 4465 elif scale == 1:466 num_in_ch = num_in_ch * 16467 468 print('num_in_ch', num_in_ch)469 470 self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)471 self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)472 self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)473 # upsample474 self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)475 self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)476 if scale == 8:477 self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)478 self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)479 self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)480 481 self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)482 483 def forward(self, x):484 # print('IN FORWARD, X.shape is', x.shape)485 if self.scale == 2:486 feat = pixel_unshuffle(x, scale=2)487 elif self.scale == 1:488 feat = pixel_unshuffle(x, scale=4)489 else:490 feat = x491 # print('feat shape', feat.shape)492 # breaks here ...493 feat = self.conv_first(feat)494 body_feat = self.conv_body(self.body(feat))495 feat = feat + body_feat496 # upsample497 feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))498 feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))499 if self.scale == 8:500 feat = self.lrelu(self.conv_up3(F.interpolate(feat, scale_factor=2, mode='nearest')))501 out = self.conv_last(self.lrelu(self.conv_hr(feat)))502 return out503 504import numpy as np505import torch506from PIL import Image507import os508import io509 510def pad_reflect(image, pad_size):511 imsize = image.shape512 height, width = imsize[:2]513 print('imsize', imsize)514 new_img = np.zeros([height+pad_size*2, width+pad_size*2, imsize[2]]).astype(np.uint8)515 new_img[pad_size:-pad_size, pad_size:-pad_size, :] = image516 # print('new_img.shape 1', new_img.shape)517 518 new_img[0:pad_size, pad_size:-pad_size, :] = np.flip(image[0:pad_size, :, :], axis=0) #top519 new_img[-pad_size:, pad_size:-pad_size, :] = np.flip(image[-pad_size:, :, :], axis=0) #bottom520 new_img[:, 0:pad_size, :] = np.flip(new_img[:, pad_size:pad_size*2, :], axis=1) #left521 new_img[:, -pad_size:, :] = np.flip(new_img[:, -pad_size*2:-pad_size, :], axis=1) #right522 # print('new_img.shape 2', new_img.shape)523 524 return new_img525 526def unpad_image(image, pad_size):527 return image[pad_size:-pad_size, pad_size:-pad_size, :]528 529 530def process_array(image_array, expand=True):531 """ Process a 3-dimensional array into a scaled, 4 dimensional batch of size 1. """532 533 image_batch = image_array / 255.0534 if expand:535 image_batch = np.expand_dims(image_batch, axis=0)536 return image_batch537 538 539def process_output(output_tensor):540 """ Transforms the 4-dimensional output tensor into a suitable image format. """541 542 sr_img = output_tensor.clip(0, 1) * 255543 sr_img = np.uint8(sr_img)544 return sr_img545 546 547def pad_patch(image_patch, padding_size, channel_last=True):548 """ Pads image_patch with with padding_size edge values. """549 550 if channel_last:551 return np.pad(552 image_patch,553 ((padding_size, padding_size), (padding_size, padding_size), (0, 0)),554 'edge',555 )556 else:557 return np.pad(558 image_patch,559 ((0, 0), (padding_size, padding_size), (padding_size, padding_size)),560 'edge',561 )562 563 564def unpad_patches(image_patches, padding_size):565 return image_patches[:, padding_size:-padding_size, padding_size:-padding_size, :]566 567 568def split_image_into_overlapping_patches(image_array, patch_size, padding_size=2):569 """ Splits the image into partially overlapping patches.570 The patches overlap by padding_size pixels.571 Pads the image twice:572 - first to have a size multiple of the patch size,573 - then to have equal padding at the borders.574 Args:575 image_array: numpy array of the input image.576 patch_size: size of the patches from the original image (without padding).577 padding_size: size of the overlapping area.578 """579 580 xmax, ymax, _ = image_array.shape581 x_remainder = xmax % patch_size582 y_remainder = ymax % patch_size583 584 # modulo here is to avoid extending of patch_size instead of 0585 x_extend = (patch_size - x_remainder) % patch_size586 y_extend = (patch_size - y_remainder) % patch_size587 588 # make sure the image is divisible into regular patches589 extended_image = np.pad(image_array, ((0, x_extend), (0, y_extend), (0, 0)), 'edge')590 591 # add padding around the image to simplify computations592 padded_image = pad_patch(extended_image, padding_size, channel_last=True)593 594 xmax, ymax, _ = padded_image.shape595 patches = []596 597 x_lefts = range(padding_size, xmax - padding_size, patch_size)598 y_tops = range(padding_size, ymax - padding_size, patch_size)599 600 for x in x_lefts:601 for y in y_tops:602 x_left = x - padding_size603 y_top = y - padding_size604 x_right = x + patch_size + padding_size605 y_bottom = y + patch_size + padding_size606 patch = padded_image[x_left:x_right, y_top:y_bottom, :]607 patches.append(patch)608 609 return np.array(patches), padded_image.shape610 611 612def stich_together(patches, padded_image_shape, target_shape, padding_size=4):613 """ Reconstruct the image from overlapping patches.614 After scaling, shapes and padding should be scaled too.615 Args:616 patches: patches obtained with split_image_into_overlapping_patches617 padded_image_shape: shape of the padded image contructed in split_image_into_overlapping_patches618 target_shape: shape of the final image619 padding_size: size of the overlapping area.620 """621 622 xmax, ymax, _ = padded_image_shape623 patches = unpad_patches(patches, padding_size)624 patch_size = patches.shape[1]625 n_patches_per_row = ymax // patch_size626 627 complete_image = np.zeros((xmax, ymax, 3))628 629 row = -1630 col = 0631 for i in range(len(patches)):632 if i % n_patches_per_row == 0:633 row += 1634 col = 0635 complete_image[636 row * patch_size: (row + 1) * patch_size, col * patch_size: (col + 1) * patch_size,:637 ] = patches[i]638 col += 1639 return complete_image[0: target_shape[0], 0: target_shape[1], :]640 641 642 643class EndpointHandler():644 def __init__(self, path="."):645 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")646 self.model = RealESRGAN(self.device, scale=2)647 self.model.load_weights('/repository/RealESRGAN_x2.pth', download=True) 648 649 650 def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:651 """652 data args:653 images (:obj:`PIL.Image`)654 candiates (:obj:`list`)655 Return:656 A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}657 """658 inputs = data.pop("inputs", data)659 if isinstance(inputs['image'], list):660 input_images = []661 for base64_string in inputs['image']:662 image = Image.open(BytesIO(base64.b64decode(base64_string)))663 input_images.append(image)664 665 for i in range(len(input_images)):666 input_images[i] = input_images[i].resize((224, 224))667 668 numpy_images = [np.array(img) for img in input_images]669 output_images = self.model.predict(numpy_images)670 671 base64_strings = []672 for output_image in output_images:673 buffered = BytesIO()674 output_image = output_image.convert('RGB')675 output_image.save(buffered, format="png")676 img_str = base64.b64encode(buffered.getvalue())677 base64_strings.append(img_str.decode('utf-8'))678 679 return base64_strings680 681 else:682 # decode base64 image to PIL683 image = Image.open(BytesIO(base64.b64decode(inputs['image'])))684 685 # forward pass686 output_image = self.model.predict([image])687 688 if isinstance(output_image, list):689 output_image = output_image[0]690 691 # base64 encode output692 buffered = BytesIO()693 output_image = output_image.convert('RGB')694 output_image.save(buffered, format="png")695 img_str = base64.b64encode(buffered.getvalue())696 697 # postprocess the prediction698 return {"image": img_str.decode()}