Update app.py

app.py

@@ -1,77 +1,101 @@
-import torch
-from transformers import AutoModelForDepthEstimation, AutoImageProcessor
 import gradio as gr
+import cv2
 import numpy as np
 from PIL import Image
+import noise  # Requires `pynoise` package: `pip install noise`
 
-# Load the pretrained depth estimation model from Hugging Face
-model_name = "Intel/dpt-large"
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForDepthEstimation.from_pretrained(model_name)
+# 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
+    if use_bilateral:
+        img = cv2.bilateralFilter(img, 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)
+    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)
 
-def generate_material_maps(image):
-    """
-    Generate diffuse, height, and normal maps from an input image.
-    
-    Args:
-        image: PIL Image object (input image)
-    
-    Returns:
-        tuple: Diffuse map, height map, and normal map as numpy arrays
-    """
-    # Preprocess the image for the model
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Run depth estimation
-    with torch.no_grad():
-        outputs = model(**inputs)
-        depth_map = outputs.predicted_depth
-    
-    # Convert depth map to numpy and normalize to [0, 1]
-    depth_map = depth_map.squeeze().cpu().numpy()
-    height_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    
-    # Compute normal map from height map using gradients
-    gy, gx = np.gradient(height_map)
-    normal_map = np.dstack((gx, gy, np.ones_like(height_map)))  # RGB for x, y, z normals
-    normal_map = (normal_map + 1) / 2  # Normalize to [0, 1] for visualization
-    
-    # Use the original image as the diffuse map
-    diffuse_map = np.array(image) / 255.0  # Normalize to [0, 1]
-    
-    return diffuse_map, height_map, normal_map
+# Advanced Displacement Map generation
+def generate_displacement_map(image, contrast=1.0, add_noise=False, noise_scale=0.1):
+    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+    # Enhance contrast with histogram equalization
+    img = cv2.equalizeHist(img)
+    # Adjust contrast
+    img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
+    # Optional: Add Perlin noise for realism
+    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)
 
-def gradio_interface(image):
-    """
-    Interface function for Gradio to process the input and return images.
-    
-    Args:
-        image: PIL Image object
-    
-    Returns:
-        tuple: PIL Images for diffuse, height, and normal maps
-    """
-    diffuse, height, normal = generate_material_maps(image)
-    
-    # Convert to PIL Images and scale to [0, 255] for display
-    diffuse_img = Image.fromarray((diffuse * 255).astype(np.uint8))
-    height_img = Image.fromarray((height * 255).astype(np.uint8), mode="L")  # Grayscale
-    normal_img = Image.fromarray((normal * 255).astype(np.uint8))
-    
-    return diffuse_img, height_img, normal_img
+# Advanced Roughness Map generation
+def generate_roughness_map(image, invert=True, sharpness=1.0, detail_boost=0.5):
+    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+    # Edge-preserving smoothing
+    img = cv2.bilateralFilter(img, 9, 75, 75)
+    # Optional inversion
+    if invert:
+        img = 255 - img
+    # Sharpen the image
+    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)
 
-# Create Gradio interface
+# 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)
+    return normal_map, displacement_map, roughness_map
+
+# Gradio Interface with advanced controls
 interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
+    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 (if not bilateral)"),
+        gr.Checkbox(value=False, label="Use Bilateral Filter for Normal Map"),
+        # 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"),
+        # 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")
+    ],
     outputs=[
-        gr.Image(type="pil", label="Diffuse Map"),
-        gr.Image(type="pil", label="Height Map"),
-        gr.Image(type="pil", label="Normal Map")
+        gr.Image(type="pil", label="Normal Map"),
+        gr.Image(type="pil", label="Displacement Map"),
+        gr.Image(type="pil", label="Roughness Map")
     ],
-    title="Material Map Generator",
-    description="Upload an image to generate diffuse, height, and normal maps using a pretrained Hugging Face model."
+    title="Advanced Material Map Generator",
+    description="Upload an image to generate highly customizable Normal, Displacement, and Roughness maps."
 )
 
-# Launch the interface locally (for testing)
 interface.launch()