erase-lama / verify_against_official.py

big-lama for Core ML, verified against the official weights

ef3cdb3 verified 23 days ago

4.5 kB

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}')

	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}')