Spaces:
Running
Running
| import cv2 | |
| import torch | |
| import torchvision | |
| import numpy as np | |
| import torch.nn as nn | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import torch.nn.functional as F | |
| import torchvision.transforms as transforms | |
| from . model import BiSeNet | |
| class SoftErosion(nn.Module): | |
| def __init__(self, kernel_size=15, threshold=0.6, iterations=1): | |
| super(SoftErosion, self).__init__() | |
| r = kernel_size // 2 | |
| self.padding = r | |
| self.iterations = iterations | |
| self.threshold = threshold | |
| # Create kernel | |
| y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size)) | |
| dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2) | |
| kernel = dist.max() - dist | |
| kernel /= kernel.sum() | |
| kernel = kernel.view(1, 1, *kernel.shape) | |
| self.register_buffer('weight', kernel) | |
| def forward(self, x): | |
| batch_size = x.size(0) # Get the batch size | |
| output = [] | |
| for i in tqdm(range(batch_size), desc="Soft-Erosion", leave=False): | |
| input_tensor = x[i:i+1] # Take one input tensor from the batch | |
| input_tensor = input_tensor.float() # Convert input to float tensor | |
| input_tensor = input_tensor.unsqueeze(1) # Add a channel dimension | |
| for _ in range(self.iterations - 1): | |
| input_tensor = torch.min(input_tensor, F.conv2d(input_tensor, weight=self.weight, | |
| groups=input_tensor.shape[1], | |
| padding=self.padding)) | |
| input_tensor = F.conv2d(input_tensor, weight=self.weight, groups=input_tensor.shape[1], | |
| padding=self.padding) | |
| mask = input_tensor >= self.threshold | |
| input_tensor[mask] = 1.0 | |
| input_tensor[~mask] /= input_tensor[~mask].max() | |
| input_tensor = input_tensor.squeeze(1) # Remove the extra channel dimension | |
| output.append(input_tensor.detach().cpu().numpy()) | |
| return np.array(output) | |
| transform = transforms.Compose([ | |
| transforms.Resize((512, 512)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) | |
| ]) | |
| def init_parsing_model(model_path, device="cpu"): | |
| net = BiSeNet(19) | |
| net.to(device) | |
| net.load_state_dict(torch.load(model_path)) | |
| net.eval() | |
| return net | |
| def transform_images(imgs): | |
| tensor_images = torch.stack([transform(Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))) for img in imgs], dim=0) | |
| return tensor_images | |
| def get_parsed_mask(net, imgs, classes=[1, 2, 3, 4, 5, 10, 11, 12, 13], device="cpu", batch_size=8, softness=20): | |
| if softness > 0: | |
| smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=softness).to(device) | |
| masks = [] | |
| for i in tqdm(range(0, len(imgs), batch_size), total=len(imgs) // batch_size, desc="Face-parsing"): | |
| batch_imgs = imgs[i:i + batch_size] | |
| tensor_images = transform_images(batch_imgs).to(device) | |
| with torch.no_grad(): | |
| out = net(tensor_images)[0] | |
| # parsing = out.argmax(dim=1) | |
| # arget_classes = torch.tensor(classes).to(device) | |
| # batch_masks = torch.isin(parsing, target_classes).to(device) | |
| ## torch.isin was slightly slower in my test, so using np.isin | |
| parsing = out.argmax(dim=1).detach().cpu().numpy() | |
| batch_masks = np.isin(parsing, classes).astype('float32') | |
| if softness > 0: | |
| # batch_masks = smooth_mask(batch_masks).transpose(1,0,2,3)[0] | |
| mask_tensor = torch.from_numpy(batch_masks.copy()).float().to(device) | |
| batch_masks = smooth_mask(mask_tensor).transpose(1,0,2,3)[0] | |
| yield batch_masks | |
| #masks.append(batch_masks) | |
| #if len(masks) >= 1: | |
| # masks = np.concatenate(masks, axis=0) | |
| # masks = np.repeat(np.expand_dims(masks, axis=1), 3, axis=1) | |
| # for i, mask in enumerate(masks): | |
| # cv2.imwrite(f"mask/{i}.jpg", (mask * 255).astype("uint8")) | |
| #return masks | |