initial implementation

Dockerfile

@@ -0,0 +1,25 @@
+# Choose a specific Python version. Using a slim image is good for size.
+# You can change the version to match your needs, e.g., python:3.10-slim
+FROM python:3.12-slim
+
+# Set the working directory inside the container
+WORKDIR /code
+
+# Copy the file with your Python package dependencies
+COPY requirements.txt .
+
+# Install system-level dependencies if you need them.
+# For example, if you need ffmpeg for audio processing:
+# RUN apt-get update && apt-get install -y ffmpeg
+
+# Install the Python dependencies from your requirements.txt
+# --no-cache-dir makes the image smaller
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy your application code into the container
+COPY app.py .
+
+# Command to run your Gradio app.
+# The --server-name 0.0.0.0 makes it accessible from outside the container.
+# The --server-port 7860 is the default Gradio port that Hugging Face expects.
+CMD ["python", "app.py", "--server-name", "0.0.0.0", "--server-port", "7860"]

app.py

@@ -0,0 +1,153 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForDepthEstimation
+import tempfile
+from stl import mesh
+from sklearn.decomposition import PCA
+
+MAX_RESOLUTION = 1.5e6
+MODEL_ID = "depth-anything/Depth-Anything-V2-Large-hf"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID).to(device)
+print("Model loaded successfully.")
+
+def create_3d_model(
+    input_image: np.ndarray,
+    texture_strength: float,
+    max_z: float,
+    min_z: float,
+    x_length: float,
+    do_pca_correction: bool,
+    depth_map_smoothing: int,
+    texture_smoothing: int
+):
+    if input_image is None: raise gr.Error("Please upload an image.")
+    if max_z <= min_z: raise gr.Error("Max Z-height must be greater than Min Z-height.")
+
+    # resize large images
+    h, w, _ = input_image.shape
+    if h * w > MAX_RESOLUTION:
+        gr.Info("Image is large, downsampling to improve performance...")
+        ratio = (MAX_RESOLUTION / (h * w)) ** 0.5
+        new_w, new_h = int(w * ratio), int(h * ratio)
+        input_image = cv2.resize(input_image, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
+    image = Image.fromarray(input_image).convert("RGB")
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        outputs = model(**inputs)
+        predicted_depth = outputs.predicted_depth
+        depth = torch.nn.functional.interpolate(
+            predicted_depth.unsqueeze(1), size=(image.height, image.width),
+            mode="bicubic", align_corners=False,
+        ).squeeze().cpu().numpy()
+    depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())
+
+    # smoothing base depthmap
+    if depth_map_smoothing > 1:
+        ksize = int(depth_map_smoothing)
+        if ksize % 2 == 0: ksize += 1
+        depth_normalized = cv2.GaussianBlur(depth_normalized, (ksize, ksize), 0)
+
+    gray_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY)
+    brightness_normalized = gray_image.astype(float) / 255.0
+    
+    # smoothing brightness
+    if texture_smoothing > 1:
+        ksize = int(texture_smoothing)
+        if ksize % 2 == 0: ksize += 1
+        brightness_normalized = cv2.GaussianBlur(brightness_normalized, (ksize, ksize), 0)
+
+    if brightness_normalized.shape != depth_normalized.shape:
+        brightness_normalized = cv2.resize(
+            brightness_normalized, (depth_normalized.shape[1], depth_normalized.shape[0]),
+            interpolation=cv2.INTER_LINEAR
+        )
+
+    combined_map = depth_normalized + (brightness_normalized * texture_strength)
+    c_min, c_max = combined_map.min(), combined_map.max()
+    if c_max > c_min:
+        combined_map_rescaled = (combined_map - c_min) / (c_max - c_min)
+    else:
+        combined_map_rescaled = np.zeros_like(combined_map)
+
+    z_data = min_z + combined_map_rescaled * (max_z - min_z)
+
+    # Planar correction with PCA
+    if do_pca_correction:
+        height, width = z_data.shape
+        y_length = x_length * (height / width)
+        x_coords_1d, y_coords_1d = np.linspace(0, x_length, width), np.linspace(y_length, 0, height)
+        x_grid, y_grid = np.meshgrid(x_coords_1d, y_coords_1d)
+        points = np.stack([x_grid.flatten(), y_grid.flatten(), z_data.flatten()], axis=1)
+        n_points, n_samples = points.shape[0], min(points.shape[0], 50000)
+        sample_indices = np.random.choice(n_points, n_samples, replace=False)
+        pca = PCA(n_components=3)
+        pca.fit(points[sample_indices])
+        normal = pca.components_[2]
+        if normal[2] < 0:
+            normal *= -1
+        p0 = pca.mean_
+        z_plane = p0[2] - (normal[0] * (x_grid - p0[0]) + normal[1] * (y_grid - p0[1])) / normal[2]
+        corrected_z = z_data - z_plane
+        cz_min, cz_max = corrected_z.min(), corrected_z.max()
+        if cz_max > cz_min:
+            z_data = min_z + (corrected_z - cz_min) / (cz_max - cz_min) * (max_z - min_z)
+
+    # STL mesh
+    height, width = z_data.shape
+    y_length = x_length * (height / width)
+    vertices = np.zeros((height, width, 3))
+    x_coords, y_coords = np.linspace(0, x_length, width), np.linspace(y_length, 0, height)
+    vertices[:, :, 0], vertices[:, :, 1], vertices[:, :, 2] = x_coords[np.newaxis, :], y_coords[:, np.newaxis], z_data
+    faces = []
+    for i in range(height-1):
+        for j in range(width-1):
+            v1,v2,v3,v4 = vertices[i,j], vertices[i+1,j], vertices[i+1,j+1], vertices[i,j+1]
+            faces.extend([[v1, v2, v3], [v1, v3, v4]])
+    v_tl, v_tr, v_bl, v_br = vertices[0,0], vertices[0,width-1], vertices[height-1,0], vertices[height-1,width-1]
+    b_tl,b_tr,b_bl,b_br = np.array([v_tl[0],v_tl[1],0]), np.array([v_tr[0],v_tr[1],0]), np.array([v_bl[0],v_bl[1],0]), np.array([v_br[0],v_br[1],0])
+    faces.extend([[v_tl,b_tl,b_tr],[v_tl,b_tr,v_tr], [v_br,b_br,b_bl],[v_br,b_bl,v_bl], [v_bl,b_bl,b_tl],[v_bl,b_tl,v_tl], [v_tr,b_tr,b_br],[v_tr,b_br,v_br], [b_tl,b_br,b_bl],[b_tl,b_tr,b_br]])
+    surface = mesh.Mesh(np.zeros(len(faces), dtype=mesh.Mesh.dtype))
+    surface.vectors = np.array(faces)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".stl") as tmpfile:
+        surface.save(tmpfile.name)
+        return tmpfile.name, tmpfile.name
+
+with gr.Blocks(theme='base') as demo:
+    gr.Markdown("# Image to 3D Relief Generator")
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="numpy", label="Upload Image")
+            gr.Markdown("### Model Parameters")
+            texture_strength = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.001, label="Brightness Texture Strength")
+            depth_map_smoothing = gr.Slider(minimum=0, maximum=51, value=0, step=1, label="Depth Map Smoothing", info="Smooths the base geometry. 0 = none.")
+            texture_smoothing = gr.Slider(minimum=0, maximum=51, value=0, step=1, label="Texture Smoothing", info="Smooths the brightness texture. 0 = none.")
+            
+            x_length = gr.Number(value=100, label="X Length (units)")
+            min_z = gr.Number(value=0.5, label="Min Z-Height (units)")
+            max_z = gr.Number(value=5.0, label="Max Z-Height (units)")
+            do_pca = gr.Checkbox(value=True, label="Enable PCA Planar Correction")
+            
+            generate_btn = gr.Button("Generate STL", variant="primary")
+
+        with gr.Column(scale=1):
+            gr.Markdown("### 3D Model Output")
+            output_model = gr.Model3D(label="Generated 3D Model")
+            output_file = gr.File(label="Download STL File")
+
+    generate_btn.click(
+        fn=create_3d_model,
+        inputs=[
+            input_image, texture_strength, max_z, min_z, x_length, do_pca,
+            depth_map_smoothing, texture_smoothing
+        ],
+        outputs=[output_model, output_file]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+transformers
+torch
+torchvision
+accelerate
+opencv-python-headless
+numpy-stl
+Pillow
+scikit-learn