Spaces:
Build error
Build error
File size: 2,552 Bytes
fa8453f |
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 |
import os
import cv2
import default_paths as dp
from upscaler.GPEN import GPEN
from upscaler.GFPGAN import GFPGAN
from upscaler.codeformer import CodeFormer
from upscaler.restoreformer import RestoreFormer
def gfpgan_runner(img, model):
img = model.enhance(img)
return img
def codeformer_runner(img, model):
img = model.enhance(img, w=0.9)
return img
def gpen_runner(img, model):
img = model.enhance(img)
return img
def restoreformer_runner(img, model):
img = model.enhance(img)
return img
supported_upscalers = {
"CodeFormer": (dp.CODEFORMER_PATH, codeformer_runner),
"GFPGANv1.4": (dp.GFPGAN_V14_PATH, gfpgan_runner),
"GFPGANv1.3": (dp.GFPGAN_V13_PATH, gfpgan_runner),
"GFPGANv1.2": (dp.GFPGAN_V12_PATH, gfpgan_runner),
"GPEN-BFR-512": (dp.GPEN_BFR_512_PATH, gpen_runner),
"GPEN-BFR-256": (dp.GPEN_BFR_256_PATH, gpen_runner),
"RestoreFormer": (dp.RESTOREFORMER_PATH, gpen_runner),
}
cv2_upscalers = ["LANCZOS4", "CUBIC", "NEAREST"]
def get_available_upscalers_names():
available = []
for name, data in supported_upscalers.items():
if os.path.exists(data[0]):
available.append(name)
return available
def load_face_upscaler(name='GFPGAN', provider=["CPUExecutionProvider"], session_options=None):
assert name in get_available_upscalers_names() + cv2_upscalers, f"Face upscaler {name} unavailable."
if name in supported_upscalers.keys():
model_path, model_runner = supported_upscalers.get(name)
if name == 'CodeFormer':
model = CodeFormer(model_path=model_path, provider=provider, session_options=session_options)
elif name.startswith('GFPGAN'):
model = GFPGAN(model_path=model_path, provider=provider, session_options=session_options)
elif name.startswith('GPEN'):
model = GPEN(model_path=model_path, provider=provider, session_options=session_options)
elif name == "RestoreFormer":
model = RestoreFormer(model_path=model_path, provider=provider, session_options=session_options)
elif name == 'LANCZOS4':
model = None
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_LANCZOS4)
elif name == 'CUBIC':
model = None
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_CUBIC)
elif name == 'NEAREST':
model = None
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_NEAREST)
else:
model = None
return (model, model_runner)
|