wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-30 08:10:21 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

Add expand option to rotate

Browse Source
This commit is contained in:
missionfloyd 2023-04-10 21:20:42 -06:00 committed by GitHub
parent bc54b69c59
commit 9b40cd3f89
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
comfy_extras/nodes_post_processing.py
View File

@ -270,6 +270,7 @@ class Rotate:
"Bilinear", "Bilinear",
"Bicubic", "Bicubic",
],), ],),
"expand": (["disabled", "enabled"],),
"fill_color": ("STRING", {"default": "#000000"}), "fill_color": ("STRING", {"default": "#000000"}),
}, },
} }
@ -279,9 +280,8 @@ class Rotate:
CATEGORY = "image/postprocessing" CATEGORY = "image/postprocessing"
def rotate(self, image: torch.Tensor, angle: int, resample: str, fill_color: str): def rotate(self, image: torch.Tensor, angle: int, resample: str, expand: str, fill_color: str):
batch_size, height, width, _ = image.shape batch_size, _, _, _ = image.shape
result = torch.zeros_like(image)
resamplers = { resamplers = {
"Nearest Neighbor": Image.Resampling.NEAREST, "Nearest Neighbor": Image.Resampling.NEAREST,
@ -289,29 +289,30 @@ class Rotate:
"Bicubic": Image.Resampling.BICUBIC, "Bicubic": Image.Resampling.BICUBIC,
} }
for b in range(batch_size): tensor_image = image[0]
tensor_image = image[b] img = (tensor_image * 255).to(torch.uint8).numpy()
img = (tensor_image * 255).to(torch.uint8).numpy() pil_image = Image.fromarray(img, mode='RGB')
pil_image = Image.fromarray(img, mode='RGB')
expand = True if expand == "enabled" else False
fill_color = fill_color or "#000000" fill_color = fill_color or "#000000"
def parse_palette(color_str): def parse_palette(color_str):
if re.match(r'^#[a-fA-F0-9]{6}$', color_str) or color_str.lower() in ImageColor.colormap: if re.match(r'^#[a-fA-F0-9]{6}$', color_str) or color_str.lower() in ImageColor.colormap:
return ImageColor.getrgb(color_str) return ImageColor.getrgb(color_str)
color_rgb = re.match(r'^\(?(\d{1,3}),(\d{1,3}),(\d{1,3})\)?$', color_str) color_rgb = re.match(r'^\(?(\d{1,3}),(\d{1,3}),(\d{1,3})\)?$', color_str)
if color_rgb and int(color_rgb.group(1)) <= 255 and int(color_rgb.group(2)) <= 255 and int(color_rgb.group(3)) <= 255: if color_rgb and int(color_rgb.group(1)) <= 255 and int(color_rgb.group(2)) <= 255 and int(color_rgb.group(3)) <= 255:
return tuple(map(int, re.findall(r'\d{1,3}', color_str))) return tuple(map(int, re.findall(r'\d{1,3}', color_str)))
else: else:
raise ValueError(f"Invalid color format: {color_str}") raise ValueError(f"Invalid color format: {color_str}")
color = fill_color.replace(" ", "") color = fill_color.replace(" ", "")
color = parse_palette(color) color = parse_palette(color)
rotated_image = pil_image.rotate(angle=angle, resample=resamplers[resample], expand=False, fillcolor=color) rotated_image = pil_image.rotate(angle=angle, resample=resamplers[resample], expand=expand, fillcolor=color)
height, width = rotated_image.size
rotated_array = torch.tensor(np.array(rotated_image.convert("RGB"))).float() / 255 result = torch.zeros(batch_size, height, width, 3)
result[b] = rotated_array rotated_array = torch.tensor(np.array(rotated_image.convert("RGB"))).float() / 255
result[0] = rotated_array
return (result,) return (result,)