Enable loading and processing OBJ meshes from memory.

Also remove dependency on modules from data_utils to make it more hermetic.

PiperOrigin-RevId: 406120785

polygen/data_utils.py

@@ -14,7 +14,6 @@
 
 """Mesh data utilities."""
 import matplotlib.pyplot as plt
-import modules
 from mpl_toolkits import mplot3d  # pylint: disable=unused-import
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 import networkx as nx
@@ -88,8 +87,16 @@ def make_face_model_dataset(
       example['faces'] = tf.cast(
           tf.gather(face_permutation, example['faces']), tf.int64)
 
+    def _dequantize_verts(verts, n_bits):
+      min_range = -0.5
+      max_range = 0.5
+      range_quantize = 2**n_bits - 1
+      verts = tf.cast(verts, tf.float32)
+      verts = verts * (max_range - min_range) / range_quantize + min_range
+      return verts
+
     # Vertices are quantized. So convert to floats for input to face model
-    example['vertices'] = modules.dequantize_verts(vertices, quantization_bits)
+    example['vertices'] = _dequantize_verts(vertices, quantization_bits)
     example['vertices_mask'] = tf.ones_like(
         example['vertices'][..., 0], dtype=tf.float32)
     example['faces_mask'] = tf.ones_like(example['faces'], dtype=tf.float32)
@@ -97,14 +104,13 @@ def make_face_model_dataset(
   return ds.map(_face_model_map_fn)
 
 
-def read_obj(obj_path):
+def read_obj_file(obj_file):
-  """Read vertices and faces from .obj file."""
+  """Read vertices and faces from already opened file."""
   vertex_list = []
   flat_vertices_list = []
   flat_vertices_indices = {}
   flat_triangles = []
 
-  with open(obj_path) as obj_file:
   for line in obj_file:
     tokens = line.split()
     if not tokens:
@@ -135,6 +141,13 @@ def read_obj(obj_path):
   return np.array(flat_vertices_list, dtype=np.float32), flat_triangles
 
 
+def read_obj(obj_path):
+  """Open .obj file from the path provided and read vertices and faces."""
+
+  with open(obj_path) as obj_file:
+    return read_obj_file(obj_file)
+
+
 def write_obj(vertices, faces, file_path, transpose=True, scale=1.):
   """Write vertices and faces to obj."""
   if transpose:
@@ -286,10 +299,8 @@ def quantize_process_mesh(vertices, faces, tris=None, quantization_bits=8):
   return vertices, faces, tris
 
 
-def load_process_mesh(mesh_obj_path, quantization_bits=8):
+def process_mesh(vertices, faces, quantization_bits=8):
-  """Load obj file and process."""
+  """Process mesh vertices and faces."""
-  # Load mesh
-  vertices, faces = read_obj(mesh_obj_path)
 
   # Transpose so that z-axis is vertical.
   vertices = vertices[:, [2, 0, 1]]
@@ -316,6 +327,13 @@ def load_process_mesh(mesh_obj_path, quantization_bits=8):
   }
 
 
+def load_process_mesh(mesh_obj_path, quantization_bits=8):
+  """Load obj file and process."""
+  # Load mesh
+  vertices, faces = read_obj(mesh_obj_path)
+  return process_mesh(vertices, faces, quantization_bits)
+
+
 def plot_meshes(mesh_list,
                 ax_lims=0.3,
                 fig_size=4,