2023-04-10 22:49:56 +00:00
|
|
|
from logging import getLogger
|
|
|
|
from os import path
|
2023-11-19 00:13:13 +00:00
|
|
|
from typing import Optional
|
2023-04-10 22:49:56 +00:00
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
from ..models.onnx import OnnxModel
|
2023-12-03 18:13:45 +00:00
|
|
|
from ..params import DeviceParams, ImageParams, SizeChart, StageParams, UpscaleParams
|
2023-07-03 16:33:56 +00:00
|
|
|
from ..server import ModelTypes, ServerContext
|
2023-04-10 22:49:56 +00:00
|
|
|
from ..utils import run_gc
|
|
|
|
from ..worker import WorkerContext
|
2023-11-18 23:18:23 +00:00
|
|
|
from .base import BaseStage
|
2023-11-19 00:08:38 +00:00
|
|
|
from .result import StageResult
|
2023-04-10 22:49:56 +00:00
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
|
|
|
2023-07-02 23:21:21 +00:00
|
|
|
class UpscaleSwinIRStage(BaseStage):
|
2023-12-03 18:13:45 +00:00
|
|
|
max_tile = SizeChart.micro
|
2023-07-01 12:10:53 +00:00
|
|
|
|
|
|
|
def load(
|
|
|
|
self,
|
|
|
|
server: ServerContext,
|
|
|
|
_stage: StageParams,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
device: DeviceParams,
|
|
|
|
):
|
|
|
|
# must be within the load function for patch to take effect
|
|
|
|
model_path = path.join(server.model_path, "%s.onnx" % (upscale.upscale_model))
|
|
|
|
cache_key = (model_path,)
|
2023-07-03 16:33:56 +00:00
|
|
|
cache_pipe = server.cache.get(ModelTypes.upscaling, cache_key)
|
2023-07-01 12:10:53 +00:00
|
|
|
|
|
|
|
if cache_pipe is not None:
|
|
|
|
logger.info("reusing existing SwinIR pipeline")
|
|
|
|
return cache_pipe
|
|
|
|
|
|
|
|
logger.debug("loading SwinIR model from %s", model_path)
|
|
|
|
|
|
|
|
pipe = OnnxModel(
|
|
|
|
server,
|
|
|
|
model_path,
|
2024-01-13 00:58:26 +00:00
|
|
|
provider=device.ort_provider("swinir"),
|
2023-07-01 12:10:53 +00:00
|
|
|
sess_options=device.sess_options(),
|
|
|
|
)
|
|
|
|
|
2023-07-03 16:33:56 +00:00
|
|
|
server.cache.set(ModelTypes.upscaling, cache_key, pipe)
|
2023-07-01 12:10:53 +00:00
|
|
|
run_gc([device])
|
|
|
|
|
|
|
|
return pipe
|
|
|
|
|
|
|
|
def run(
|
|
|
|
self,
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-07-01 12:10:53 +00:00
|
|
|
server: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
_params: ImageParams,
|
2023-11-19 00:08:38 +00:00
|
|
|
sources: StageResult,
|
2023-07-01 12:10:53 +00:00
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
stage_source: Optional[Image.Image] = None,
|
|
|
|
**kwargs,
|
2023-11-19 00:08:38 +00:00
|
|
|
) -> StageResult:
|
2023-07-01 12:10:53 +00:00
|
|
|
upscale = upscale.with_args(**kwargs)
|
|
|
|
|
|
|
|
if upscale.upscale_model is None:
|
2023-07-16 00:00:52 +00:00
|
|
|
logger.warning("no correction model given, skipping")
|
2023-07-04 18:29:58 +00:00
|
|
|
return sources
|
2023-07-01 12:10:53 +00:00
|
|
|
|
|
|
|
logger.info("correcting faces with SwinIR model: %s", upscale.upscale_model)
|
2023-07-15 23:54:54 +00:00
|
|
|
device = worker.get_device()
|
2023-07-01 12:10:53 +00:00
|
|
|
swinir = self.load(server, stage, upscale, device)
|
|
|
|
|
2023-07-04 18:29:58 +00:00
|
|
|
outputs = []
|
2024-01-06 02:11:58 +00:00
|
|
|
for source in sources.as_arrays():
|
2023-07-04 18:29:58 +00:00
|
|
|
# TODO: add support for grayscale (1-channel) images
|
2023-11-19 00:08:38 +00:00
|
|
|
image = source / 255.0
|
2023-07-04 18:29:58 +00:00
|
|
|
image = image[:, :, [2, 1, 0]].astype(np.float32).transpose((2, 0, 1))
|
|
|
|
image = np.expand_dims(image, axis=0)
|
|
|
|
logger.trace("SwinIR input shape: %s", image.shape)
|
|
|
|
|
|
|
|
scale = upscale.outscale
|
2023-11-19 00:13:13 +00:00
|
|
|
logger.trace(
|
|
|
|
"SwinIR output shape: %s",
|
|
|
|
(
|
2023-07-04 18:29:58 +00:00
|
|
|
image.shape[0],
|
|
|
|
image.shape[1],
|
|
|
|
image.shape[2] * scale,
|
|
|
|
image.shape[3] * scale,
|
2023-11-19 00:13:13 +00:00
|
|
|
),
|
|
|
|
)
|
2023-07-04 18:29:58 +00:00
|
|
|
|
2023-11-19 00:08:38 +00:00
|
|
|
output = swinir(image)
|
|
|
|
output = np.clip(np.squeeze(output, axis=0), 0, 1)
|
|
|
|
output = output[[2, 1, 0], :, :].transpose((1, 2, 0))
|
|
|
|
output = (output * 255.0).round().astype(np.uint8)
|
2023-07-01 12:10:53 +00:00
|
|
|
|
2023-11-19 00:08:38 +00:00
|
|
|
logger.info("output image size: %s", output.shape)
|
2023-07-04 18:29:58 +00:00
|
|
|
outputs.append(output)
|
2023-07-01 12:10:53 +00:00
|
|
|
|
2024-01-06 02:11:58 +00:00
|
|
|
return StageResult(images=outputs, metadata=sources.metadata)
|