run.py

# -----------------------------------------------------------------------------
# Created by: Quentin Lengele (16/10/2025)
# -----------------------------------------------------------------------------

import sys
import os
from pathlib import Path
import cv2
import time
import json
import math
import shutil
import warnings
import numpy as np
import torch
from PIL import Image, ImageFilter


import utils.imgops as ops
import utils.architecture.architecture as arch

warnings.filterwarnings('ignore')

# ====================================================================================================================

NORMAL_MAP_MODEL = 'utils/models/1x_NormalMapGenerator-CX-Lite_200000_G.pth'
OTHER_MAP_MODEL = 'utils/models/1x_FrankenMapGenerator-CX-Lite_215000_G.pth'

DEVICE = torch.device('cuda')

MATGEN_MODELS = []

# ====================================================================================================================

def process_matgen(img, model):
    global DEVICE
    img = img * 1. / np.iinfo(img.dtype).max
    img = img[:, :, [2, 1, 0]]
    img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
    img_LR = img.unsqueeze(0)
    img_LR = img_LR.to(DEVICE)

    output = model(img_LR).data.squeeze(0).float().cpu().clamp_(0, 1).numpy()
    output = output[[2, 1, 0], :, :]
    output = np.transpose(output, (1, 2, 0))
    output = (output * 255.).round()
    return output


def load_matgen_model(model_path):
    global DEVICE
    state_dict = torch.load(model_path)
    model = arch.RRDB_Net(3, 3, 32, 12, gc=32, upscale=1, norm_type=None, act_type='leakyrelu',
                            mode='CNA', res_scale=1, upsample_mode='upconv')
    model.load_state_dict(state_dict, strict=True)
    del state_dict
    model.eval()
    for k, v in model.named_parameters():
        v.requires_grad = False
    return model.to(DEVICE)


def generate_pbr(albedo_filepath,
                 output_dir,
                 override=False,
                 tile_size=512,
                 seamless=False,
                 mirror=False,
                 replicate=False):

    global MATGEN_MODELS

    file_path = Path(albedo_filepath)
    basename = file_path.stem
    file_ext = file_path.suffix

    # output_folder = file_path.parent
    output_folder = output_dir

    # copy original albedo and name it _BaseColor

    destination = Path(f"{output_folder}/{basename}_BaseColor{file_ext}")
    shutil.copy2(file_path, destination)

    # read image
    try:
        img = cv2.imread(albedo_filepath, cv2.cv2.IMREAD_COLOR)
    except:
        img = cv2.imread(albedo_filepath, cv2.IMREAD_COLOR)

    # Seamless modes
    if seamless:
        img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_WRAP)
    elif mirror:
        img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_REFLECT_101)
    elif replicate:
        img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_REPLICATE)

    img_height, img_width = img.shape[:2]

    # Whether to perform the split/merge action
    do_split = img_height > tile_size or img_width > tile_size

    if do_split:
        rlts = ops.esrgan_launcher_split_merge(img, process_matgen, MATGEN_MODELS, scale_factor=1, tile_size=tile_size)
    else:
        rlts = [process_matgen(img, model) for model in MATGEN_MODELS]

    if seamless or mirror or replicate:
        rlts = [ops.crop_seamless(rlt) for rlt in rlts]

    normal_map = rlts[0]
    roughness = rlts[1][:, :, 1]
    displacement = rlts[1][:, :, 0]

    normal_name = '{:s}_Normal.png'.format(basename)
    cv2.imwrite(os.path.join(output_folder, normal_name), normal_map)

    rough_name = '{:s}_Roughness.png'.format(basename)
    rough_img = roughness
    cv2.imwrite(os.path.join(output_folder, rough_name), rough_img)

    displace_name = '{:s}_Displacement.png'.format(basename)
    cv2.imwrite(os.path.join(output_folder, displace_name), displacement)

    metallic_file = f"{output_folder}/{basename}_Metallic{file_ext}"
    generate_metallic_map(albedo_filepath, metallic_file)

    displace_file = f"{output_folder}/{basename}_Displacement{file_ext}"
    ao_file = f"{output_folder}/{basename}_AO{file_ext}"
    generate_ao_map(displace_file, ao_file)

    roughness_file = f"{output_folder}/{basename}_Roughness{file_ext}"
    orm_file = f"{output_folder}/{basename}_ORM{file_ext}"

    generate_orm_map(ao_file, roughness_file, metallic_file, orm_file)

    emissive_file = f"{output_folder}/{basename}_Emissive{file_ext}"

    generate_emissive_map(albedo_filepath, emissive_file)
    

