upload

app.py

@@ -0,0 +1,498 @@
+
+import cv2
+import torch
+from PIL import Image, ImageOps
+import numpy as np
+import gradio as gr
+import math
+import os
+import zipfile
+import trimesh
+import pygltflib
+from scipy.ndimage import median_filter
+import requests # Import requests for downloading
+
+
+# Depth-Anything V2 model setup
+from depth_anything_v2.dpt import DepthAnythingV2
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
+}
+
+encoder = 'vitl' # or 'vits', 'vitb', 'vitg'
+
+model = DepthAnythingV2(**model_configs[encoder])
+
+# Define model directory and path
+MODEL_DIR = "models"
+os.makedirs(MODEL_DIR, exist_ok=True)
+model_filename = f'depth_anything_v2_{encoder}.pth'
+model_path = os.path.join(MODEL_DIR, model_filename)
+
+# Add code to download model weights if not exists
+if not os.path.exists(model_path):
+    print(f"Downloading {model_path}...")
+    url = f"https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/{model_filename}"
+    response = requests.get(url, stream=True)
+    with open(model_path, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    print("Download complete.")
+
+model.load_state_dict(torch.load(model_path, map_location='cpu'))
+model = model.to(DEVICE).eval()
+
+# Helper functions (from your notebook)
+def quaternion_multiply(q1, q2):
+    x1, y1, z1, w1 = q1
+    x2, y2, z2, w2 = q2
+    return [
+        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+    ]
+
+
+def glb_add_lights(path_input, path_output):
+    """
+    Adds directional lights in the horizontal plane to the glb file.
+    :param path_input: path to input glb
+    :param path_output: path to output glb
+    :return: None
+    """
+    glb = pygltflib.GLTF2().load(path_input)
+
+    N = 3  # default max num lights in Babylon.js is 4
+    angle_step = 2 * math.pi / N
+    elevation_angle = math.radians(75)
+
+    light_colors = [
+        [1.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0],
+        [0.0, 0.0, 1.0],
+    ]
+
+    lights_extension = {
+        "lights": [
+            {"type": "directional", "color": light_colors[i], "intensity": 2.0}
+            for i in range(N)
+        ]
+    }
+
+    if "KHR_lights_punctual" not in glb.extensionsUsed:
+        glb.extensionsUsed.append("KHR_lights_punctual")
+    glb.extensions["KHR_lights_punctual"] = lights_extension
+
+    light_nodes = []
+    for i in range(N):
+        angle = i * angle_step
+
+        pos_rot = [0.0, 0.0, math.sin(angle / 2), math.cos(angle / 2)]
+        elev_rot = [
+            math.sin(elevation_angle / 2),
+            0.0,
+            0.0,
+            math.cos(elevation_angle / 2),
+        ]
+        rotation = quaternion_multiply(pos_rot, elev_rot)
+
+        node = {
+            "rotation": rotation,
+            "extensions": {"KHR_lights_punctual": {"light": i}},
+        }
+        light_nodes.append(node)
+
+    light_node_indices = list(range(len(glb.nodes), len(glb.nodes) + N))
+    glb.nodes.extend(light_nodes)
+
+    root_node_index = glb.scenes[glb.scene].nodes[0]
+    root_node = glb.nodes[root_node_index]
+    if hasattr(root_node, "children"):
+        root_node.children.extend(light_node_indices)
+    else:
+        root_node.children = light_node_indices
+
+    glb.save(path_output)
+
+
+def extrude_depth_3d(
+    path_rgb,
+    path_depth,
+    path_out_base="../",
+    alpha=1.0,
+    invert=0,
+    output_model_scale=100,
+    filter_size=3,
+    coef_near=0.0,
+    coef_far=1.0,
+    emboss=0.3,
+    f_thic=0.05,
+    f_near=-0.15,
+    f_back=0.01,
+    vertex_colors=True,
+    scene_lights=True,
+    prepare_for_3d_printing=False,
+    zip_outputs=False,
+):
+    f_far_inner = -emboss
+    f_far_outer = f_far_inner - f_back
+
+    f_near = max(f_near, f_far_inner)
+
+    depth_image = Image.open(path_depth)
+    mono_image = Image.open(path_rgb).convert("L")
+
+    if invert==1:
+        mono_image = ImageOps.invert(mono_image)
+
+    w, h = depth_image.size
+    d_max = max(w, h)
+    depth_image = np.array(depth_image).astype(np.double)
+    mono_image = np.array(mono_image).astype(np.double)
+    z_min, z_max = np.min(depth_image), np.max(depth_image)
+    m_min, m_max = np.min(mono_image), np.max(mono_image)
+    depth_image = (depth_image.astype(np.double) - z_min) / (z_max - z_min)
+    depth_image[depth_image < coef_near] = coef_near
+    depth_image[depth_image > coef_far] = coef_far
+    z_min, z_max = np.min(depth_image), np.max(depth_image)
+    depth_image = (depth_image - z_min) / (z_max - z_min)
+    mono_image = median_filter(mono_image, size=5)
+    mono_image = (mono_image.astype(np.double) - m_min) / (m_max - m_min)
+    mono_image_new = np.where(depth_image == coef_far, 1, mono_image)
+    m_min=np.min(mono_image_new)
+    mono_image_new = np.where(depth_image == coef_far, 0, mono_image)
+    m_max=np.max(mono_image_new)
+    mono_image = np.where(depth_image == coef_far, m_min, mono_image)
+    mono_image = (mono_image - m_min) / (m_max - m_min)
+    depth_image = np.where(depth_image != 1.0, (1-alpha) * depth_image + alpha * mono_image, depth_image)
+    #depth_image_new[depth_image < coef_near] = 0
+    #depth_image_new[depth_image > coef_far] = 1
+    #depth_image_new[depth_image_new < 0] = 0
+    depth_image = median_filter(depth_image, size=filter_size)
+    depth_image = emboss*(depth_image - np.min(depth_image)) / (np.max(depth_image) - np.min(depth_image))
+    Image.fromarray((depth_image * 255).astype(np.uint8)).convert("L").save(path_out_base+".png")
+    rgb_image = np.array(
+        Image.open(path_rgb).convert("RGB").resize((w, h), Image.Resampling.LANCZOS)
+    )
+
+    w_norm = w / float(d_max - 1)
+    h_norm = h / float(d_max - 1)
+    w_half = w_norm / 2
+    h_half = h_norm / 2
+
+    x, y = np.meshgrid(np.arange(w), np.arange(h))
+    x = x / float(d_max - 1) - w_half  # [-w_half, w_half]
+    y = -y / float(d_max - 1) + h_half  # [-h_half, h_half]
+    z = -depth_image  # -depth_emboss (far) - 0 (near)
+    vertices_2d = np.stack((x, y, z), axis=-1)
+    vertices = vertices_2d.reshape(-1, 3)
+    colors = rgb_image[:, :, :3].reshape(-1, 3) / 255.0
+
+    faces = []
+    for y in range(h - 1):
+        for x in range(w - 1):
+            idx = y * w + x
+            faces.append([idx, idx + w, idx + 1])
+            faces.append([idx + 1, idx + w, idx + 1 + w])
+
+    # OUTER frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],  # 00
+            [-w_half - f_thic, -h_half - f_thic, f_far_outer],  # 01
+            [w_half + f_thic, -h_half - f_thic, f_near],  # 02
+            [w_half + f_thic, -h_half - f_thic, f_far_outer],  # 03
+            [w_half + f_thic, h_half + f_thic, f_near],  # 04
+            [w_half + f_thic, h_half + f_thic, f_far_outer],  # 05
+            [-w_half - f_thic, h_half + f_thic, f_near],  # 06
+            [-w_half - f_thic, h_half + f_thic, f_far_outer],  # 07
+        ],
+        axis=0,
+    )
+    faces.extend(
+        [
+            [nv + 0, nv + 1, nv + 2],
+            [nv + 2, nv + 1, nv + 3],
+            [nv + 2, nv + 3, nv + 4],
+            [nv + 4, nv + 3, nv + 5],
+            [nv + 4, nv + 5, nv + 6],
+            [nv + 6, nv + 5, nv + 7],
+            [nv + 6, nv + 7, nv + 0],
+            [nv + 0, nv + 7, nv + 1],
+        ]
+    )
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * 8, axis=0)
+
+    # INNER frame
+
+    nv = len(vertices)
+    vertices_left_data = vertices_2d[:, 0]  # H x 3
+    vertices_left_frame = vertices_2d[:, 0].copy()  # H x 3
+    vertices_left_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_left_data, axis=0)
+    vertices = np.append(vertices, vertices_left_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
+    for i in range(h - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + h
+        faces.append([nvi_d, nvi_f, nvi_d + 1])
+        faces.append([nvi_d + 1, nvi_f, nvi_f + 1])
+
+    nv = len(vertices)
+    vertices_right_data = vertices_2d[:, -1]  # H x 3
+    vertices_right_frame = vertices_2d[:, -1].copy()  # H x 3
+    vertices_right_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_right_data, axis=0)
+    vertices = np.append(vertices, vertices_right_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
+    for i in range(h - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + h
+        faces.append([nvi_d, nvi_d + 1, nvi_f])
+        faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
+
+    nv = len(vertices)
+    vertices_top_data = vertices_2d[0, :]  # H x 3
+    vertices_top_frame = vertices_2d[0, :].copy()  # H x 3
+    vertices_top_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_top_data, axis=0)
+    vertices = np.append(vertices, vertices_top_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
+    for i in range(w - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + w
+        faces.append([nvi_d, nvi_d + 1, nvi_f])
+        faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
+
+    nv = len(vertices)
+    vertices_bottom_data = vertices_2d[-1, :]  # H x 3
+    vertices_bottom_frame = vertices_2d[-1, :].copy()  # H x 3
+    vertices_bottom_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_bottom_data, axis=0)
+    vertices = np.append(vertices, vertices_bottom_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
+    for i in range(w - 1):
+        faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
+    faces.append([nv, nv + 1, nv + w + 1])
+
+    # FRONT frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],
+            [-w_half - f_thic, h_half + f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_left_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
+    for i in range(h - 1):
+        faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
+    faces.append([nv, nv + 2, nv + 1])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [w_half + f_thic, h_half + f_thic, f_near],
+            [w_half + f_thic, -h_half - f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_right_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
+    for i in range(h - 1):
+        faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
+    faces.append([nv, nv + h + 1, nv + 1])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [w_half + f_thic, h_half + f_thic, f_near],
+            [-w_half - f_thic, h_half + f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_top_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
+    for i in range(w - 1):
+        faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
+    faces.append([nv, nv + 1, nv + 2])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],
+            [w_half + f_thic, -h_half - f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_bottom_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
+    for i in range(w - 1):
+        faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
+    faces.append([nv, nv + 1, nv + w + 1])
+
+    # BACK frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_far_outer],  # 00
+            [w_half + f_thic, -h_half - f_thic, f_far_outer],  # 01
+            [w_half + f_thic, h_half + f_thic, f_far_outer],  # 02
+            [-w_half - f_thic, h_half + f_thic, f_far_outer],  # 03
+        ],
+        axis=0,
+    )
+    faces.extend(
+        [
+            [nv + 0, nv + 2, nv + 1],
+            [nv + 2, nv + 0, nv + 3],
+        ]
+    )
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * 4, axis=0)
+
+
+    trimesh_kwargs = {}
+    if vertex_colors:
+        trimesh_kwargs["vertex_colors"] = colors
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, **trimesh_kwargs)
+
+    mesh.merge_vertices()
+
+    current_max_dimension = max(mesh.extents)
+    scaling_factor = output_model_scale / current_max_dimension
+    mesh.apply_scale(scaling_factor)
+
+    if prepare_for_3d_printing:
+        rotation_mat = trimesh.transformations.rotation_matrix(
+            np.radians(0), [0.5, 0, 0]
+        )
+        mesh.apply_transform(rotation_mat)
+
+    if path_out_base is None:
+        path_out_base = os.path.splitext(path_depth)[0].replace("_16bit", "")
+    path_out_glb = path_out_base + ".glb"
+    path_out_stl = path_out_base + ".stl"
+    path_out_obj = path_out_base + ".obj"
+
+    mesh.export(path_out_stl, file_type="stl")
+    """
+    mesh.export(path_out_glb, file_type="glb")
+    if scene_lights:
+        glb_add_lights(path_out_glb, path_out_glb)
+    mesh.export(path_out_obj, file_type="obj")
+
+    if zip_outputs:
+        with zipfile.ZipFile(path_out_glb + ".zip", "w", zipfile.ZIP_DEFLATED) as zipf:
+            arcname = os.path.basename(os.path.splitext(path_out_glb)[0]) + ".glb"
+            zipf.write(path_out_glb, arcname=arcname)
+            path_out_glb = path_out_glb + ".zip"
+        with zipfile.ZipFile(path_out_stl + ".zip", "w", zipfile.ZIP_DEFLATED) as zipf:
+            arcname = os.path.basename(os.path.splitext(path_out_stl)[0]) + ".stl"
+            zipf.write(path_out_stl, arcname=arcname)
+            path_out_stl = path_out_stl + ".zip"
+        with zipfile.ZipFile(path_out_obj + ".zip", "w", zipfile.ZIP_DEFLATED) as zipf:
+            arcname = os.path.basename(os.path.splitext(path_out_obj)[0]) + ".obj"
+            zipf.write(path_out_obj, arcname=arcname)
+            path_out_obj = path_out_obj + ".zip"
+    """
+    return path_out_glb, path_out_stl, path_out_obj
+
+def scale_to_width(img, length):
+      if img.width < img.height:
+        width = length
+        height = round(img.height * length / img.width)
+      else:
+        width = round(img.width * length / img.height)
+        height = length
+      return (width,height)
+
+
+# Gradio Interface function
+def process_image_and_generate_stl(image_input, depth_near, depth_far, thickness, alpha):
+    # Depth Estimation
+    raw_img = cv2.imread(image_input)
+    depth = model.infer_image(raw_img) # HxW raw depth map in numpy
+
+    # Save depth map temporarily
+    depth_output_path = "output_depth.png"
+    cv2.imwrite(depth_output_path, depth)
+
+    # Prepare images for 3D model generation
+    img_rgb = image_input
+    img_depth = depth_output_path
+    inv = 0 # Assuming no inversion for now, based on previous code
+    # Image.open(img_rgb).convert("L").save("example_1_black.png") # This line might not be necessary for the final output
+    size = scale_to_width(Image.open(img_rgb), 512)
+    Image.open(img_rgb).resize(size, Image.Resampling.LANCZOS).save("one.png") # Use Resampling.LANCZOS
+    if inv == 1:
+        Image.open(img_depth).convert(mode="F").resize(size, Image.Resampling.BILINEAR).convert("I").save("two.png") # Use Resampling.BILINEAR
+    else:
+        img=Image.open(img_depth).convert(mode="F").resize(size, Image.Resampling.BILINEAR).convert("I") # Use Resampling.BILINEAR
+        img = np.array(img).astype(np.double)
+        im_max=np.max(img)
+        im_min=np.min(img)
+        img=(1-(img-im_min)/(im_max-im_min))*im_max
+        img=Image.fromarray(img)
+        img.convert("I").save("two.png")
+
+
+    # 3D Model Generation
+    output_path_base = "generated_relief"
+    glb_path, stl_path, obj_path = extrude_depth_3d(
+        "one.png",
+        "two.png",
+        alpha=alpha,
+        invert=inv,
+        path_out_base=output_path_base,
+        output_model_scale=100,
+        filter_size=5, # Using 5 based on previous code
+        coef_near=depth_near,
+        coef_far=depth_far,
+        emboss=thickness,
+        f_thic=0.0, # Using 0.0 based on previous code
+        f_near=-thickness, # Using -thickness based on previous code
+        f_back=0.01, # Using 0.01 based on previous code
+        vertex_colors=True,
+        scene_lights=True,
+        prepare_for_3d_printing=True,
+    )
+
+    return stl_path # Return the path to the generated STL file
+
+
+# Gradio Interface definition
+iface = gr.Interface(
+    fn=process_image_and_generate_stl,
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Slider(minimum=0, maximum=1.0, value=0, label="Depth Near"),
+        gr.Slider(minimum=0, maximum=1.0, value=1.0, label="Depth Far"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.3, label="Thickness"),
+        gr.Slider(minimum=0, maximum=1.0, value=0.05, label="Alpha"),
+    ],
+    outputs=gr.File(label="Download STL File"), # Use gr.File() for file downloads
+    title="Image to 2.5D Relief Model Generator",
+    description="Upload an image, set parameters, and generate a 2.5D relief model (.stl file)."
+)
+
+# Launch the interface (for local testing)
+if __name__ == "__main__":
+    iface.launch(debug=True)