Spaces:
Running
Running
File size: 2,644 Bytes
c64fb9f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
## Imports
import os, time, argparse, base64, requests, os, json, sys, datetime
from itertools import product
import warnings
warnings.filterwarnings("ignore")
# import cv2
from PIL import Image
import numpy as np
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import ImageNet
import torchvision.transforms as T
def readImg(p):
return Image.open(p)
def toImg(t):
return T.ToPILImage()(t)
def invtrans(mask, image, method = Image.BICUBIC):
return mask.resize(image.size, method)
def merge(mask, image, grap_scale = 200):
gray = np.ones((image.size[1], image.size[0], 3))*grap_scale
image_np = np.array(image).astype(np.float32)[..., :3]
mask_np = np.array(mask).astype(np.float32)
mask_np = mask_np / 255.0
blended_np = image_np * mask_np[:, :, None] + (1 - mask_np[:, :, None]) * gray
blended_image = Image.fromarray((blended_np).astype(np.uint8))
return blended_image
def normalize(mat, method = "max"):
if method == "max":
return (mat.max() - mat) / (mat.max() - mat.min())
elif method == "min":
return (mat - mat.min()) / (mat.max() - mat.min())
else:
raise NotImplementedError
def enhance(mat, coe=10):
mat = mat - mat.mean()
mat = mat / mat.std()
mat = mat * coe
mat = torch.sigmoid(mat)
mat = mat.clamp(0,1)
return mat
def revise_mask(patch_mask, kernel_size = 3, enhance_coe = 10):
patch_mask = normalize(patch_mask, "min")
patch_mask = enhance(patch_mask, coe = enhance_coe)
assert kernel_size % 2 == 1
padding_size = int((kernel_size - 1) / 2)
conv = torch.nn.Conv2d(1,1,kernel_size = kernel_size, padding = padding_size, padding_mode = "replicate", stride = 1, bias = False)
conv.weight.data = torch.ones_like(conv.weight.data) / kernel_size**2
conv.to(patch_mask.device)
patch_mask = conv(patch_mask.unsqueeze(0))[0]
mask = patch_mask
return mask
def blend_mask(image_path_or_pil_image, mask, enhance_coe, kernel_size, interpolate_method_name, grayscale):
mask = revise_mask(mask.float(), kernel_size = kernel_size, enhance_coe = enhance_coe)
mask = mask.detach().cpu()
mask = toImg(mask.reshape(1,24,24))
if isinstance(image_path_or_pil_image, Image.Image):
image = image_path_or_pil_image
else:
raise NotImplementedError
interpolate_method = getattr(Image, interpolate_method_name)
mask = invtrans(mask, image, method = interpolate_method)
merged_image = merge(mask.convert("L"), image.convert("RGB"), grayscale).convert("RGB")
return merged_image |