Modify post-processing

infer_asset.py

@@ -328,7 +328,7 @@ def compute_part_components_for_mesh_cc(mesh, mesh_cc_faces, current_face_part_i
     return components
 
 
-def refine_part_ids_for_faces(mesh, face_part_ids):
+def refine_part_ids_nonstrict(mesh, face_part_ids):
     """
     Refine face part IDs to ensure each part ID forms a single connected component.
     
@@ -346,8 +346,6 @@ def refine_part_ids_for_faces(mesh, face_part_ids):
     Returns:
         refined_face_part_ids: refined part ID for each face [num_faces]
     """
-    face_part_ids_final_strict = refine_part_ids_strict(mesh, face_part_ids)
-    
     face_part_ids_final = face_part_ids.copy()  # Don't modify the input
     
     # Step 1: Find connected components of the original mesh (immutable structure)
@@ -398,10 +396,10 @@ def refine_part_ids_for_faces(mesh, face_part_ids):
                     face_part_ids_final[comps[comp_idx]['faces']] = chosen_part_id
                     break
     
-    return face_part_ids_final_strict, face_part_ids_final
+    return face_part_ids_final
 
 
-def find_part_ids_for_faces(mesh, part_ids, face_indices):
+def find_part_ids_for_faces(mesh, part_ids, face_indices, strict=False):
     """
     Assign part IDs to each face in the mesh.
     
@@ -433,9 +431,11 @@ def find_part_ids_for_faces(mesh, part_ids, face_indices):
         counts = np.bincount(point_part_ids)
         majority_part_id = np.argmax(counts)
         face_part_ids[face_idx] = majority_part_id
-    
-    face_part_ids_refined_strict, face_part_ids_refined = refine_part_ids_for_faces(mesh, face_part_ids)
-    return face_part_ids, face_part_ids_refined_strict, face_part_ids_refined
+
+    if strict:
+        return refine_part_ids_strict(mesh, face_part_ids)
+    else:
+        return refine_part_ids_nonstrict(mesh, face_part_ids)
 
 
 @torch.no_grad()
@@ -501,12 +501,12 @@ def save_articulated_meshes(mesh, face_indices, outputs, output_path, strict, an
     prismatic_axis = outputs[hyp_idx]['prismatic_axis']
     prismatic_range = outputs[hyp_idx]['prismatic_range']
 
-    _, face_part_ids_refined_strict, face_part_ids_refined = find_part_ids_for_faces(
+    face_part_ids = find_part_ids_for_faces(
         mesh,
         part_ids,
-        face_indices
+        face_indices,
+        strict=strict
     )
-    face_part_ids = face_part_ids_refined_strict if strict else face_part_ids_refined
     unique_part_ids = np.unique(face_part_ids)
     num_parts = len(unique_part_ids)
     print(f"Found {num_parts} unique parts")