apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): handle blending numpy arrays

This commit is contained in:
Sean Sube 2023-11-25 21:19:11 -06:00
parent 1c3b2f8dfc
commit 83ebd77c3d
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 16 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
api/onnx_web/chain/tile.py
View File

@ -2,7 +2,7 @@ import itertools
from enum import Enum from enum import Enum
from logging import getLogger from logging import getLogger
from math import ceil from math import ceil
from typing import Any, Callable, List, Optional, Protocol, Tuple from typing import Any, Callable, List, Optional, Protocol, Tuple, Union
import numpy as np import numpy as np
from PIL import Image from PIL import Image
@ -138,6 +138,20 @@ def make_tile_mask(
return mask return mask
def get_channels(image: Union[np.ndarray, Image.Image]) -> int:
if isinstance(image, np.ndarray):
return image.shape[-1]
if image.mode == "RGBA":
return 4
elif image.mode == "RGB":
return 3
elif image.mode == "L":
return 1
raise ValueError("unknown image format")
def blend_tiles( def blend_tiles(
tiles: List[Tuple[int, int, Image.Image]], tiles: List[Tuple[int, int, Image.Image]],
scale: int, scale: int,
@ -151,7 +165,7 @@ def blend_tiles(
"adjusting tile size from %s to %s based on %s overlap", tile, adj_tile, overlap "adjusting tile size from %s to %s based on %s overlap", tile, adj_tile, overlap
) )
channels = max([4 if tile_image.mode == "RGBA" else 3 for _left, _top, tile_image in tiles]) channels = max([get_channels(tile_image) for _left, _top, tile_image in tiles])
scaled_size = (height * scale, width * scale, channels) scaled_size = (height * scale, width * scale, channels)
count = np.zeros(scaled_size) count = np.zeros(scaled_size)