Spaces:
Running
Running
File size: 4,890 Bytes
2514fb4 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import datetime
import hashlib
import numpy as np
import os
import subprocess
from pathlib import Path
from typing import Any, Dict
import cv2
import gradio as gr
from joblib import Parallel, delayed
from numpy.typing import NDArray
from PIL import Image
def _run_in_subprocess(command: str, wd: str) -> Any:
p = subprocess.Popen(command, shell=True, cwd=wd)
(output, err) = p.communicate()
p_status = p.wait()
print("Status of subprocess: ", p_status)
return p_status
SWIN_IR_WD = "KAIR"
SWINIR_CKPT_DIR: str = Path("KAIR/model_zoo/")
MODEL_NAME_TO_PATH: Dict[str, Path] = {
"LambdaSwinIR_v0.1": Path(str(SWINIR_CKPT_DIR) + "/805000_G.pth"),
}
SWINIR_NAME_TO_PATCH_SIZE: Dict[str, int] = {
"LambdaSwinIR_v0.1": 96,
}
SWINIR_NAME_TO_SCALE: Dict[str, int] = {
"LambdaSwinIR_v0.1": 2,
}
SWINIR_NAME_TO_LARGE_MODEL: Dict[str, bool] = {
"LambdaSwinIR_v0.1": False,
}
def _run_swin_ir(
image: NDArray,
model_path: Path,
patch_size: int,
scale: int,
is_large_model: bool,
):
print("model_path: ", str(model_path))
m = hashlib.sha256()
now_time = datetime.datetime.utcnow()
m.update(bytes(str(model_path), encoding='utf-8') +
bytes(now_time.strftime("%Y-%m-%d %H:%M:%S.%f"), encoding='utf-8'))
random_id = m.hexdigest()[0:20]
cwd = os.getcwd()
input_root = Path(cwd + "/sr_interactive_tmp")
input_root.mkdir(parents=True, exist_ok=True)
Image.fromarray(image).save(str(input_root) + "/gradio_img.png")
command = f"python main_test_swinir.py --scale {scale} " + \
f"--folder_lq {input_root} --task real_sr " + \
f"--model_path {cwd}/{model_path} --training_patch_size {patch_size}"
if is_large_model:
command += " --large_model"
print("COMMAND: ", command)
status = _run_in_subprocess(command, wd=cwd + "/" + SWIN_IR_WD)
print("STATUS: ", status)
if scale == 2:
str_scale = "2"
if scale == 4:
str_scale = "4_large"
output_img = Image.open(f"{cwd}/KAIR/results/swinir_real_sr_x{str_scale}/gradio_img_SwinIR.png")
output_root = Path("./sr_interactive_tmp_output")
output_root.mkdir(parents=True, exist_ok=True)
output_img.save(str(output_root) + "/SwinIR_" + random_id + ".png")
print("SAVING: SwinIR_" + random_id + ".png")
result = np.array(output_img)
return result
def _bilinear_upsample(image: NDArray):
result = cv2.resize(
image,
dsize=(image.shape[1] * 2, image.shape[0] * 2),
interpolation=cv2.INTER_LANCZOS4
)
return result
def _decide_sr_algo(model_name: str, image: NDArray):
# if "SwinIR" in model_name:
# result = _run_swin_ir(image,
# model_path=MODEL_NAME_TO_PATH[model_name],
# patch_size=SWINIR_NAME_TO_PATCH_SIZE[model_name],
# scale=SWINIR_NAME_TO_SCALE[model_name],
# is_large_model=("SwinIR-L" in model_name))
# else:
# result = _bilinear_upsample(image)
# elif algo == SR_OPTIONS[1]:
# result = _run_maxine(image, mode="SR")
# elif algo == SR_OPTIONS[2]:
# result = _run_maxine(image, mode="UPSCALE")
# return result
result = _run_swin_ir(image,
model_path=MODEL_NAME_TO_PATH[model_name],
patch_size=SWINIR_NAME_TO_PATCH_SIZE[model_name],
scale=SWINIR_NAME_TO_SCALE[model_name],
is_large_model=SWINIR_NAME_TO_LARGE_MODEL[model_name])
return result
def _super_resolve(model_name: str, input_img):
# futures = []
# with ThreadPoolExecutor(max_workers=4) as executor:
# for model_name in model_names:
# futures.append(executor.submit(_decide_sr_algo, model_name, input_img))
# return [f.result() for f in futures]
# return Parallel(n_jobs=2, prefer="threads")(
# delayed(_decide_sr_algo)(model_name, input_img)
# for model_name in model_names
# )
return _decide_sr_algo(model_name, input_img)
def _gradio_handler(sr_option: str, input_img: NDArray):
return _super_resolve(sr_option, input_img)
gr.close_all()
SR_OPTIONS = ["LambdaSwinIR_v0.1"]
examples = [
["LambdaSwinIR_v0.1", "examples/oldphoto6.png"],
["LambdaSwinIR_v0.1", "examples/Lincoln.png"],
["LambdaSwinIR_v0.1", "examples/OST_009.png"],
["LambdaSwinIR_v0.1", "examples/00003.png"],
["LambdaSwinIR_v0.1", "examples/00000067_cropped.png"],
]
ui = gr.Interface(fn=_gradio_handler,
inputs=[
gr.Radio(SR_OPTIONS),
gr.Image(image_mode="RGB")
],
outputs=["image"],
live=False,
examples=examples,
cache_examples=True)
ui.launch(enable_queue=True)
|