stb_image.py

import ctypes
from ctypes import (
    c_char_p, c_wchar_p, c_int, c_ubyte, c_ushort, c_float,
    POINTER, c_void_p
)
import os
import platform
import sys

# ---------- Load DLL ----------
def load_stb_image_dll(path=None):
    if path:
        return ctypes.CDLL(path)

    system = platform.system()
    if system == 'Windows':
        return ctypes.CDLL('stb_image.dll')
    elif system == 'Linux':
        return ctypes.CDLL('./libstb_image.so')
    elif system == 'Darwin':
        return ctypes.CDLL('./libstb_image.dylib')
    else:
        raise RuntimeError(f"Unsupported OS: {system}")

lib = load_stb_image_dll()

# ---------- Type setup ----------
lib.load_image.argtypes = [c_char_p, POINTER(c_int), POINTER(c_int), POINTER(c_int), c_int]
lib.load_image.restype = POINTER(c_ubyte)

lib.load_16bit_image.argtypes = [c_char_p, POINTER(c_int), POINTER(c_int), POINTER(c_int), c_int]
lib.load_16bit_image.restype = POINTER(c_ushort)

lib.load_hdr_image.argtypes = [c_char_p, POINTER(c_int), POINTER(c_int), POINTER(c_int), c_int]
lib.load_hdr_image.restype = POINTER(c_float)

lib.free_image.argtypes = [c_void_p]
lib.free_image.restype = None

lib.get_image_info.argtypes = [c_char_p, POINTER(c_int), POINTER(c_int), POINTER(c_int)]
lib.get_image_info.restype = c_int

lib.is_hdr.argtypes = [c_char_p]
lib.is_hdr.restype = c_int

if platform.system() == 'Windows':
    lib.load_image_w.argtypes = [c_wchar_p, POINTER(c_int), POINTER(c_int), POINTER(c_int), c_int]
    lib.load_image_w.restype = POINTER(c_ubyte)

# ---------- High-level helper ----------
def load_image(filename, force_channels=0):
    width = c_int()
    height = c_int()
    channels = c_int()

    ptr = lib.load_image(filename.encode('utf-8'), ctypes.byref(width), ctypes.byref(height), ctypes.byref(channels), force_channels)
    if not ptr:
        reason = lib.get_failure_reason().decode('utf-8')
        raise RuntimeError(f"Failed to load image: {reason}")

    size = width.value * height.value * (force_channels or channels.value)
    data = ctypes.string_at(ptr, size)
    lib.free_image(ptr)

    return data, width.value, height.value, force_channels or channels.value

def load_hdr(filename):
    width = c_int()
    height = c_int()
    channels = c_int()

    ptr = lib.load_hdr_image(filename.encode('utf-8'), ctypes.byref(width), ctypes.byref(height), ctypes.byref(channels), 0)
    if not ptr:
        raise RuntimeError("Failed to load HDR image")

    size = width.value * height.value * channels.value * ctypes.sizeof(c_float)
    data = ctypes.string_at(ptr, size)
    lib.free_image(ptr)

    return data, width.value, height.value, channels.value

def get_image_info(filename):
    width = c_int()
    height = c_int()
    channels = c_int()

    result = lib.get_image_info(filename.encode('utf-8'), ctypes.byref(width), ctypes.byref(height), ctypes.byref(channels))
    if not result:
        return None
    return width.value, height.value, channels.value

def is_hdr_file(filename):
    return bool(lib.is_hdr(filename.encode('utf-8')))

# Optional: Wide-char loading on Windows
def load_image_wide(filename, force_channels=0):
    if platform.system() != 'Windows':
        raise RuntimeError("Wide-char loading is only supported on Windows")
    width = c_int()
    height = c_int()
    channels = c_int()

    ptr = lib.load_image_w(filename, ctypes.byref(width), ctypes.byref(height), ctypes.byref(channels), force_channels)
    if not ptr:
        raise RuntimeError("Failed to load image")
    size = width.value * height.value * (force_channels or channels.value)
    data = ctypes.string_at(ptr, size)
    lib.free_image(ptr)

    return data, width.value, height.value, force_channels or channels.value

# ---------- Example usage ----------
if __name__ == '__main__':
    file = 'test.png'
    if not os.path.exists(file):
        print("Put a test image (e.g., test.png) next to this script.")
        sys.exit(1)

    print(f"Loading: {file}")
    try:
        data, w, h, ch = load_image(file)
        print(f"Loaded {file}: {w}x{h} pixels, {ch} channels, {len(data)} bytes")
    except Exception as e:
        print("Error:", e)

    if is_hdr_file(file):
        print("HDR image detected.")
    else:
        print("Not an HDR image.")

    info = get_image_info(file)
    if info:
        print("Image info:", info)