From 4332bca744b60fec32c731af970f1db6ec5e950e Mon Sep 17 00:00:00 2001 From: shkr Date: Tue, 6 Aug 2024 17:41:34 +0530 Subject: [PATCH] fix --- .gitattributes | 2 + .gitignore | 2 + __init__.py | 3 + nodes/__init__.py | 27 +++ nodes/channel_select.py | 31 +++ nodes/config.py | 9 + nodes/fetch_surreal.py | 30 +++ nodes/florence.py | 459 +++++++++++++++++++++++++++++++++++++++ nodes/foreground_mask.py | 192 ++++++++++++++++ nodes/mask_image.py | 28 +++ nodes/save_surreal.py | 53 +++++ nodes/save_to_s3.py | 63 ++++++ nodes/surreal.py | 7 + pyproject.toml | 15 ++ requirements.txt | 4 + 15 files changed, 925 insertions(+) create mode 100644 .gitattributes create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 nodes/__init__.py create mode 100644 nodes/channel_select.py create mode 100644 nodes/config.py create mode 100644 nodes/fetch_surreal.py create mode 100644 nodes/florence.py create mode 100644 nodes/foreground_mask.py create mode 100644 nodes/mask_image.py create mode 100644 nodes/save_surreal.py create mode 100644 nodes/save_to_s3.py create mode 100644 nodes/surreal.py create mode 100644 pyproject.toml create mode 100644 requirements.txt diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..eba1110 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Auto detect text files and perform LF normalization +* text=auto \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6b1cf0c --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +config_.py \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3d998ea --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..d3d628e --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,27 @@ +from .channel_select import ChannelSelector +from .mask_image import MaskImage +from .save_surreal import SaveJsonToSurreal, SaveTextToSurreal +from .fetch_surreal import FetchJsonFromSurreal +from .foreground_mask import ForegroundMask +from .save_to_s3 import SaveImageToS3 + +NODE_CLASS_MAPPINGS = { + "ChannelSelector": ChannelSelector, + "MaskImage": MaskImage, + "SaveImageToS3": SaveImageToS3, + "SaveJsonToSurreal": SaveJsonToSurreal, + "SaveTextToSurreal": SaveTextToSurreal, + "FetchJsonFromSurreal": FetchJsonFromSurreal, + "ForegroundMask": ForegroundMask, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ChannelSelector":"ChannelSelector", + "MaskImage": "MaskImage", + "SaveImageToS3": "SaveImageToS3", + "SaveJsonToSurreal": "SaveJsonToSurreal", + "SaveTextToSurreal": "SaveTextToSurreal", + "FetchJsonFromSurreal": "FetchJsonFromSurreal", + "ForegroundMask": "ForegroundMask", +} diff --git a/nodes/channel_select.py b/nodes/channel_select.py new file mode 100644 index 0000000..f818b6a --- /dev/null +++ b/nodes/channel_select.py @@ -0,0 +1,31 @@ +import torch + +class ChannelSelector: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "channel": ("INT", { + "default": 0, + "min": 0, #Minimum value + "max": 100, #Maximum value + "step": 1, #Slider's step + "display": "number" # Cosmetic only: display as "number" or "slider" + }) + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "main" + + CATEGORY = "tensorops" + + def main(self, image, channel): + # Select the specified channel and add a new dimension at position 0 + mask = image[channel].unsqueeze(0) + + return (mask,) + diff --git a/nodes/config.py b/nodes/config.py new file mode 100644 index 0000000..278424d --- /dev/null +++ b/nodes/config.py @@ -0,0 +1,9 @@ +SURREAL_URL = "" +SURREAL_NAMESPACE = "" +SURREAL_USER = "" +SURREAL_PASSWORD = "" +DATALAKE_AWS_ENDPOINT_URL= "" +DATALAKE_AWS_ACCESS_KEY_ID= "" +DATALAKE_AWS_SECRET_ACCESS_KEY= "" +DATALAKE_AWS_REGION= "" +BUCKET = "" diff --git a/nodes/fetch_surreal.py b/nodes/fetch_surreal.py new file mode 100644 index 0000000..dcbcd25 --- /dev/null +++ b/nodes/fetch_surreal.py @@ -0,0 +1,30 @@ +from .surreal import surreal_connect + +SURREAL_TABLE = "processor" + +class FetchJsonFromSurreal: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "database": ("STRING", {"multiline": False}), + "id": ("STRING", {"multiline": False}), + "key": ("STRING", {"multiline": False}) + }, + } + + RETURN_TYPES = ("JSON",) + + FUNCTION = "main" + OUTPUT_NODE = True + CATEGORY = "database_ops" + + def main(self, database: str, id: str, key: str): + connection = surreal_connect(database) + query = f"SELECT {key} from {SURREAL_TABLE}:`{id}`" + results = connection.query(query).result + data = results[0][key] + print("results", data) + return [data] + diff --git a/nodes/florence.py b/nodes/florence.py new file mode 100644 index 0000000..407bd9f --- /dev/null +++ b/nodes/florence.py @@ -0,0 +1,459 @@ +import torch +import torchvision.transforms.functional as F +import io +import os +from typing import List +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from PIL import Image, ImageDraw, ImageColor, ImageFont +import random +import numpy as np +import re + +#workaround for unnecessary flash_attn requirement +from unittest.mock import patch +from transformers.dynamic_module_utils import get_imports + +def fixed_get_imports(filename: str | os.PathLike) -> list[str]: + if not str(filename).endswith("modeling_florence2.py"): + return get_imports(filename) + imports = get_imports(filename) + imports.remove("flash_attn") + return imports + + +import comfy.model_management as mm +from comfy.utils import ProgressBar +import folder_paths + +script_directory = os.path.dirname(os.path.abspath(__file__)) + +from transformers import AutoModelForCausalLM, AutoProcessor + +class DownloadAndLoadFlorence2Model: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ( + [ + 'microsoft/Florence-2-base', + 'microsoft/Florence-2-base-ft', + 'microsoft/Florence-2-large', + 'microsoft/Florence-2-large-ft', + 'HuggingFaceM4/Florence-2-DocVQA' + ], + { + "default": 'microsoft/Florence-2-base' + }), + "precision": ([ 'fp16','bf16','fp32'], + { + "default": 'fp16' + }), + "attention": ( + [ 'flash_attention_2', 'sdpa', 'eager'], + { + "default": 'sdpa' + }), + + }, + } + + RETURN_TYPES = ("FL2MODEL",) + RETURN_NAMES = ("florence2_model",) + FUNCTION = "loadmodel" + CATEGORY = "Florence2" + + def loadmodel(self, model, precision, attention): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] + + model_name = model.rsplit('/', 1)[-1] + model_path = os.path.join(folder_paths.models_dir, "LLM", model_name) + + if not os.path.exists(model_path): + print(f"Downloading Lumina model to: {model_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id=model, + local_dir=model_path, + local_dir_use_symlinks=False) + + print(f"using {attention} for attention") + with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement + model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True) + processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) + + florence2_model = { + 'model': model, + 'processor': processor, + 'dtype': dtype + } + + return (florence2_model,) + +def calculate_bounding_box(width, height, flat_points) -> List[float]: + """ + Calculate the bounding box for a polygon. + + Args: + flat_points (list of int): Flat list of x, y coordinates defining the polygon points. + + Returns: + tuple: (min_x, min_y, max_x, max_y) defining the bounding box. + """ + if not flat_points or len(flat_points) % 2 != 0: + raise ValueError("The list of points must be non-empty and have an even number of elements") + + x_coords = flat_points[0::2] + y_coords = flat_points[1::2] + + min_x = min(x_coords) + max_x = max(x_coords) + min_y = min(y_coords) + max_y = max(y_coords) + + return [min_x / width, min_y / height, max_x / width, max_y / height] + +class Florence2Run: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "florence2_model": ("FL2MODEL", ), + "text_input": ("STRING", {"default": "", "multiline": True}), + "task": ( + [ + 'region_caption', + 'dense_region_caption', + 'region_proposal', + 'caption', + 'detailed_caption', + 'more_detailed_caption', + 'caption_to_phrase_grounding', + 'referring_expression_segmentation', + 'ocr', + 'ocr_with_region', + 'docvqa' + ], + ), + "fill_mask": ("BOOLEAN", {"default": True}), + }, + "optional": { + "keep_model_loaded": ("BOOLEAN", {"default": False}), + "max_new_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}), + "num_beams": ("INT", {"default": 3, "min": 1, "max": 64}), + "do_sample": ("BOOLEAN", {"default": True}), + "output_mask_select": ("STRING", {"default": ""}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "STRING", "JSON") + RETURN_NAMES =("image", "mask", "caption", "data") + FUNCTION = "encode" + CATEGORY = "Florence2" + + def encode(self, image, text_input, florence2_model, task, fill_mask, keep_model_loaded=False, + num_beams=3, max_new_tokens=1024, do_sample=True, output_mask_select=""): + device = mm.get_torch_device() + _, height, width, _ = image.shape + offload_device = mm.unet_offload_device() + annotated_image_tensor = None + mask_tensor = None + processor = florence2_model['processor'] + model = florence2_model['model'] + dtype = florence2_model['dtype'] + model.to(device) + + colormap = ['blue','orange','green','purple','brown','pink','olive','cyan','red', + 'lime','indigo','violet','aqua','magenta','gold','tan','skyblue'] + + prompts = { + 'region_caption': '', + 'dense_region_caption': '', + 'region_proposal': '', + 'caption': '', + 'detailed_caption': '', + 'more_detailed_caption': '', + 'caption_to_phrase_grounding': '', + 'referring_expression_segmentation': '', + 'ocr': '', + 'ocr_with_region': '', + 'docvqa': '' + } + task_prompt = prompts.get(task, '') + + if (task not in ['referring_expression_segmentation', 'caption_to_phrase_grounding', 'docvqa']) and text_input: + raise ValueError("Text input (prompt) is only supported for 'referring_expression_segmentation', 'caption_to_phrase_grounding', and 'docvqa'") + + if text_input != "": + prompt = task_prompt + " " + text_input + else: + prompt = task_prompt + + image = image.permute(0, 3, 1, 2) + + out = [] + out_masks = [] + out_results = [] + out_data = [] + pbar = ProgressBar(len(image)) + for img in image: + image_pil = F.to_pil_image(img) + inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device) + + generated_ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=max_new_tokens, + do_sample=do_sample, + num_beams=num_beams, + ) + + results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] + print(results) + # cleanup the special tokens from the final list + if task == 'ocr_with_region': + clean_results = str(results) + cleaned_string = re.sub(r'|<[^>]*>', '\n', clean_results) + clean_results = re.sub(r'\n+', '\n', cleaned_string) + else: + clean_results = str(results) + clean_results = clean_results.replace('', '') + clean_results = clean_results.replace('', '') + + #return single string if only one image for compatibility with nodes that can't handle string lists + if len(image) == 1: + out_results = clean_results + else: + out_results.append(clean_results) + + W, H = image_pil.size + + parsed_answer = processor.post_process_generation(results, task=task_prompt, image_size=(W, H)) + + if task == 'region_caption' or task == 'dense_region_caption' or task == 'caption_to_phrase_grounding' or task == 'region_proposal': + fig, ax = plt.subplots(figsize=(W / 100, H / 100), dpi=100) + fig.subplots_adjust(left=0, right=1, top=1, bottom=0) + ax.imshow(image_pil) + bboxes = parsed_answer[task_prompt]['bboxes'] + labels = parsed_answer[task_prompt]['labels'] + + mask_indexes = [] + # Determine mask indexes outside the loop + if output_mask_select != "": + mask_indexes = [n for n in output_mask_select.split(",")] + print(mask_indexes) + else: + mask_indexes = [str(i) for i in range(len(bboxes))] + + # Initialize mask_layer only if needed + if fill_mask: + mask_layer = Image.new('RGB', image_pil.size, (0, 0, 0)) + mask_draw = ImageDraw.Draw(mask_layer) + + for index, (bbox, label) in enumerate(zip(bboxes, labels)): + # Modify the label to include the index + indexed_label = f"{index}.{label}" + + if fill_mask: + if str(index) in mask_indexes: + print("match index:", str(index), "in mask_indexes:", mask_indexes) + mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255)) + if label in mask_indexes: + print("match label") + mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255)) + + # Create a Rectangle patch + rect = patches.Rectangle( + (bbox[0], bbox[1]), # (x,y) - lower left corner + bbox[2] - bbox[0], # Width + bbox[3] - bbox[1], # Height + linewidth=1, + edgecolor='r', + facecolor='none', + label=indexed_label + ) + # Calculate text width with a rough estimation + text_width = len(label) * 6 # Adjust multiplier based on your font size + text_height = 12 # Adjust based on your font size + + # Initial text position + text_x = bbox[0] + text_y = bbox[1] - text_height # Position text above the top-left of the bbox + + # Adjust text_x if text is going off the left or right edge + if text_x < 0: + text_x = 0 + elif text_x + text_width > W: + text_x = W - text_width + + # Adjust text_y if text is going off the top edge + if text_y < 0: + text_y = bbox[3] # Move text below the bottom-left of the bbox if it doesn't overlap with bbox + + # Add the rectangle to the plot + ax.add_patch(rect) + facecolor = random.choice(colormap) if len(image) == 1 else 'red' + # Add the label + plt.text( + text_x, + text_y, + indexed_label, + color='white', + fontsize=12, + bbox=dict(facecolor=facecolor, alpha=0.5) + ) + if fill_mask: + mask_tensor = F.to_tensor(mask_layer) + mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + mask_tensor = mask_tensor.mean(dim=0, keepdim=True) + mask_tensor = mask_tensor.repeat(1, 1, 1, 3) + mask_tensor = mask_tensor[:, :, :, 0] + out_masks.append(mask_tensor) + + # Remove axis and padding around the image + ax.axis('off') + ax.margins(0,0) + ax.get_xaxis().set_major_locator(plt.NullLocator()) + ax.get_yaxis().set_major_locator(plt.NullLocator()) + fig.canvas.draw() + buf = io.BytesIO() + plt.savefig(buf, format='png', pad_inches=0) + buf.seek(0) + annotated_image_pil = Image.open(buf) + + annotated_image_tensor = F.to_tensor(annotated_image_pil) + out_tensor = annotated_image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + out.append(out_tensor) + + + pbar.update(1) + + plt.close(fig) + + elif task == 'referring_expression_segmentation': + # Create a new black image + mask_image = Image.new('RGB', (W, H), 'black') + mask_draw = ImageDraw.Draw(mask_image) + + predictions = parsed_answer[task_prompt] + + # Iterate over polygons and labels + for polygons, label in zip(predictions['polygons'], predictions['labels']): + color = random.choice(colormap) + for _polygon in polygons: + _polygon = np.array(_polygon).reshape(-1, 2) + # Clamp polygon points to image boundaries + _polygon = np.clip(_polygon, [0, 0], [W - 1, H - 1]) + if len(_polygon) < 3: + print('Invalid polygon:', _polygon) + continue + + _polygon = _polygon.reshape(-1).tolist() + + # Draw the polygon + if fill_mask: + overlay = Image.new('RGBA', image_pil.size, (255, 255, 255, 0)) + image_pil = image_pil.convert('RGBA') + draw = ImageDraw.Draw(overlay) + color_with_opacity = ImageColor.getrgb(color) + (180,) + draw.polygon(_polygon, outline=color, fill=color_with_opacity, width=3) + image_pil = Image.alpha_composite(image_pil, overlay) + else: + draw = ImageDraw.Draw(image_pil) + draw.polygon(_polygon, outline=color, width=3) + + #draw mask + mask_draw.polygon(_polygon, outline="white", fill="white") + + image_tensor = F.to_tensor(image_pil) + image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + out.append(image_tensor) + + mask_tensor = F.to_tensor(mask_image) + mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + mask_tensor = mask_tensor.mean(dim=0, keepdim=True) + mask_tensor = mask_tensor.repeat(1, 1, 1, 3) + mask_tensor = mask_tensor[:, :, :, 0] + out_masks.append(mask_tensor) + pbar.update(1) + + elif task == 'ocr_with_region': + try: + font = ImageFont.load_default().font_variant(size=24) + except: + font = ImageFont.load_default() + predictions = parsed_answer[task_prompt] + scale = 1 + draw = ImageDraw.Draw(image_pil) + bboxes, labels = predictions['quad_boxes'], predictions['labels'] + + for box, label in zip(bboxes, labels): + bbox = calculate_bounding_box(width, height, box) + out_data.append({"label": label, "polygon": box, "box": bbox}) + color = random.choice(colormap) + new_box = (np.array(box) * scale).tolist() + draw.polygon(new_box, width=3, outline=color) + draw.text((new_box[0]+8, new_box[1]+2), + "{}".format(label), + align="right", + font=font, + fill=color) + + image_tensor = F.to_tensor(image_pil) + image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + out.append(image_tensor) + + elif task == 'docvqa': + if text_input == "": + raise ValueError("Text input (prompt) is required for 'docvqa'") + prompt = " " + text_input + + inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device) + generated_ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=max_new_tokens, + do_sample=do_sample, + num_beams=num_beams, + ) + + results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] + clean_results = results.replace('', '').replace('', '') + + if len(image) == 1: + out_results = clean_results + else: + out_results.append(clean_results) + + out.append(F.to_tensor(image_pil).unsqueeze(0).permute(0, 2, 3, 1).cpu().float()) + + pbar.update(1) + + if len(out) > 0: + out_tensor = torch.cat(out, dim=0) + else: + out_tensor = torch.zeros((1, 64,64, 3), dtype=torch.float32, device="cpu") + if len(out_masks) > 0: + out_mask_tensor = torch.cat(out_masks, dim=0) + else: + out_mask_tensor = torch.zeros((1,64,64), dtype=torch.float32, device="cpu") + + if not keep_model_loaded: + print("Offloading model...") + model.to(offload_device) + mm.soft_empty_cache() + + return (out_tensor, out_mask_tensor, out_results, out_data) + +NODE_CLASS_MAPPINGS = { + "DownloadAndLoadFlorence2Model": DownloadAndLoadFlorence2Model, + "Florence2Run": Florence2Run, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "DownloadAndLoadFlorence2Model": "DownloadAndLoadFlorence2Model", + "Florence2Run": "Florence2Run", +} \ No newline at end of file diff --git a/nodes/foreground_mask.py b/nodes/foreground_mask.py new file mode 100644 index 0000000..c925c7d --- /dev/null +++ b/nodes/foreground_mask.py @@ -0,0 +1,192 @@ +from typing import List +from PIL import Image, ImageDraw +import numpy as np +from collections import Counter +import torchvision.transforms.functional as F +import json + +def calculate_bounding_box(points) -> List[float]: + """ + Calculate the bounding box for a polygon. + + Args: + flat_points (list of int): Flat list of x, y coordinates defining the polygon points. + + Returns: + tuple: (min_x, min_y, max_x, max_y) defining the bounding box. + """ + if not points or len(points) % 2 != 0: + raise ValueError("The list of points must be non-empty and have an even number of elements") + + x_coords = points[0::2] + y_coords = points[1::2] + + min_x = min(x_coords) + max_x = max(x_coords) + min_y = min(y_coords) + max_y = max(y_coords) + + return [min_x, min_y, max_x, max_y] + + +def find_mode_color(image: Image.Image): + """ + Identify the most frequent (mode) color in a PIL image. + + Parameters: + image_path (str): The path to the input image. + + Returns: + tuple: The mode color in the image as an (R, G, B) tuple. + """ + # Convert image to RGB mode if it's not already + image = image.convert('RGB') + + # Get the list of pixels + pixels = list(image.getdata()) + + # Use Counter to count the frequency of each color + counter = Counter(pixels) + + # Find the most common color + mode_color = counter.most_common(1)[0][0] + + return mode_color + +def separate_foreground_background(image): + """ + Separate the Pillow image into foreground and background using the mode color and distance clustering. + + Parameters: + image_path (str): The path to the input image. + output_foreground (str): The path to save the foreground image. + output_background (str): The path to save the background image. + + Returns: + None + """ + # Convert image to RGBA mode to handle transparency + image = image.convert('RGBA') + pixels = np.array(image) + + # Calculate the Euclidean distance of each pixel to the mode color + background_color = find_mode_color(image) + print("Background color:", background_color) + mode_color_array = np.array(background_color) + distances = np.linalg.norm(pixels[:, :, :3] - mode_color_array, axis=2) + + # Determine the threshold distance for clustering + threshold_distance = np.mean(distances) + + print("Threshold distance:", threshold_distance) + # Create masks for foreground and background + foreground_mask = distances > threshold_distance + background_mask = distances <= threshold_distance + + # Create empty arrays for the new images + foreground_image = np.zeros_like(pixels) + background_image = np.zeros_like(pixels) + + # Copy the pixels to the new images based on the masks + foreground_image[foreground_mask] = pixels[foreground_mask] + background_image[background_mask] = pixels[background_mask] + + # Find the fg color + fg_color = find_mode_color(Image.fromarray(foreground_image, 'RGBA')) + + # Set foreground pixels with alpha == 255 to black + alpha_channel = foreground_image[:, :, 3] == 255 + foreground_image[alpha_channel, :3] = [255, 255, 255] + foreground_image[:, :, 3] = 255 + + # Convert back to PIL images + foreground_image = Image.fromarray(foreground_image, 'RGBA') + background_image = Image.fromarray(background_image, 'RGBA') + + # Invert Foreground As White + # foreground_image = ImageOps.invert(foreground_image.convert("RGB")) + + return foreground_image, fg_color + +def crop_polygon(image, points): + """ + Create a white mask on a black image of size width x height using a list of polygon points. + + Args: + points (list of tuples): List of (x, y) tuples defining the polygon points. + width (int): Width of the image. + height (int): Height of the image. + + Returns: + Image: Pillow Image object with the polygon mask. + """ + x_min, y_min, x_max, y_max = calculate_bounding_box(points) + image_crop = image.crop((x_min, y_min, x_max, y_max)) + return image_crop + +def mask_polygon(image, points): + """ + Crop a polygon from a Pillow image. + + Args: + image (PIL.Image): The input image. + flat_points (list of int): Flat list of x, y coordinates defining the polygon points. + + Returns: + PIL.Image: Cropped image of the polygon. + """ + if not points or len(points) % 2 != 0: + raise ValueError("The list of points must be non-empty and have an even number of elements") + + # Create a mask + mask = Image.new('L', image.size, 0) + draw = ImageDraw.Draw(mask) + new_box = (np.array(points) * 1.0).tolist() + draw.polygon(new_box, fill="white") + + # Apply the mask to the image + masked_image = Image.composite(image.convert("RGBA"), mask.convert("RGBA"), mask) + return masked_image + + +import torch + +class ForegroundMask: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "json_data": ("JSON",), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "main" + + CATEGORY = "tensorops" + + def main(self, image: torch.Tensor, json_data: str): + print(json_data) + image = image.permute(0, 3, 1, 2) + print("Image", image.shape) + print("JsonData", json_data, len(json_data)) + image_pil = F.to_pil_image(image[0]) + full_image = Image.new("RGBA", image_pil.size, (0, 0, 0, 255)) + print("ForegroundMask") + print("ImageShape", image.shape) + print("OutImageShape", full_image.size) + for item in json_data: + points = item["polygon"] + print("polygon", points) + masked_image = mask_polygon(image_pil, points) + print("NewMaskImageShape", masked_image.size) + masked_image_crop = crop_polygon(image_pil, points) + fg_image, fg_color = separate_foreground_background(masked_image_crop) + x_min, y_min, x_max, y_max = calculate_bounding_box(points) + full_image.paste(fg_image, (int(x_min), int(y_min))) + out_image = F.to_tensor(full_image) + return (out_image,) + diff --git a/nodes/mask_image.py b/nodes/mask_image.py new file mode 100644 index 0000000..f6a578c --- /dev/null +++ b/nodes/mask_image.py @@ -0,0 +1,28 @@ +import torch + +class MaskImage: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "main" + + CATEGORY = "tensorops" + + def main(self, image: torch.Tensor, mask: torch.Tensor): + mask = mask.unsqueeze(-1) + new_image = image * mask + print("MaskImage") + print("ImageShape", image.shape) + print("MaskShape", mask.shape) + print("NewImageShape", new_image.shape) + return (new_image,) + diff --git a/nodes/save_surreal.py b/nodes/save_surreal.py new file mode 100644 index 0000000..401b7b9 --- /dev/null +++ b/nodes/save_surreal.py @@ -0,0 +1,53 @@ +from .surreal import surreal_connect + +SURREAL_TABLE = "processor" + +class SaveJsonToSurreal: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "database": ("STRING", {"multiline": False}), + "json": ("JSON",), + "id": ("STRING", {"multiline": False}), + "key": ("STRING", {"multiline": False}) + }, + } + + RETURN_TYPES = () + + FUNCTION = "main" + OUTPUT_NODE = True + CATEGORY = "database_ops" + + def main(self, database: str, id: str, key: str, json: str): + connection = surreal_connect(database) + query = f"UPDATE {SURREAL_TABLE}:`{id}` CONTENT {{{key}: {json}}};" + connection.query(query) + return () + +class SaveTextToSurreal: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "database": ("STRING", {"multiline": False}), + "text": ("STRING",{"forceInput": True}), + "id": ("STRING", {"multiline": False}), + "key": ("STRING", {"multiline": False}) + }, + } + + RETURN_TYPES = () + + FUNCTION = "main" + OUTPUT_NODE = True + CATEGORY = "database_ops" + + def main(self, database: str, id: str, key: str, text: str): + connection = surreal_connect(database) + query = f"UPDATE {SURREAL_TABLE}:`{id}` CONTENT {{{key}: '{text}'}};" + connection.query(query) + return () \ No newline at end of file diff --git a/nodes/save_to_s3.py b/nodes/save_to_s3.py new file mode 100644 index 0000000..afde027 --- /dev/null +++ b/nodes/save_to_s3.py @@ -0,0 +1,63 @@ +from .config import DATALAKE_AWS_ACCESS_KEY_ID, DATALAKE_AWS_ENDPOINT_URL, DATALAKE_AWS_SECRET_ACCESS_KEY, BUCKET +import boto3 +from botocore.client import Config +import torch +import PIL +import PIL.Image +import io +import logging + + +logger = logging.getLogger(__name__) + +S3_RESOURCE = None + +def init_s3(): + global S3_RESOURCE + if not S3_RESOURCE: + S3_ENDPOINT_URL = DATALAKE_AWS_ENDPOINT_URL + AWS_ACCESS_KEY_ID = DATALAKE_AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY = DATALAKE_AWS_SECRET_ACCESS_KEY + # Initialize the S3 client + S3_RESOURCE = boto3.resource( + "s3", + endpoint_url=S3_ENDPOINT_URL, + aws_access_key_id=AWS_ACCESS_KEY_ID, + aws_secret_access_key=AWS_SECRET_ACCESS_KEY, + config=Config(signature_version="s3v4"), + ) + +def store_image(key: str, image: PIL.Image.Image): + global S3_RESOURCE + init_s3() + try: + s3_bucket = S3_RESOURCE.Bucket(BUCKET) + file_content = io.BytesIO() + image.save(file_content, format="webp") + file_content.seek(0) + logger.info(f"StoreImage: {key} in {BUCKET}") + s3_bucket.put_object(Key=key, Body=file_content) + except Exception as e: + logger.error(f"StoreContentError {key}: {e}") + +class SaveImageToS3: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "database": ("STRING", {"multiline": False}), + "key": ("STRING", {"multiline": False}), + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = () + + FUNCTION = "main" + OUTPUT_NODE = True + CATEGORY = "database_ops" + + def main(self, database: str, key: str, image: torch.Tensor): + store_image(f"{database}/{key}.webp", PIL.Image.fromarray(image.mul(255).byte().permute(1, 2, 0).cpu().numpy())) + return () \ No newline at end of file diff --git a/nodes/surreal.py b/nodes/surreal.py new file mode 100644 index 0000000..7ebf14e --- /dev/null +++ b/nodes/surreal.py @@ -0,0 +1,7 @@ +from surrealist import Surreal +from .config import SURREAL_URL, SURREAL_NAMESPACE, SURREAL_USER, SURREAL_PASSWORD + +def surreal_connect(database: str): + surreal_client = Surreal(SURREAL_URL, namespace=SURREAL_NAMESPACE, database=database, credentials=(SURREAL_USER, SURREAL_PASSWORD), use_http=True, timeout=10) + surreal_connection = surreal_client.connect() + return surreal_connection \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ba593b6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "comfyui-tensorops" +description = "Nodes to perform tensor operations in ComfyUI" +version = "1.0.0" +license = "MIT" +dependencies = ["transformers>=4.38.0"] + +[project.urls] +Repository = "https://github.com/un-seen/comfyui-tensorops" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "shkr" +DisplayName = "comfyui-tensorops" +Icon = "" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..2cf3f59 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +tqdm +numpy +surrealist +boto3==1.34.86 \ No newline at end of file