Update app.py

app.py

@@ -4,40 +4,59 @@ import numpy as np
 from PIL import Image
 import noise  # Requires `pynoise` package: `pip install noise`
 
-# Advanced Normal Map generation
-def generate_normal_map(image, strength=1.0, blur_size=5, use_bilateral=False):
-    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-    # Noise reduction: Bilateral filter (edge-preserving) or Gaussian blur
+# 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:
-        img = cv2.bilateralFilter(img, 9, 75, 75)
+        gray = cv2.bilateralFilter(gray, 9, 75, 75)
     else:
-        img = cv2.GaussianBlur(img, (blur_size, blur_size), 0)
-    # Compute gradients with Scharr (better than Sobel for fine details)
-    sobel_x = cv2.Scharr(img, cv2.CV_64F, 1, 0)
-    sobel_y = cv2.Scharr(img, cv2.CV_64F, 0, 1)
-    # Normalize gradients with strength adjustment
-    sobel_x = cv2.normalize(sobel_x, None, -strength, strength, cv2.NORM_MINMAX)
-    sobel_y = cv2.normalize(sobel_y, None, -strength, strength, cv2.NORM_MINMAX)
-    # Compute normal vectors
-    normal_map = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.float32)
-    normal_map[..., 0] = sobel_x  # X (Red)
-    normal_map[..., 1] = sobel_y  # Y (Green)
-    normal_map[..., 2] = 1.0      # Z (Blue)
+        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 (e.g., red affects X, green affects Y)
+    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)
-    # Convert to RGB
     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 generation
-def generate_displacement_map(image, contrast=1.0, add_noise=False, noise_scale=0.1):
+# Advanced Displacement Map with adaptive thresholding and edge enhancement
+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)
-    # Enhance contrast with histogram equalization
-    img = cv2.equalizeHist(img)
+    # Adaptive histogram equalization for better contrast
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    img = clahe.apply(img)
     # Adjust contrast
     img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
-    # Optional: Add Perlin noise for realism
+    # Edge enhancement with Laplacian
+    laplacian = cv2.Laplacian(img, cv2.CV_64F)
+    img = cv2.addWeighted(img, 1.0, laplacian, edge_boost, 0)
+    img = np.clip(img, 0, 255).astype(np.uint8)
+    # Optional Perlin noise
     if add_noise:
         height, width = img.shape
         noise_map = np.zeros((height, width), dtype=np.float32)
@@ -47,28 +66,33 @@ def generate_displacement_map(image, contrast=1.0, add_noise=False, noise_scale=
         img = cv2.add(img, noise_map.astype(np.uint8))
     return Image.fromarray(img)
 
-# Advanced Roughness Map generation
-def generate_roughness_map(image, invert=True, sharpness=1.0, detail_boost=0.5):
+# Advanced Roughness Map with frequency separation and texture analysis
+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)
-    # Edge-preserving smoothing
-    img = cv2.bilateralFilter(img, 9, 75, 75)
+    # Frequency separation: Low-frequency (smooth) and high-frequency (detail)
+    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
-    # Sharpen the image
+    # Sharpening and detail boost
     blurred = cv2.GaussianBlur(img, (5, 5), 0)
     img = cv2.addWeighted(img, 1.0 + sharpness, blurred, -sharpness, 0)
-    # Boost fine details with unsharp masking
     img = cv2.addWeighted(img, 1.0 + detail_boost, blurred, -detail_boost, 0)
     return Image.fromarray(img)
 
-# Process function with all parameters
-def process_image(input_image, normal_strength=1.0, normal_blur=5, normal_bilateral=False,
-                  disp_contrast=1.0, disp_noise=False, disp_noise_scale=0.1,
-                  rough_invert=True, rough_sharpness=1.0, rough_detail=0.5):
-    normal_map = generate_normal_map(input_image, strength=normal_strength, blur_size=normal_blur, use_bilateral=normal_bilateral)
-    displacement_map = generate_displacement_map(input_image, contrast=disp_contrast, add_noise=disp_noise, noise_scale=disp_noise_scale)
-    roughness_map = generate_roughness_map(input_image, invert=rough_invert, sharpness=rough_sharpness, detail_boost=rough_detail)
+# Process function with advanced parameters
+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 with advanced controls
@@ -78,16 +102,19 @@ interface = gr.Interface(
         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 (if not bilateral)"),
+        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 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"),
@@ -95,7 +122,7 @@ interface = gr.Interface(
         gr.Image(type="pil", label="Roughness Map")
     ],
     title="Advanced Material Map Generator",
-    description="Upload an image to generate highly customizable Normal, Displacement, and Roughness maps."
+    description="Upload an image to generate highly accurate Normal, Displacement, and Roughness maps with advanced controls."
 )
 
 interface.launch()