app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image
import noise  # Requires `pynoise` package: `pip install noise`

# Advanced Normal Map with multi-scale gradients and color influence
def generate_normal_map(image, strength=1.0, blur_size=5, use_bilateral=False, color_influence=0.3):
    img = np.array(image)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    # Noise reduction
    if use_bilateral:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
    else:
        gray = cv2.GaussianBlur(gray, (blur_size, blur_size), 0)
    
    # Multi-scale gradient computation using Laplacian pyramid
    levels = 3
    normal_map = np.zeros((gray.shape[0], gray.shape[1], 3), dtype=np.float32)
    for i in range(levels):
        scale = 1 / (2 ** i)
        resized = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
        sobel_x = cv2.Scharr(resized, cv2.CV_64F, 1, 0)
        sobel_y = cv2.Scharr(resized, cv2.CV_64F, 0, 1)
        sobel_x = cv2.resize(sobel_x, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
        sobel_y = cv2.resize(sobel_y, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
        normal_map[..., 0] += sobel_x * (1.0 / levels)
        normal_map[..., 1] += sobel_y * (1.0 / levels)
    
    # Normalize with strength
    normal_map[..., 0] = cv2.normalize(normal_map[..., 0], None, -strength, strength, cv2.NORM_MINMAX)
    normal_map[..., 1] = cv2.normalize(normal_map[..., 1], None, -strength, strength, cv2.NORM_MINMAX)
    normal_map[..., 2] = 1.0
    
    # Add color influence
    color_factor = color_influence * strength
    normal_map[..., 0] += (img[..., 0] / 255.0 - 0.5) * color_factor
    normal_map[..., 1] += (img[..., 1] / 255.0 - 0.5) * color_factor
    
    # Normalize to unit vectors
    norm = np.linalg.norm(normal_map, axis=2, keepdims=True)
    normal_map = np.divide(normal_map, norm, out=np.zeros_like(normal_map), where=norm != 0)
    normal_map = (normal_map + 1) * 127.5
    normal_map = np.clip(normal_map, 0, 255).astype(np.uint8)
    return Image.fromarray(normal_map)

# Advanced Displacement Map with fixed type handling
def generate_displacement_map(image, contrast=1.0, add_noise=False, noise_scale=0.1, edge_boost=1.0):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    # Adaptive histogram equalization
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img = clahe.apply(img)
    # Adjust contrast
    img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
    # Edge enhancement with Laplacian
    laplacian = cv2.Laplacian(img, cv2.CV_64F)
    # Convert Laplacian to uint8 after scaling and clipping
    laplacian = cv2.convertScaleAbs(laplacian, alpha=edge_boost, beta=0)
    # Ensure both arrays are uint8 before blending
    img = cv2.addWeighted(img, 1.0, laplacian, 0.5 * edge_boost, 0)  # Reduced weight for stability
    # Optional Perlin noise
    if add_noise:
        height, width = img.shape
        noise_map = np.zeros((height, width), dtype=np.float32)
        for y in range(height):
            for x in range(width):
                noise_map[y, x] = noise.pnoise2(x / 50.0, y / 50.0, octaves=6) * noise_scale * 255
        img = cv2.add(img, noise_map.astype(np.uint8))
    return Image.fromarray(img)

# Advanced Roughness Map with frequency separation
def generate_roughness_map(image, invert=True, sharpness=1.0, detail_boost=0.5, frequency_weight=0.5):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    # Frequency separation
    low_freq = cv2.bilateralFilter(img, 9, 75, 75)
    high_freq = cv2.subtract(img, low_freq)
    # Combine with weight
    img = cv2.addWeighted(low_freq, 1.0 - frequency_weight, high_freq, frequency_weight, 0)
    # Optional inversion
    if invert:
        img = 255 - img
    # Sharpening and detail boost
    blurred = cv2.GaussianBlur(img, (5, 5), 0)
    img = cv2.addWeighted(img, 1.0 + sharpness, blurred, -sharpness, 0)
    img = cv2.addWeighted(img, 1.0 + detail_boost, blurred, -detail_boost, 0)
    return Image.fromarray(img)

# Process function
def process_image(input_image, normal_strength=1.0, normal_blur=5, normal_bilateral=False, normal_color=0.3,
                  disp_contrast=1.0, disp_noise=False, disp_noise_scale=0.1, disp_edge=1.0,
                  rough_invert=True, rough_sharpness=1.0, rough_detail=0.5, rough_freq=0.5):
    normal_map = generate_normal_map(input_image, strength=normal_strength, blur_size=normal_blur, 
                                     use_bilateral=normal_bilateral, color_influence=normal_color)
    displacement_map = generate_displacement_map(input_image, contrast=disp_contrast, add_noise=disp_noise, 
                                                 noise_scale=disp_noise_scale, edge_boost=disp_edge)
    roughness_map = generate_roughness_map(input_image, invert=rough_invert, sharpness=rough_sharpness, 
                                           detail_boost=rough_detail, frequency_weight=rough_freq)
    return normal_map, displacement_map, roughness_map

# Gradio Interface
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        # Normal Map Controls
        gr.Slider(minimum=0.1, maximum=5.0, step=0.1, value=1.0, label="Normal Map Strength"),
        gr.Slider(minimum=3, maximum=15, step=2, value=5, label="Normal Map Blur Size (odd numbers)"),
        gr.Checkbox(value=False, label="Use Bilateral Filter for Normal Map"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.3, label="Normal Map Color Influence"),
        # Displacement Map Controls
        gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Displacement Map Contrast"),
        gr.Checkbox(value=False, label="Add Noise to Displacement Map"),
        gr.Slider(minimum=0.05, maximum=0.5, step=0.05, value=0.1, label="Displacement Noise Scale"),
        gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.0, label="Displacement Edge Boost"),
        # Roughness Map Controls
        gr.Checkbox(value=True, label="Invert Roughness Map"),
        gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.0, label="Roughness Sharpness"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.5, label="Roughness Detail Boost"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.5, label="Roughness Frequency Weight")
    ],
    outputs=[
        gr.Image(type="pil", label="Normal Map"),
        gr.Image(type="pil", label="Displacement Map"),
        gr.Image(type="pil", label="Roughness Map")
    ],
    title="Advanced Material Map Generator",
    description="Upload an image to generate highly accurate Normal, Displacement, and Roughness maps with advanced controls."
)

interface.launch()