def generate_metallic_map(albedo_path: str,
                          output_path: str = "metallic.png",
                          grayness_thresh: float = 0.08,
                          saturation_thresh: float = 0.25,
                          smooth: bool = True) -> Image.Image:
    """
    Generate a metallic map from an albedo texture using color heuristics.

    Parameters:
        albedo_path (str): Path to the albedo (base color) image.
        output_path (str): Path to save the resulting metallic map.
        grayness_thresh (float): Max color std deviation for 'gray' detection (0–1).
        saturation_thresh (float): Max saturation value for metal classification (0–1).
        smooth (bool): If True, apply Gaussian-like blur for smoother mask.

    Returns:
        PIL.Image.Image: The generated metallic map as a grayscale image.
    """
    # Load image
    albedo = Image.open(albedo_path).convert("RGB")
    arr = np.array(albedo, dtype=np.float32) / 255.0

    # Compute color statistics
    mean_rgb = np.mean(arr, axis=2)
    std_rgb = np.std(arr, axis=2)
    max_rgb = np.max(arr, axis=2)
    min_rgb = np.min(arr, axis=2)

    # Compute saturation (approx from RGB)
    saturation = (max_rgb - min_rgb) / (max_rgb + 1e-6)
    grayness = std_rgb  # measure of how neutral the color is

    # Core metallic heuristic
    metallic_mask = np.where(
        (grayness < grayness_thresh) & (saturation < saturation_thresh) & (mean_rgb > 0.1),
        1.0,
        0.0
    )

    # Optional soft blending
    if smooth:
        from scipy.ndimage import gaussian_filter
        metallic_mask = gaussian_filter(metallic_mask, sigma=1.2)
        metallic_mask = np.clip(metallic_mask, 0, 1)

    # Save output
    metallic_img = Image.fromarray((metallic_mask * 255).astype(np.uint8))
    metallic_img.save(output_path)
    return metallic_img


def generate_ao_map(displacement_path: str, output_path: str = None, blur_radius: int = 2):
    """
    Generates a rough Ambient Occlusion (AO) map from a displacement (height) map.

    Args:
        displacement_path (str): Path to the displacement (height) map PNG.
        output_path (str, optional): Path to save the AO map.
                                     If None, saves alongside input with '_AO' suffix.
        blur_radius (int): Radius for Gaussian blur to simulate soft occlusion.

    Returns:
        str: Path of the saved AO map.
    """
    # Load displacement map as grayscale
    disp_img = Image.open(displacement_path).convert("L")
    disp_arr = np.array(disp_img, dtype=np.float32) / 255.0  # normalize 0-1

    # Invert height: high areas = less occluded
    inv_height = 1.0 - disp_arr

    # Convert back to image
    inv_height_img = Image.fromarray((inv_height * 255).astype(np.uint8))

    # Apply Gaussian blur to simulate ambient occlusion
    ao_img = inv_height_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))

    # Determine output path
    if output_path is None:
        base, ext = os.path.splitext(displacement_path)
        output_path = f"{base}_AO.png"

    # Save AO map
    ao_img.save(output_path)
    print(f"AO map saved to: {output_path}")
    return output_path


