polyHavenLoader.py

'''
@brief Module to fetch MaterialX assets from PolyHaven API and download them.
'''
import requests
import json
from pathlib import Path
import zipfile
import logging
import io
import tempfile
import os
import MaterialX as mx
 
HAVE_OPENIMAGEIO = False
try:
    import OpenImageIO as oiio
    import numpy as np
    HAVE_OPENIMAGEIO = True
except ImportError:
    print("OpenImageIO or numpy not installed. EXR image conversion not supported.")    
 
class PolyHavenLoader:
    """
    @brief Fetch MaterialX assets from PolyHaven API and download them.
 
    This class provides methods to fetch and download MaterialX assets and textures from PolyHaven.
    """
    def __init__(self):
        '''
        Initialize the PolyHavenLoader with API endpoints and headers.
        '''
        self.BASE_API = "https://api.polyhaven.com"
        self.ASSET_API = "https://api.polyhaven.com/assets"
        self.INFO_API = "https://api.polyhaven.com/info"
        self.FILES_API = "https://api.polyhaven.com/files"
        self.HEADERS = {
            "User-Agent": "MTLX_Polyaven_Loader/1.0",  # Required by PolyHaven API
        }
 
        self.logger = logging.getLogger('PolyH')
        logging.basicConfig(level=self.logger.info)
 
    def fetch_materialx_assets(self, max_items=1, download_id=None):
        '''
        Fetch MaterialX assets from PolyHaven API and filter them by resolution.
 
        @param max_items Maximum number of assets to fetch. If None, fetch all.
        @param download_id If set, only fetch asset with this ID.
        @return Tuple (materialx_assets, all_assets, filtered_polyhaven_assets):
            - materialx_assets: Dictionary of MaterialX assets with URLs and texture files.
            - all_assets: All assets returned by PolyHaven API.
            - filtered_polyhaven_assets: Filtered assets containing only MaterialX files.
        '''
        parameters = {
            "type": "textures"
        }
 
        resp = requests.get(self.ASSET_API, headers=self.HEADERS, params=parameters)
        resp.raise_for_status()
        all_assets = resp.json()
 
        materialx_assets = {}
        filtered_polyhaven_assets = {}
 
        item_count = 0;
        for id, data in all_assets.items():
 
            if download_id and id != download_id:
                #self.logger.info(f"Skipping asset id: '{id}' (not matching {download_id})")
                continue
 
            # Get the thumbnail
            thumbnail_url = data.get("thumbnail_url")
 
            resp = requests.get(f"{self.FILES_API}/{id}", headers=self.HEADERS)
            resp.raise_for_status()
            files_data = resp.json()
            #json_string = json.dumps(files_data, indent=4)
 
            # Remove all keys other than "mtlx"
            files_data = {k: v for k, v in files_data.items() if k == "mtlx"}            
 
            mtlx_files = files_data.get("mtlx", [])
            if mtlx_files:
 
                self.logger.info(f"Found MaterialX data for '{id}'") 
 
                # Look for 1K, 2K , 4K, and 8K versions
                resolutions = {
                    "1k": None,
                    "2k": None,
                    "4k": None,
                    "8k": None
                }
                for resolution_key in resolutions.keys():
                    res = mtlx_files.get(resolution_key, None)
                    if not res:
                        continue
 
                    n_k_mtlx = res.get("mtlx")
                    texture_struct = {}
                    if n_k_mtlx:
                        #self.logger.info(f"Found MaterialX files for '{one_k}'")
                        include_files = n_k_mtlx.get("include", {})
                        #self.logger.info(f"Found include files for '{id}': {one_k}")
                        for path, data in include_files.items():
                            texture_url = data.get("url")
                            #self.logger.info("Texture path:", path, "URL:", texture_url)
                            texture_struct[path] = texture_url
                    mtlx_url = n_k_mtlx.get("url")
                    res_id = id + '_' + resolution_key
                    if mtlx_url:
                        materialx_assets[res_id] = {
                            "url": mtlx_url,
                            "texture_files": texture_struct,
                            "thumbnail_url": thumbnail_url
                        }
                        json_string = json.dumps(materialx_assets[res_id], indent=4)
                        #self.logger.info(f"Found MaterialX for '{res_id}': {json_string}") 
                        #self.logger.info(f"Found MaterialX for '{res_id}'") 
                    # Create folder poly_have_data
 
                    # Save asset data to JSON file
                    #with open(f"polyhaven_data/{id}_data.json", "w") as f:
                    #    json.dump(asset_data, f, indent=4)
                    #    self.logger.info(f"Saved asset data for '{id}' to polyhaven_data/{id}_data.json")
 
                filtered_polyhaven_assets[id] = files_data
 
            # Halt if download_id is specified and matches the current asset ID        
            if download_id == id:
                break
 
            # Halt if max_items is specified and reached
            if not download_id and max_items:
                item_count += 1
                if item_count >= max_items:
                    break            
 
            #if "materialx" in formats:
            #    materialx_assets[slug] = formats["materialx"]
 
        return materialx_assets, all_assets, filtered_polyhaven_assets
 
    def download_asset(self, asset_list, convert_exr_to_png=True):
        '''
        Download MaterialX asset and its textures from PolyHaven.
 
        @param asset_list Dictionary of MaterialX assets with URLs and texture files.
        @param convert_exr_to_png If True, attempt to use PNG images instead of EXR if available other attempt 
        to convert using OpenImageIO if installed. Default is to use PNG if possible, and convert EXR to PNG if OpenImageIO is available.
        @return Tuple (id, mtlx_string, texture_binaries):
            - id: ID of the downloaded asset.
            - mtlx_string: The MaterialX document as a string.
            - texture_binaries: List of tuples (path, binary content) for textures and thumbnails.
        '''
        for id, asset in asset_list.items():
            url = asset.get("url")
            if not url:
                self.logger.info(f"No MaterialX URL found for '{id}'")
                continue
 
            resp = requests.get(url, headers=self.HEADERS)
            resp.raise_for_status()
            mtlx_string = resp.text
            self.logger.info(f"Download MaterialX document {url}, length: {len(mtlx_string)} characters")
 
            texture_binaries = []
            download_texture_names = []
            for path, texture_url in asset.get("texture_files", {}).items():
                # Get texture files
                ext = Path(path).suffix.lower()
                texture_content = None
 
                if ext == ".exr" and convert_exr_to_png:
                    # Replace /exr and .exr with /png and .png in the URL to check if a PNG version is available  
                    old_texture_url = texture_url  
                    texture_url = texture_url.replace("/exr/", "/png/").replace(".exr", ".png")
                    texture_resp = requests.get(texture_url, headers=self.HEADERS)
                    texture_resp.raise_for_status()            
                    texture_content = texture_resp.content
                    if texture_content:
                        self.logger.info(f"Download {texture_url} instead of {old_texture_url} SUCCESSFUL")
                        ext = ".png"
                        # Update extension in the path to .png
                        old_path = path
                        path = str(Path(path).with_suffix(ext))
                        self.logger.info(f"- Updating texture path from {old_path} to {path}")
 
                    else:
                        self.logger.info(f"Download {texture_url} instead of {old_texture_url} FAILED")
 
                if not texture_content:
                    self.logger.info(f"Download texture from {texture_url} ...")
                    texture_resp = requests.get(texture_url, headers=self.HEADERS)
                    texture_resp.raise_for_status()            
                    texture_content = texture_resp.content
 
                name = Path(path).stem
 
                if ext == ".exr":
                    if HAVE_OPENIMAGEIO and convert_exr_to_png:
                        self.logger.info(f"Converting EXR to PNG for texture: {path}")
                        # Write EXR bytes to a temporary file
                        with tempfile.NamedTemporaryFile(suffix=".exr", delete=False) as tmp_exr:
                            tmp_exr.write(texture_content)
                            tmp_exr_path = tmp_exr.name
                        try:
                            inbuf = oiio.ImageInput.open(tmp_exr_path)
                            if inbuf:
                                spec = inbuf.spec()
                                pixels = inbuf.read_image(format=oiio.UINT8)
                                inbuf.close()
                                # Write PNG to another temporary file
                                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_png:
                                    outbuf = oiio.ImageOutput.create(tmp_png.name)
                                    if outbuf:
                                        outbuf.open(tmp_png.name, spec)
                                        outbuf.write_image(pixels)
                                        outbuf.close()
                                        tmp_png.seek(0)
                                        png_bytes = tmp_png.read()
                                        png_name = f"{name}.png"
                                        texture_binaries.append((png_name, png_bytes))
 
                                        continue  # Skip adding the original EXR
                            else:
                                self.logger.info("Failed to read EXR with OpenImageIO")
                        finally:
                            os.remove(tmp_exr_path)
                            if 'tmp_png' in locals():
                                os.remove(tmp_png.name)
                    self.logger.info(f"  WARNING: EXR file present which may not be supported by MaterialX texture loader: {path}")
 
                # Get file name from path
                download_texture_name = path.split('/')[-1]
                download_texture_names.append(download_texture_name)
 
                texture_binaries.append((path, texture_content))
 
            thumbnail_url = asset.get("thumbnail_url")
            if thumbnail_url:
                self.logger.info(f"Download thumbnail from {thumbnail_url} ...")                
                thumbnail_resp = requests.get(thumbnail_url, headers=self.HEADERS)
                thumbnail_resp.raise_for_status()
                
                clean_url = thumbnail_url
                # Strip any ? or # from the URL
                clean_url = clean_url.split('?')[0].split('#')[0]
                clean_url = clean_url.split('/')[-1]  # Get the last part of the URL
                extension = Path(clean_url).suffix.lower()
                texture_binaries.append((f"{id}_thumbnail{extension}", thumbnail_resp.content))
 
            # Replace .exr with .png in mtlx_string
            for name in download_texture_names:
                extension = Path(name).suffix.lower()                
                exr_name = name.replace(extension, ".exr")
                # Replace exr_name with name in the mtlx_string
                mtlx_string = mtlx_string.replace(exr_name, name)
 
            before_mtlx_string = mtlx_string
            mtlx_string = mtlx_string.replace(".exr", ".png")
            if (before_mtlx_string != mtlx_string):
                self.logger.info(f"Updated MaterialX string to reference PNG texture instead of EXR for {path}")
                #self.logger.info(mtlx_string)
 
            return id, mtlx_string, texture_binaries
 
    def save_materialx_with_textures(self, id, mtlx_string, texture_binaries, data_folder, extract_zip=False):
        '''
        Save MaterialX string and texture binaries to a zip file.
 
        @param id The ID of the MaterialX asset.
        @param mtlx_string The MaterialX string content.
        @param texture_binaries List of tuples (path, binary content) for textures and thumbnails.
        @param data_folder Folder to save the zip file.
        @param extract_zip If True, extract the zip file after saving zip.
        @return None
        '''
        # Create a zip file with MaterialX and textures
        filename = f"{id}_materialx.zip"
        filename = Path(data_folder) / filename
        with zipfile.ZipFile(filename, "w") as zipf:
            # Write MaterialX file
            zipf.writestr(f"{id}.mtlx", mtlx_string)
            # Write texture files
            for path, content in texture_binaries:
                zipf.writestr(path, content)
        self.logger.info(f"Saved zip: {filename}")
 
        # Save zip contents to folder
        extract_folder = Path(data_folder) / f"{id}_materialx"
        if extract_zip:
            with zipfile.ZipFile(filename, "r") as zipf:
                zipf.extractall(extract_folder)
            self.logger.info(f"Extracted zip contents to folder: {extract_folder}")