103 lines
3.5 KiB
Python
103 lines
3.5 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
class CalculateTilingSpec_Tiled:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": 512, "min": 1, "max": 16384}),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": 16384}),
|
|
"rows": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"columns": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"overlap": ("INT", {"default": 64, "min": 0, "max": 1024}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*", "INT", "INT")
|
|
RETURN_NAMES = ("tiling_spec", "tile_width", "tile_height")
|
|
FUNCTION = "calculate"
|
|
CATEGORY = "SplitImage"
|
|
|
|
def calculate(self, width, height, rows, columns, overlap):
|
|
tile_w = math.ceil((width + (columns - 1) * overlap) / columns)
|
|
tile_h = math.ceil((height + (rows - 1) * overlap) / rows)
|
|
|
|
pad_w = (columns * tile_w - (columns - 1) * overlap) - width
|
|
pad_h = (rows * tile_h - (rows - 1) * overlap) - height
|
|
|
|
tiling_spec = {
|
|
"rows": rows,
|
|
"columns": columns,
|
|
"overlap": overlap,
|
|
"tile_w": tile_w,
|
|
"tile_h": tile_h,
|
|
"orig_w": width,
|
|
"orig_h": height,
|
|
"padded_w": width + pad_w,
|
|
"padded_h": height + pad_h
|
|
}
|
|
|
|
return (tiling_spec, tile_w, tile_h)
|
|
|
|
class SplitImage_Tiled:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"rows": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"columns": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"overlap": ("INT", {"default": 64, "min": 0, "max": 1024}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "*")
|
|
RETURN_NAMES = ("tiles", "rows", "columns", "overlap", "tiling_spec")
|
|
OUTPUT_IS_LIST = (True, False, False, False, False)
|
|
FUNCTION = "split"
|
|
CATEGORY = "SplitImage"
|
|
|
|
def split(self, image, rows, columns, overlap):
|
|
B, H, W, C = image.shape
|
|
|
|
tile_w = math.ceil((W + (columns - 1) * overlap) / columns)
|
|
tile_h = math.ceil((H + (rows - 1) * overlap) / rows)
|
|
|
|
pad_w = (columns * tile_w - (columns - 1) * overlap) - W
|
|
pad_h = (rows * tile_h - (rows - 1) * overlap) - H
|
|
|
|
if pad_w > 0 or pad_h > 0:
|
|
image_padded = image.permute(0, 3, 1, 2)
|
|
image_padded = F.pad(image_padded, (0, pad_w, 0, pad_h), mode='constant', value=0)
|
|
image = image_padded.permute(0, 2, 3, 1)
|
|
_, H, W, _ = image.shape
|
|
|
|
tiles = []
|
|
for r in range(rows):
|
|
for c in range(columns):
|
|
x1 = c * (tile_w - overlap)
|
|
y1 = r * (tile_h - overlap)
|
|
x2 = x1 + tile_w
|
|
y2 = y1 + tile_h
|
|
|
|
tile = image[:, y1:y2, x1:x2, :]
|
|
# Store each tile as a single-image tensor in the list
|
|
for i in range(tile.shape[0]):
|
|
tiles.append(tile[i:i+1])
|
|
|
|
tiling_spec = {
|
|
"rows": rows,
|
|
"columns": columns,
|
|
"overlap": overlap,
|
|
"tile_w": tile_w,
|
|
"tile_h": tile_h,
|
|
"orig_w": W - pad_w,
|
|
"orig_h": H - pad_h,
|
|
"padded_w": W,
|
|
"padded_h": H
|
|
}
|
|
|
|
return (tiles, rows, columns, overlap, tiling_spec)
|