Upload BAM_img_to_bam.py

BAM_img_to_bam.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, List
+
+from PIL import Image, ExifTags
+import folder_paths
+
+
+_BAM_RE = re.compile(
+    r"GPT_BAM_START###.*?###GPT_BAM_END",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+
+_IMAGE_EXTS = {
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".webp",
+    ".bmp",
+    ".tif",
+    ".tiff",
+    ".gif",
+}
+
+
+def _safe_decode_bytes(data: bytes) -> str:
+    if not data:
+        return ""
+
+    # EXIF UserComment prefixes
+    prefixes = [
+        b"ASCII\x00\x00\x00",
+        b"UNICODE\x00",
+        b"JIS\x00\x00\x00\x00\x00",
+    ]
+    for p in prefixes:
+        if data.startswith(p):
+            data = data[len(p):]
+            break
+
+    for enc in ("utf-8", "utf-16", "utf-16le", "latin-1"):
+        try:
+            return data.decode(enc, errors="ignore")
+        except Exception:
+            pass
+
+    return ""
+
+
+def _collect_strings(value: Any, out: List[str]) -> None:
+    if value is None:
+        return
+
+    if isinstance(value, str):
+        if value:
+            out.append(value)
+        return
+
+    if isinstance(value, bytes):
+        decoded = _safe_decode_bytes(value)
+        if decoded:
+            out.append(decoded)
+        return
+
+    if isinstance(value, dict):
+        for k, v in value.items():
+            if isinstance(k, str) and k:
+                out.append(k)
+            _collect_strings(v, out)
+        return
+
+    if isinstance(value, (list, tuple, set)):
+        for item in value:
+            _collect_strings(item, out)
+        return
+
+    # Ignore plain numbers / objects by default
+    return
+
+
+def _extract_from_text_blob(text: str) -> str:
+    if not text:
+        return ""
+    m = _BAM_RE.search(text)
+    if not m:
+        return ""
+    return m.group(0).replace("\x00", "").strip()
+
+
+def _extract_from_pil_metadata(image_path: str) -> str:
+    text_parts: List[str] = []
+
+    with Image.open(image_path) as img:
+        # Standard PIL info dict
+        _collect_strings(getattr(img, "info", {}), text_parts)
+
+        # PNG text chunks if present
+        if hasattr(img, "text"):
+            _collect_strings(getattr(img, "text"), text_parts)
+
+        # EXIF
+        try:
+            exif = img.getexif()
+            if exif:
+                named_exif = {}
+                for tag_id, value in exif.items():
+                    tag_name = ExifTags.TAGS.get(tag_id, str(tag_id))
+                    named_exif[tag_name] = value
+                _collect_strings(named_exif, text_parts)
+        except Exception:
+            pass
+
+    combined = "\n".join(text_parts)
+    return _extract_from_text_blob(combined)
+
+
+def _extract_from_raw_file(image_path: str) -> str:
+    try:
+        with open(image_path, "rb") as f:
+            raw = f.read()
+    except Exception:
+        return ""
+
+    # Try direct byte search first
+    try:
+        start = raw.find(b"GPT_BAM_START###")
+        if start != -1:
+            end_marker = b"###GPT_BAM_END"
+            end = raw.find(end_marker, start)
+            if end != -1:
+                end += len(end_marker)
+                found = raw[start:end].decode("utf-8", errors="ignore").replace("\x00", "").strip()
+                if found:
+                    return found
+    except Exception:
+        pass
+
+    # Decode whole file with a few encodings and regex-search
+    for enc in ("utf-8", "utf-16", "utf-16le", "latin-1"):
+        try:
+            text = raw.decode(enc, errors="ignore")
+        except Exception:
+            continue
+
+        found = _extract_from_text_blob(text)
+        if found:
+            return found
+
+    return ""
+
+
+class img_meta_to_BAM:
+    @classmethod
+    def INPUT_TYPES(cls):
+        input_dir = folder_paths.get_input_directory()
+        files = []
+
+        if os.path.isdir(input_dir):
+            for f in os.listdir(input_dir):
+                full = os.path.join(input_dir, f)
+                if os.path.isfile(full):
+                    ext = os.path.splitext(f)[1].lower()
+                    if ext in _IMAGE_EXTS:
+                        files.append(f)
+
+        return {
+            "required": {
+                "image": (sorted(files), {"image_upload": True}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("BAM_format",)
+    FUNCTION = "extract_bam"
+    CATEGORY = "BAM"
+
+    def extract_bam(self, image):
+        image_path = folder_paths.get_annotated_filepath(image)
+
+        bam = ""
+
+        # First try proper metadata extraction
+        try:
+            bam = _extract_from_pil_metadata(image_path)
+        except Exception:
+            bam = ""
+
+        # Fallback: raw file scan
+        if not bam:
+            bam = _extract_from_raw_file(image_path)
+
+        return (bam,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "img_meta_to_BAM": img_meta_to_BAM,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "img_meta_to_BAM": "img_meta_to_BAM",
+}