2023-01-27 23:08:36 +00:00
|
|
|
from PIL import Image
|
|
|
|
from os import path
|
2023-01-28 14:19:40 +00:00
|
|
|
from typing import Any, List, Optional, Protocol, Tuple
|
2023-01-27 23:08:36 +00:00
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
from ..params import (
|
2023-01-28 05:28:14 +00:00
|
|
|
ImageParams,
|
2023-01-28 04:48:06 +00:00
|
|
|
StageParams,
|
|
|
|
)
|
|
|
|
from ..utils import (
|
2023-01-27 23:08:36 +00:00
|
|
|
ServerContext,
|
|
|
|
)
|
2023-01-28 14:19:40 +00:00
|
|
|
from .utils import (
|
|
|
|
process_tiles,
|
|
|
|
)
|
2023-01-27 23:08:36 +00:00
|
|
|
|
|
|
|
|
|
|
|
class StageCallback(Protocol):
|
|
|
|
def __call__(
|
|
|
|
self,
|
|
|
|
ctx: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
params: ImageParams,
|
|
|
|
source: Image.Image,
|
|
|
|
**kwargs: Any
|
|
|
|
) -> Image.Image:
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
PipelineStage = Tuple[StageCallback, StageParams, Optional[dict]]
|
2023-01-27 23:08:36 +00:00
|
|
|
|
|
|
|
|
|
|
|
class ChainPipeline:
|
|
|
|
'''
|
|
|
|
Run many stages in series, passing the image results from each to the next, and processing
|
|
|
|
tiles as needed.
|
|
|
|
'''
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
self,
|
2023-01-27 23:38:21 +00:00
|
|
|
stages: List[PipelineStage] = [],
|
2023-01-27 23:08:36 +00:00
|
|
|
):
|
|
|
|
'''
|
|
|
|
Create a new pipeline that will run the given stages.
|
|
|
|
'''
|
|
|
|
self.stages = stages
|
|
|
|
|
|
|
|
def append(self, stage: PipelineStage):
|
|
|
|
'''
|
|
|
|
Append an additional stage to this pipeline.
|
|
|
|
'''
|
|
|
|
self.stages.append(stage)
|
|
|
|
|
|
|
|
def __call__(self, ctx: ServerContext, params: ImageParams, source: Image.Image) -> Image.Image:
|
|
|
|
'''
|
|
|
|
TODO: handle List[Image] outputs
|
|
|
|
'''
|
|
|
|
print('running pipeline on source image with dimensions %sx%s' %
|
|
|
|
source.size)
|
|
|
|
image = source
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
for stage_pipe, stage_params, stage_kwargs in self.stages:
|
2023-01-28 05:28:14 +00:00
|
|
|
name = stage_params.name or stage_pipe.__name__
|
2023-01-28 04:48:06 +00:00
|
|
|
kwargs = stage_kwargs or {}
|
2023-01-27 23:38:21 +00:00
|
|
|
print('running pipeline stage %s on result image with dimensions %sx%s' %
|
|
|
|
(name, image.width, image.height))
|
2023-01-28 04:48:06 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
if image.width > stage_params.tile_size or image.height > stage_params.tile_size:
|
2023-01-27 23:38:21 +00:00
|
|
|
print('source image larger than tile size of %s, tiling stage' % (
|
|
|
|
stage_params.tile_size))
|
2023-01-27 23:08:36 +00:00
|
|
|
|
|
|
|
def stage_tile(tile: Image.Image) -> Image.Image:
|
2023-01-28 04:48:06 +00:00
|
|
|
tile = stage_pipe(ctx, stage_params, params, tile,
|
|
|
|
**kwargs)
|
2023-01-27 23:08:36 +00:00
|
|
|
tile.save(path.join(ctx.output_path, 'last-tile.png'))
|
|
|
|
return tile
|
|
|
|
|
|
|
|
image = process_tiles(
|
|
|
|
image, stage_params.tile_size, stage_params.outscale, [stage_tile])
|
|
|
|
else:
|
2023-01-27 23:38:21 +00:00
|
|
|
print('source image within tile size, running stage')
|
2023-01-28 04:48:06 +00:00
|
|
|
image = stage_pipe(ctx, stage_params, params, image,
|
|
|
|
**kwargs)
|
2023-01-27 23:08:36 +00:00
|
|
|
|
2023-01-27 23:38:21 +00:00
|
|
|
print('finished running pipeline stage %s, result size: %sx%s' %
|
|
|
|
(name, image.width, image.height))
|
2023-01-27 23:08:36 +00:00
|
|
|
|
|
|
|
print('finished running pipeline, result size: %sx%s' % image.size)
|
|
|
|
return image
|