BAM_img_to_bam.py

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",
}