From 51d8bc4490caa1cc478b82b5fddbd99b530f75a4 Mon Sep 17 00:00:00 2001 From: Aleix Date: Mon, 30 Dec 2024 19:28:01 +0100 Subject: [PATCH] initial commit --- __init__.py | 18 +++++ gcp_storage_tools.py | 167 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 __init__.py create mode 100644 gcp_storage_tools.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..0eb44c2 --- /dev/null +++ b/__init__.py @@ -0,0 +1,18 @@ +import importlib +import pip + +def install_gcp_storage(): + print("Checking for google-cloud-storage installation..") + try: + importlib.import_module("google.cloud.storage") + except ImportError: + print("google-cloud-storage not found. Installing...") + pip.main(["install", "google-cloud-storage"]) + + +install_gcp_storage() + +from .gcp_storage_tools import NODE_CLASS_MAPPINGS +from .gcp_storage_tools import NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/gcp_storage_tools.py b/gcp_storage_tools.py new file mode 100644 index 0000000..403be7a --- /dev/null +++ b/gcp_storage_tools.py @@ -0,0 +1,167 @@ +import os +import io +import numpy as np +from PIL import Image, ImageOps +from google.cloud import storage +import torch + +class GCPReadImageNode: + def __init__(self): + # Build a temp folder 2 levels above this script + script_dir = os.path.dirname(os.path.abspath(__file__)) + self.output_dir = os.path.normpath(os.path.join(script_dir, "temp")) + os.makedirs(self.output_dir, exist_ok=True) + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "bucket_name": ("STRING", {"default": "my-bucket"}), + "bucket_path": ("STRING", {"default": "some/folder"}), + "file_name": ("STRING", {"default": "my_image.png"}), + "gcp_service_json": ("STRING", {"default": "/path/to/service_account.json"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "download_from_gcp" + OUTPUT_NODE = True # Mark as output node + CATEGORY = "gcp_storage" + + + def download_from_gcp(self, bucket_name, bucket_path, file_name, gcp_service_json): + try: + + script_dir = os.path.dirname(os.path.abspath(__file__)) # folder where this script lives + if not os.path.isabs(gcp_service_json): + gcp_service_json = os.path.join(script_dir, gcp_service_json) + + # Check if credentials file exists + if not os.path.exists(gcp_service_json): + raise FileNotFoundError(f"Credential file not found: {gcp_service_json}") + + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_service_json + + + + if not file_name.lower().endswith(".png"): + file_name += ".png" + + local_file_path = os.path.join(self.output_dir, os.path.basename(file_name)) + gcp_path = f"{bucket_path}/{os.path.basename(file_name)}".strip("/") + + # Download from GCP to local file + storage_client = storage.Client() + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(gcp_path) + print(f"[GCP READ] Downloading gs://{bucket_name}/{gcp_path} → {local_file_path}") + blob.download_to_filename(local_file_path) + + # Load and process image + i = Image.open(local_file_path) + i = ImageOps.exif_transpose(i) + image = i.convert("RGB") + + # Convert to numpy array and normalize + image_np = np.array(image).astype(np.float32) / 255.0 + + # Ensure shape is [batch, height, width, channels] + if len(image_np.shape) == 3: + image_np = image_np[None, ...] + + image = torch.from_numpy(image_np) + image = image.contiguous().float() + + os.remove(local_file_path) + + return (image,) + except Exception as e: + print(str(e)) + raise e + + +class GCPWriteImageNode: + def __init__(self): + self.compress_level = 4 + self.type = "output" + script_dir = os.path.dirname(os.path.abspath(__file__)) + self.output_dir = os.path.normpath(os.path.join(script_dir, "out")) + os.makedirs(self.output_dir, exist_ok=True) + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "bucket_name": ("STRING", {"default": "my-bucket"}), + "bucket_path": ("STRING", {"default": "some/folder"}), + "file_name": ("STRING", {"default": "my_output.png"}), + "gcp_service_json": ("STRING", {"default": "/path/to/service_account.json"}), + } + } + + # We'll define two outputs: + # 1) The original image if success + # 2) A string error message if failure + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("images", "error_message") + FUNCTION = "store_image_in_gcp" + OUTPUT_NODE = True # Mark as output node + CATEGORY = "gcp_storage" + + def store_image_in_gcp(self, images, bucket_name, bucket_path, file_name, gcp_service_json): + error_message = "" + try: + script_dir = os.path.dirname(os.path.abspath(__file__)) # folder where this script lives + if not os.path.isabs(gcp_service_json): + gcp_service_json = os.path.join(script_dir, gcp_service_json) + + if not os.path.exists(gcp_service_json): + raise FileNotFoundError(f"Credential file not found: {gcp_service_json}") + + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_service_json + + subfolder = os.path.dirname(os.path.normpath(file_name)) + full_output_folder = os.path.join(self.output_dir, subfolder) + full_file_path = os.path.join(full_output_folder, file_name) + + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + img.save(os.path.join(full_output_folder, file_name), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file_name, + "subfolder": subfolder, + "type": self.type + }) + + storage_client = storage.Client() + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(f"{bucket_path}/{file_name}") + print(f"Uploading blob to {bucket_name}/{bucket_path}/{file_name}..") + blob.upload_from_filename(full_file_path) + + # Cleanup the temp file if you want + os.remove(full_file_path) + + return {"ui": {"images": results}} + except Exception as e: + error_message = f"Upload failed: {str(e)}" + print(error_message) + # Return None for image, and the error message in the second output + return (None, error_message) + + + +NODE_CLASS_MAPPINGS = { + "GCPReadImageNode": GCPReadImageNode, + "GCPWriteImageNode": GCPWriteImageNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GCPReadImageNode": "GCP: Read Image", + "GCPWriteImageNode": "GCP: Write Images", +} \ No newline at end of file