File size: 4,495 Bytes
ef3cdb3 | 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 83 | import sys, types, os, math
import numpy as np, torch
from PIL import Image, ImageDraw
L = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, f'{L}/lama')
# Stub the training-only imports so the generator module loads without kornia/lightning.
st = types.ModuleType('saicinpainting.training.modules.spatial_transform')
class LearnableSpatialTransformWrapper(torch.nn.Module): pass
st.LearnableSpatialTransformWrapper = LearnableSpatialTransformWrapper
sys.modules['saicinpainting.training.modules.spatial_transform'] = st
ut = types.ModuleType('saicinpainting.utils')
ut.get_shape = lambda t: tuple(t.shape)
sys.modules['saicinpainting.utils'] = ut
from saicinpainting.training.modules.ffc import FFCResNetGenerator
gen = FFCResNetGenerator(4, 3, ngf=64, n_downsampling=3, n_blocks=18, add_out_act='sigmoid',
init_conv_kwargs=dict(ratio_gin=0, ratio_gout=0, enable_lfu=False),
downsample_conv_kwargs=dict(ratio_gin=0, ratio_gout=0, enable_lfu=False),
resnet_conv_kwargs=dict(ratio_gin=0.75, ratio_gout=0.75, enable_lfu=False))
import pickle
class _Any(dict):
def __init__(self, *a, **k): pass
def __setstate__(self, state): pass
class _Unpickler(pickle.Unpickler):
def find_class(self, module, name):
if module.startswith('pytorch_lightning') or module.startswith('lightning'):
return _Any
return super().find_class(module, name)
_pm = types.ModuleType('lightning_free_pickle'); _pm.Unpickler = _Unpickler; _pm.load = pickle.load
ck = torch.load(f'{L}/big-lama/models/best.ckpt', map_location='cpu', weights_only=False, pickle_module=_pm)
sd = {k[len('generator.'):]: v for k, v in ck['state_dict'].items() if k.startswith('generator.')}
missing, unexpected = gen.load_state_dict(sd, strict=True), None
gen.eval()
import coremltools as ct
cm = ct.models.MLModel(f'{L}/tp/LaMa.mlpackage', compute_units=ct.ComputeUnit.CPU_ONLY)
N = 800
def scenes():
rng = np.random.default_rng(7)
# 1: smooth gradient with shapes
a = np.zeros((N, N, 3), np.float32)
yy, xx = np.mgrid[0:N, 0:N] / N
a[..., 0] = xx; a[..., 1] = yy; a[..., 2] = 0.5 + 0.5 * np.sin(6 * xx)
im = Image.fromarray((a * 255).astype(np.uint8)); d = ImageDraw.Draw(im)
d.ellipse([200, 200, 420, 420], fill=(240, 40, 40)); d.rectangle([500, 450, 700, 700], fill=(30, 30, 200))
yield 'gradient', im
# 2: noisy texture
yield 'texture', Image.fromarray((rng.random((N, N, 3)) * 80 + np.array([90, 140, 60])).astype(np.uint8))
# 3: a stock photo (Apple's User Pictures)
src = '/System/Library/Templates/Data/Library/User Pictures/Animals/Eagle.heic'
os.system(f"sips -s format png -z {N} {N} '{src}' --out {L}/eagle.png >/dev/null 2>&1")
yield 'eagle', Image.open(f'{L}/eagle.png').convert('RGB').resize((N, N))
def masks():
m1 = Image.new('L', (N, N), 0); ImageDraw.Draw(m1).rectangle([250, 250, 480, 520], fill=255)
m2 = Image.new('L', (N, N), 0); ImageDraw.Draw(m2).ellipse([380, 120, 640, 340], fill=255)
return [('rect', m1), ('ellipse', m2)]
def psnr(a, b, sel=None):
a = a.astype(np.float64); b = b.astype(np.float64)
if sel is not None: a, b = a[sel], b[sel]
mse = ((a - b) ** 2).mean()
return 99.0 if mse == 0 else 10 * math.log10(255 ** 2 / mse)
for sname, img in scenes():
for mname, m in masks():
x = torch.from_numpy(np.asarray(img, np.float32) / 255).permute(2, 0, 1)[None]
mk = (torch.from_numpy(np.asarray(m, np.float32) / 255) > 0.5).float()[None, None]
with torch.no_grad():
out = gen(torch.cat([x * (1 - mk), mk], 1))
pt = (out[0].permute(1, 2, 0).numpy() * 255).clip(0, 255)
cmo = np.asarray(cm.predict({'image': img, 'mask': m})['output'].convert('RGB'), np.float32)
sel = np.asarray(m) > 127
# Inside the mask is what the model invents; outside, either may pass the input through.
print(f'{sname:9} {mname:8} PSNR inside mask {psnr(pt, cmo, sel):6.2f} dB | whole frame {psnr(pt, cmo):6.2f} dB')
# Outside the mask, the package should pass the input through unchanged.
for sname, img in scenes():
m = masks()[0][1]
cmo = np.asarray(cm.predict({'image': img, 'mask': m})['output'].convert('RGB'), np.float32)
sel = np.asarray(m) <= 127
print(f'{sname:9} outside mask vs input: PSNR {psnr(np.asarray(img, np.float32), cmo, sel):6.2f} dB, max diff {np.abs(np.asarray(img, np.float32)[sel] - cmo[sel]).max():.0f}')
|