def generate_orm_map(ao_path=None, roughness_path=None, metallic_path=None, output_path=None):
    """
    Combines AO, Roughness, and Metallic textures into a single ORM map (R=AO, G=Roughness, B=Metallic).

    Args:
        ao_path (str): Path to AO texture (grayscale)
        roughness_path (str): Path to Roughness texture (grayscale)
        metallic_path (str): Path to Metallic texture (grayscale)
        output_path (str): Path to save the ORM texture. If None, saves as 'ORM.png' in the first input folder.

    Returns:
        str: Path of the saved ORM map.
    """

    # Load textures or create default gray if missing
    def load_gray(path, size=None):
        if path and os.path.exists(path):
            img = Image.open(path).convert("L")
        else:
            img = Image.new("L", size if size else (1024, 1024), 255)
        return img

    # Determine base size from first available texture
    base_size = None
    for path in [ao_path, roughness_path, metallic_path]:
        if path and os.path.exists(path):
            base_size = Image.open(path).size
            break
    if base_size is None:
        base_size = (1024, 1024)  # default

    ao_img = load_gray(ao_path, base_size)
    rough_img = load_gray(roughness_path, base_size)
    metal_img = load_gray(metallic_path, base_size)

    # Merge into RGB
    orm_img = Image.merge("RGB", (ao_img, rough_img, metal_img))

    # Determine output path
    if not output_path:
        first_path = next((p for p in [ao_path, roughness_path, metallic_path] if p), None)
        folder = os.path.dirname(first_path) if first_path else os.getcwd()
        output_path = os.path.join(folder, "ORM.png")

    orm_img.save(output_path)
    print(f"ORM map saved to: {output_path}")
    return output_path


def generate_emissive_map(input_path, output_path,
                          sat_threshold=0.4,
                          val_threshold=0.6,
                          white_exclude_threshold=0.25,
                          blur_radius=2,
                          excluded_hue_ranges=None):
    """
    Generate an emissive texture from an albedo image by detecting bright, colorful regions.
    Excludes white and light-blue (cyan/sky tones) by default.
    """
    if excluded_hue_ranges is None:
        excluded_hue_ranges = []
        
    img = Image.open(input_path).convert("RGB")
    img_np = np.array(img) / 255.0
    r, g, b = img_np[..., 0], img_np[..., 1], img_np[..., 2]

    cmax = np.max(img_np, axis=-1)
    cmin = np.min(img_np, axis=-1)
    delta = cmax - cmin

    hue = np.zeros_like(cmax)
    mask = delta != 0
    hue[mask & (cmax == r)] = (60 * ((g - b) / delta) % 360)[mask & (cmax == r)]
    hue[mask & (cmax == g)] = (60 * ((b - r) / delta) + 120)[mask & (cmax == g)]
    hue[mask & (cmax == b)] = (60 * ((r - g) / delta) + 240)[mask & (cmax == b)]

    sat = np.zeros_like(cmax)
    sat[cmax != 0] = delta[cmax != 0] / cmax[cmax != 0]
    val = cmax

    # Basic emissive mask based on saturation and brightness
    emissive_mask = (sat > sat_threshold) & (val > val_threshold) & (sat > white_exclude_threshold)

    # Exclude unwanted hue ranges (like light blues)
    for low, high in excluded_hue_ranges:
        emissive_mask &= ~((hue >= low) & (hue <= high))

    # Apply mask
    emissive_img = np.zeros_like(img_np)
    emissive_img[emissive_mask] = img_np[emissive_mask]

    emissive_img = (emissive_img * 255).astype(np.uint8)

    emissive_img = Image.fromarray(emissive_img)
    emissive_img = emissive_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    emissive_img.save(output_path)
    print(f"Emissive texture saved to: {output_path}")


def generate(input_image_path, output_dir):

    global MATGEN_MODELS

    os.makedirs(output_dir, exist_ok=True)

    # Load MatGen Models

    MATGEN_MODELS = [
        load_matgen_model(NORMAL_MAP_MODEL), # NORMAL MAP
        load_matgen_model(OTHER_MAP_MODEL)   # ROUGHNESS/DISPLACEMENT MAPS
    ]

    # Generate PBR from Albedo

    generate_pbr(input_image_path, output_dir, override=True)


if __name__ == "__main__":

    if len(sys.argv) < 1:
        sys.exit(1)

    import argparse

    parser = argparse.ArgumentParser(description="Run MaterialMapGenerator demo pipeline.")
    parser.add_argument("--image_path", type=str, default="input/example.png", help="input image path")
    parser.add_argument("--output_dir", type=str, default="output", help="output directory path")

    args = parser.parse_args()

    generate(args.image_path, args.output_dir)