improved visual

app.py

@@ -11,8 +11,7 @@ from copy import deepcopy
 import numpy as np
 import re
 from bs4 import BeautifulSoup
-import torch
-
+import trimesh.transformations as tf
 
 import logging
 
@@ -38,7 +37,6 @@ max_new_tokens = 5000
 obj_context_feature_type = "text"
 
 
-
 tokenizer, model, data_loader = load_model_and_dataloader(
     model_path=model_path,
     model_base=model_base,
@@ -46,11 +44,12 @@ tokenizer, model, data_loader = load_model_and_dataloader(
     load_4bit=load_4bit,
     load_bf16=load_bf16,
     scene_to_obj_mapping=scene_to_obj_mapping,
-    device_map='cpu',
+    device_map="cpu",
 )  # Huggingface Zero-GPU has to use .to(device) to set the device, otherwise it will fail
 
 model.to("cuda")  # Huggingface Zero-GPU requires explicit device placement
 
+
 def get_chatbot_response(user_chat_input, scene_id):
     # Get the response from the model
     prompt, response = get_model_response(
@@ -61,10 +60,11 @@ def get_chatbot_response(user_chat_input, scene_id):
         user_input=user_chat_input,
         max_new_tokens=max_new_tokens,
         temperature=0.2,
-        top_p=0.9
+        top_p=0.9,
     )
     return scene_id, prompt, response
 
+
 # def get_chatbot_response(user_chat_input):
 #     # Get the response from the chatbot
 #     scene_id = "scene0643_00"
@@ -76,6 +76,7 @@ def get_chatbot_response(user_chat_input, scene_id):
 #     """
 #     return scene_id, scene_graph, response
 
+
 # Resetting to blank
 def reset_textbox():
     return gr.update(value="")
@@ -126,7 +127,7 @@ def cylinder_frame(p0, p1):
     return transformation
 
 
-def create_cylinder_mesh(p0, p1, color, radius=0.02, resolution=20, split=1):
+def create_cylinder_mesh(p0, p1, color, radius=0.04, resolution=20, split=1):
     """Create a colored cylinder mesh between two points p0 and p1."""
     cylinder = o3d.geometry.TriangleMesh.create_cylinder(
         radius=radius, height=1, resolution=resolution, split=split
@@ -164,15 +165,24 @@ def create_bbox(center, extents, color=[1, 0, 0], radius=0.02):
     center = center.replace("[", "").replace("]", "")
     extents = [float(x.strip()) for x in extents.split(",")]
     center = [float(x.strip()) for x in center.split(",")]
+    angle = -np.pi / 2  # 90 degrees
+    axis = [1, 0, 0]  # Rotate around x-axis
+    R = tf.rotation_matrix(angle, axis)
+    center_homogeneous = np.append(center, 1)
+    extents_homogeneous = np.append(extents, 1)
 
-    sx, sy, sz = float(extents[0]), float(extents[1]), float(extents[2])
+    # Apply the rotation to the center and extents
+    rotated_center = np.dot(R, center_homogeneous)[:3]
+    rotated_extents = np.dot(R, extents_homogeneous)[:3]
+
+    sx, sy, sz = rotated_extents
     x_corners = [sx / 2, sx / 2, -sx / 2, -sx / 2, sx / 2, sx / 2, -sx / 2, -sx / 2]
     y_corners = [sy / 2, -sy / 2, -sy / 2, sy / 2, sy / 2, -sy / 2, -sy / 2, sy / 2]
     z_corners = [sz / 2, sz / 2, sz / 2, sz / 2, -sz / 2, -sz / 2, -sz / 2, -sz / 2]
     corners_3d = np.vstack([x_corners, y_corners, z_corners])
-    corners_3d[0, :] = corners_3d[0, :] + float(center[0])
-    corners_3d[1, :] = corners_3d[1, :] + float(center[1])
-    corners_3d[2, :] = corners_3d[2, :] + float(center[2])
+    corners_3d[0, :] = corners_3d[0, :] + float(rotated_center[0])
+    corners_3d[1, :] = corners_3d[1, :] + float(rotated_center[1])
+    corners_3d[2, :] = corners_3d[2, :] + float(rotated_center[2])
     corners_3d = np.transpose(corners_3d)
 
     lines = [
@@ -201,7 +211,7 @@ def highlight_clusters_in_mesh(
     centroids_extends_refer,
     mesh,
     output_dir,
-    output_file_name="highlighted_mesh.glb",
+    output_file_name="highlighted_mesh.obj",
 ):
     print("*" * 50)
     # Visualize the highlighted points by drawing 3D bounding boxes overlay on a mesh
@@ -218,7 +228,9 @@ def highlight_clusters_in_mesh(
     for center, extent in centroids_extents_detailed:
         print("center: ", center)
         print("extent: ", extent)
-        bbox = create_bbox(center, extent, color=[0, 0, 1])  # Red color for all boxes
+        bbox = create_bbox(
+            center, extent, color=[1, 1, 0]
+        )  # yellow color for all boxes
         for b in bbox:
             combined_mesh += b
 
@@ -227,7 +239,6 @@ def highlight_clusters_in_mesh(
         for b in bbox:
             combined_mesh += b
 
-    combined_mesh = prettify_mesh_for_gradio(combined_mesh)
     # Save the combined mesh
     output_file_path = os.path.join(output_path, output_file_name)
     o3d.io.write_triangle_mesh(
@@ -261,6 +272,7 @@ def get_centroids_extents(obj_list, scene_dict):
             centroids_extents.append((centroid, extent))
     return centroids_extents
 
+
 @spaces.GPU
 def language_model_forward(
     session_state, user_chat_input, top_p, temperature, dropdown_scene
@@ -283,7 +295,9 @@ def language_model_forward(
     mesh = o3d.io.read_triangle_mesh(original_model_path)
 
     # get chatbot response
-    scene_id, scene_graph, response = get_chatbot_response(user_chat_input, session_state.scene)
+    scene_id, scene_graph, response = get_chatbot_response(
+        user_chat_input, session_state.scene
+    )
 
     assert scene_id == session_state.scene  # Ensure the scene ID matches
 
@@ -323,11 +337,13 @@ def language_model_forward(
         centroids_extents_refer,
         mesh,
         session_output_dir,
-        output_file_name="highlighted_model.glb",
+        output_file_name="highlighted_model.obj",
     )
-    
+
     # Update the chat history with the response
-    last_turn = session_state.chat_history_for_display[-1] # first is user input, second is assistant response
+    last_turn = session_state.chat_history_for_display[
+        -1
+    ]  # first is user input, second is assistant response
     last_turn = (last_turn[0], response)
     session_state.chat_history_for_display[-1] = last_turn
     session_state.save()  # save the session state
@@ -335,7 +351,7 @@ def language_model_forward(
     yield session_state, highlighted_model_path, session_state.chat_history_for_display
 
 
-title = """<h1 align="center">🤖 3D-GRAND: Towards Better Grounding and Less Hallucination for 3D-LLMs 🚀</h1>
+title = """<h1 align="center">🏠💬  3D-GRAND: Towards Better Grounding and Less Hallucination for 3D-LLMs 🚀</h1>
 <p><center>
 <a href="https://3d-grand.github.io/" target="_blank">[Project Page]</a>
 <a href="https://www.dropbox.com/scl/fo/5p9nb4kalnz407sbqgemg/AG1KcxeIS_SUoJ1hoLPzv84?rlkey=weunabtbiz17jitfv3f4jpmm1&dl=0" target="_blank">[3D-GRAND Data]</a>
@@ -350,7 +366,7 @@ title = """<h1 align="center">🤖 3D-GRAND: Towards Better Grounding and Less H
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     session_state = gr.State(Session.create)
-    
+
     gr.HTML(title)
 
     with gr.Column():
@@ -394,8 +410,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 )
                 gr.HTML(
                     """<center><strong>
-                    <div style="display:inline-block; color:blue">&#9632;</div> = Landmark &nbsp;
-                    <div style="display:inline-block; color:green">&#9632;</div> = Chosen Target
+                    <div style="display:inline-block; color:green">&#9632;</div> = Chosen Target &nbsp;
+                    <div style="display:inline-block; color:yellow">&#9632;</div> = Landmarks
                     </strong></center>
                     """
                 )
@@ -422,7 +438,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         gr.Examples(
                             examples=[
                                 ["The TV on the drawer, opposing the bed."],
-                                ["the desk next to the window"]
+                                ["the desk next to the window"],
                             ],
                             inputs=user_chat_input,
                         )
@@ -455,12 +471,22 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     dropdown_scene.change(
         fn=change_scene_or_system_prompt,
         inputs=[dropdown_scene],
-        outputs=[session_state, model_3d, model_3d_grounding_result, chat_history_for_display],
+        outputs=[
+            session_state,
+            model_3d,
+            model_3d_grounding_result,
+            chat_history_for_display,
+        ],
     )
     clear_button.click(
         fn=change_scene_or_system_prompt,
         inputs=[dropdown_scene],
-        outputs=[session_state, model_3d, model_3d_grounding_result, chat_history_for_display],
+        outputs=[
+            session_state,
+            model_3d,
+            model_3d_grounding_result,
+            chat_history_for_display,
+        ],
     )
     user_chat_input.submit(
         fn=language_model_forward,

convert_mesh.ipynb

@@ -7,11 +7,28 @@
    "outputs": [],
    "source": [
     "import trimesh\n",
+    "import numpy as np\n",
+    "import trimesh.transformations as tf\n",
     "\n",
-    "def convert_ply_to_format(ply_file, output_file):\n",
+    "# adjust bright factor to increase or decrease the brightness of the colors\n",
+    "def convert_ply_to_format(ply_file, output_file, bright_factor=1.5):\n",
     "    # Load the PLY file\n",
     "    mesh = trimesh.load(ply_file)\n",
     "\n",
+    "    # Define the rotation matrix to swap y and z axes\n",
+    "    angle = -np.pi / 2  # 90 degrees\n",
+    "    axis = [1, 0, 0]  # Rotate around x-axis\n",
+    "    R = tf.rotation_matrix(angle, axis)\n",
+    "\n",
+    "    # Apply the rotation to the mesh\n",
+    "    mesh.apply_transform(R)\n",
+    "    \n",
+    "    # Adjust the brightness of vertex colors if they exist\n",
+    "    if mesh.visual.kind == 'vertex' and mesh.visual.vertex_colors is not None:\n",
+    "        vertex_colors = np.asarray(mesh.visual.vertex_colors)[:, :3]  # Ignore alpha channel if present\n",
+    "        brightened_colors = np.clip(vertex_colors * bright_factor, 0, 255).astype(np.uint8)\n",
+    "        mesh.visual.vertex_colors = brightened_colors\n",
+    "\n",
     "    # Export the mesh to the specified format\n",
     "    mesh.export(output_file)\n",
     "    print(f\"Converted {ply_file} to {output_file}\")"
@@ -103,7 +120,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.10.14"
   }
  },
  "nbformat": 4,

data/scene0025_00/scene0025_00.obj

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6cfe6e6a671560d01aaf7e9cb6bb69872aa7867c79ca9cc1a87fa1594d59aa21
-size 18953438
+oid sha256:eb45f524155df059915b0f386f493a18b6e9564571a480121572734770ff6b78
+size 18951552

data/scene0426_00/scene0426_00.obj

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8832e13adccb571f351f561c13bf34c824c7a9e9f39f3f4e701430d9b612920
-size 14657919
+oid sha256:3c8562b2bb6c62350e234930ecc118d38a09ccc4502c2680c0ee76e0a49a8223
+size 14661106

data/scene0643_00/scene0643_00.obj

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:31bf1def086561935a3e8900378cb127dc723ef8c98b191596fd2ce3e78fbbd9
-size 18171881
+oid sha256:ce4f34b6e98523533fa20c7b4aa018d7b71dde857b6ea8c3f03541e10eaa32fa
+size 18165757

requirements.txt

@@ -16,4 +16,5 @@ open3d
 bs4
 attrs
 cattrs
+trimesh