Spaces:
Running
on
L40S
Running
on
L40S
File size: 4,056 Bytes
4450790 |
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 |
from math import ceil, sqrt
from typing import cast
import torch
import torchvision.transforms.functional as TF
from PIL import Image
from ..utils import hex_to_rgb, log, pil2tensor, tensor2pil
class MTB_TransformImage:
"""Save torch tensors (image, mask or latent) to disk, useful to debug things outside comfy
it return a tensor representing the transformed images with the same shape as the input tensor
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"x": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
),
"y": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
),
"zoom": (
"FLOAT",
{"default": 1.0, "min": 0.001, "step": 0.01},
),
"angle": (
"FLOAT",
{"default": 0, "step": 1, "min": -360, "max": 360},
),
"shear": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
),
"border_handling": (
["edge", "constant", "reflect", "symmetric"],
{"default": "edge"},
),
"constant_color": ("COLOR", {"default": "#000000"}),
},
}
FUNCTION = "transform"
RETURN_TYPES = ("IMAGE",)
CATEGORY = "mtb/transform"
def transform(
self,
image: torch.Tensor,
x: float,
y: float,
zoom: float,
angle: float,
shear: float,
border_handling="edge",
constant_color=None,
):
x = int(x)
y = int(y)
angle = int(angle)
log.debug(
f"Zoom: {zoom} | x: {x}, y: {y}, angle: {angle}, shear: {shear}"
)
if image.size(0) == 0:
return (torch.zeros(0),)
transformed_images = []
frames_count, frame_height, frame_width, frame_channel_count = (
image.size()
)
new_height, new_width = (
int(frame_height * zoom),
int(frame_width * zoom),
)
log.debug(f"New height: {new_height}, New width: {new_width}")
# - Calculate diagonal of the original image
diagonal = sqrt(frame_width**2 + frame_height**2)
max_padding = ceil(diagonal * zoom - min(frame_width, frame_height))
# Calculate padding for zoom
pw = int(frame_width - new_width)
ph = int(frame_height - new_height)
pw += abs(max_padding)
ph += abs(max_padding)
padding = [
max(0, pw + x),
max(0, ph + y),
max(0, pw - x),
max(0, ph - y),
]
constant_color = hex_to_rgb(constant_color)
log.debug(f"Fill Tuple: {constant_color}")
for img in tensor2pil(image):
img = TF.pad(
img, # transformed_frame,
padding=padding,
padding_mode=border_handling,
fill=constant_color or 0,
)
img = cast(
Image.Image,
TF.affine(
img, angle=angle, scale=zoom, translate=[x, y], shear=shear
),
)
left = abs(padding[0])
upper = abs(padding[1])
right = img.width - abs(padding[2])
bottom = img.height - abs(padding[3])
# log.debug("crop is [:,top:bottom, left:right] for tensors")
log.debug("crop is [left, top, right, bottom] for PIL")
log.debug(f"crop is {left}, {upper}, {right}, {bottom}")
img = img.crop((left, upper, right, bottom))
transformed_images.append(img)
return (pil2tensor(transformed_images),)
__nodes__ = [MTB_TransformImage]
|