default to open body

app.py

@@ -13,7 +13,7 @@ pillow_heif.register_heif_opener()
 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)
+processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True)
 model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID).to(device)
 print("Model loaded successfully.")
 
@@ -25,7 +25,8 @@ def create_3d_model(
     x_length: float,
     do_pca_correction: bool,
     depth_map_smoothing: int,
-    texture_smoothing: int
+    texture_smoothing: int,
+    close_body: int
 ):
     if input_filepath 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.")
@@ -117,11 +118,22 @@ def create_3d_model(
         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]])
+
+    if close_body:
+        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], #top wall
+            [v_br, b_br, b_bl], [v_br, b_bl, v_bl], #bottom wall
+            [v_bl, b_bl, b_tl], [v_bl, b_tl, v_tl], #left
+            [v_tr, b_tr, b_br], [v_tr, b_br, v_br], #right
+            [b_tl, b_br, b_bl], [b_tl, b_tr, b_br]  #base
+        ])
+
     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
@@ -140,6 +152,7 @@ with gr.Blocks(theme='base') as demo:
             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")
+            close_body = gr.Checkbox(value=False, label="Close Body")
             
             generate_btn = gr.Button("Generate STL", variant="primary")
 
@@ -152,10 +165,10 @@ with gr.Blocks(theme='base') as demo:
         fn=create_3d_model,
         inputs=[
             input_image, texture_strength, max_z, min_z, x_length, do_pca,
-            depth_map_smoothing, texture_smoothing
+            depth_map_smoothing, texture_smoothing, close_body
         ],
         outputs=[output_model, output_file]
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)