Spaces:
Running
Running
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import tqdm | |
| import yaml | |
| from omegaconf import OmegaConf | |
| from PIL import Image | |
| from torch.utils.data._utils.collate import default_collate | |
| from saicinpainting.training.trainers import load_checkpoint | |
| from saicinpainting.evaluation.utils import move_to_device, load_image, prepare_image, pad_img_to_modulo, scale_image | |
| from saicinpainting.evaluation.refinement import refine_predict | |
| refiner_config = { | |
| 'gpu_ids': '0,', | |
| 'modulo': 8, | |
| 'n_iters': 15, | |
| 'lr': 0.002, | |
| 'min_side': 512, | |
| 'max_scales': 3, | |
| 'px_budget': 1800000 | |
| } | |
| class Inpainter(): | |
| def __init__(self, config): | |
| self.model = None | |
| self.config = config | |
| self.device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| self.scale_factor = config['scale_factor'] | |
| self.pad_out_to_modulo = config['pad_out_to_modulo'] | |
| self.predict_config = config['predict'] | |
| self.predict_config['model_path'] = 'big-lama' | |
| self.predict_config['model_checkpoint'] = 'best.ckpt' | |
| self.refiner_config = refiner_config | |
| def load_model_from_checkpoint(self, model_path, checkpoint): | |
| train_config_path = os.path.join(model_path, 'config.yaml') | |
| with open(train_config_path, 'r') as f: | |
| train_config = OmegaConf.create(yaml.safe_load(f)) | |
| train_config.training_model.predict_only = True | |
| train_config.visualizer.kind = 'noop' | |
| checkpoint_path = os.path.join(model_path, | |
| 'models', | |
| checkpoint) | |
| self.model = load_checkpoint(train_config, checkpoint_path, strict=False, map_location='cpu') | |
| def load_batch_data(self, img_, mask_): | |
| """Loads the image and mask from the given filenames. | |
| """ | |
| image = prepare_image(img_, mode='RGB') | |
| mask = prepare_image(mask_, mode='L') | |
| result = dict(image=image, mask=mask[None, ...]) | |
| if self.scale_factor is not None: | |
| result['image'] = scale_image(result['image'], self.scale_factor) | |
| result['mask'] = scale_image(result['mask'], self.scale_factor, interpolation=cv2.INTER_NEAREST) | |
| if self.pad_out_to_modulo is not None and self.pad_out_to_modulo > 1: | |
| result['unpad_to_size'] = result['image'].shape[1:] | |
| result['image'] = pad_img_to_modulo(result['image'], self.pad_out_to_modulo) | |
| result['mask'] = pad_img_to_modulo(result['mask'], self.pad_out_to_modulo) | |
| return result | |
| def inpaint_img(self, original_img, mask_img, refine=False) -> Image: | |
| """ Inpaints the image region defined by the given mask. | |
| White pixels are to be masked and black pixels kept. | |
| args: | |
| refine: if True, uses the refinement model to enhance the inpainting result, at the cost of speed. | |
| returns: the inpainted image | |
| """ | |
| # in case we are given filenames instead of images | |
| if isinstance(original_img, str): | |
| original_img = load_image(original_img, mode='RGB') | |
| mask_img = load_image(mask_img, mode='L') | |
| self.model.eval() | |
| if not refine: | |
| self.model.to(self.device) | |
| # load the image and mask | |
| batch = default_collate([self.load_batch_data(original_img, mask_img)]) | |
| if refine: | |
| assert 'unpad_to_size' in batch, "Unpadded size is required for the refinement" | |
| # image unpadding is taken care of in the refiner, so that output image | |
| # is same size as the input image | |
| cur_res = refine_predict(batch, self.model, **self.refiner_config) | |
| cur_res = cur_res[0].permute(1,2,0).detach().cpu().numpy() | |
| else: | |
| with torch.no_grad(): | |
| batch = move_to_device(batch, self.device) | |
| batch['mask'] = (batch['mask'] > 0) * 1 | |
| batch = self.model(batch) | |
| cur_res = batch[self.predict_config['out_key']][0].permute(1, 2, 0).detach().cpu().numpy() | |
| unpad_to_size = batch.get('unpad_to_size', None) | |
| if unpad_to_size is not None: | |
| orig_height, orig_width = unpad_to_size | |
| cur_res = cur_res[:orig_height, :orig_width] | |
| cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8') | |
| rslt_image = Image.fromarray(cur_res, 'RGB') | |
| #cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR) | |
| return rslt_image | |