Spaces:
Running
Running
| import torch | |
| from .model import PHNet | |
| import torchvision.transforms.functional as tf | |
| from .util import inference_img, log | |
| from .stylematte import StyleMatte | |
| import numpy as np | |
| import onnx | |
| from .engine import execute_onnx_model | |
| import cv2 | |
| from torchvision import transforms | |
| import time | |
| class Inference: | |
| def __init__(self, **kwargs): | |
| self.rank = 0 | |
| self.__dict__.update(kwargs) | |
| self.model = PHNet(enc_sizes=self.enc_sizes, | |
| skips=self.skips, | |
| grid_count=self.grid_counts, | |
| init_weights=self.init_weights, | |
| init_value=self.init_value) | |
| state = torch.load(self.checkpoint.harmonizer, | |
| map_location=self.device) | |
| self.model.load_state_dict(state, strict=True) | |
| self.model.eval() | |
| def harmonize(self, composite, mask): | |
| if len(composite.shape) < 4: | |
| composite = composite.unsqueeze(0) | |
| while len(mask.shape) < 4: | |
| mask = mask.unsqueeze(0) | |
| composite = tf.resize(composite, [self.image_size, self.image_size]) | |
| mask = tf.resize(mask, [self.image_size, self.image_size]) | |
| log(composite.shape, mask.shape) | |
| with torch.no_grad(): | |
| harmonized = self.model(composite, mask) # ['harmonized'] | |
| result = harmonized * mask + composite * (1-mask) | |
| return result | |
| class Matting: | |
| def __init__(self, **kwargs): | |
| self.rank = 0 | |
| self.__dict__.update(kwargs) | |
| if self.onnx: | |
| self.model = onnx.load(self.checkpoint.matting_onnx) | |
| else: | |
| self.model = StyleMatte().to(self.device) | |
| state = torch.load(self.checkpoint.matting, | |
| map_location=self.device) | |
| self.model.load_state_dict(state, strict=True) | |
| self.model.eval() | |
| def extract(self, inp): | |
| mask = inference_img(self.model, inp, self.device, self.onnx) | |
| inp_np = np.array(inp) | |
| fg = mask[:, :, None]*inp_np | |
| return [mask, fg] | |
| def inference_img(model, img, device='cpu', onnx=True): | |
| beg = time.time() | |
| h, w, _ = img.shape | |
| # print(img.shape) | |
| if h % 8 != 0 or w % 8 != 0: | |
| img = cv2.copyMakeBorder(img, 8-h % 8, 0, 8-w % | |
| 8, 0, cv2.BORDER_REFLECT) | |
| # print(img.shape) | |
| tensor_img = torch.from_numpy(img).permute(2, 0, 1).to(device) | |
| input_t = tensor_img/255.0 | |
| normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| input_t = normalize(input_t) | |
| input_t = input_t.unsqueeze(0).float() | |
| end_p = time.time() | |
| if onnx: | |
| out = execute_onnx_model(input_t, model) | |
| else: | |
| with torch.no_grad(): | |
| out = model(input_t).cpu().numpy() | |
| end = time.time() | |
| log(f"Inference time: {end-beg}, processing time: {end_p-beg}") | |
| # print("out",out.shape) | |
| result = out[0][:, -h:, -w:] | |
| # print(result.shape) | |
| return result[0] | |