Speed up save_obj text formatting, tidy _save helpers

Summary:
Format OBJ text through detached numpy views instead of per-element tensor indexing, and simplify the helpers accordingly.

___

Differential Revision: D119966991

fbshipit-source-id: a3c0bc069431eb1feb913069de922914c2e1ab5e

pytorch3d/io/obj_io.py

@@ -834,24 +834,22 @@ def _save(
         return
 
     verts, faces = verts.cpu(), faces.cpu()
+    vert_rows = verts.detach().cpu().numpy()
 
-    lines = ""
+    lines = []
 
     if decimal_places is None:
         float_str = "%f"
     else:
         float_str = "%" + ".%df" % decimal_places
 
-    if len(verts):
+    for vert in vert_rows:
-        V, D = verts.shape
+        lines.append("v %s\n" % " ".join(float_str % v for v in vert))
-        for i in range(V):
-            vert = [float_str % verts[i, j] for j in range(D)]
-            lines += "v %s\n" % " ".join(vert)
 
     if save_normals:
         assert normals is not None
         assert faces_normals_idx is not None
-        lines += _write_normals(normals, faces_normals_idx, float_str)
+        lines.append(_write_normals(normals, float_str))
 
     if save_uvs:
         assert faces_uvs is not None
@@ -866,15 +864,13 @@ def _save(
             raise ValueError(message)
 
         verts_uvs, faces_uvs = verts_uvs.cpu(), faces_uvs.cpu()
+        uv_rows = verts_uvs.detach().cpu().numpy()
 
         # Save verts uvs after verts
-        if len(verts_uvs):
+        for uv in uv_rows:
-            uV, uD = verts_uvs.shape
+            lines.append("vt %s\n" % " ".join(float_str % u for u in uv))
-            for i in range(uV):
-                uv = [float_str % verts_uvs[i, j] for j in range(uD)]
-                lines += "vt %s\n" % " ".join(uv)
 
-    f.write(lines)
+    f.write("".join(lines))
 
     if torch.any(faces >= verts.shape[0]) or torch.any(faces < 0):
         warnings.warn("Faces have invalid indices")
@@ -888,26 +884,16 @@ def _save(
         )
 
 
-def _write_normals(
+def _write_normals(normals: torch.Tensor, float_str: str) -> str:
-    normals: torch.Tensor, faces_normals_idx: torch.Tensor, float_str: str
-) -> str:
-    if faces_normals_idx.dim() != 2 or faces_normals_idx.size(1) != 3:
-        message = (
-            "'faces_normals_idx' should either be empty or of shape (num_faces, 3)."
-        )
-        raise ValueError(message)
-
     if normals.dim() != 2 or normals.size(1) != 3:
         message = "'normals' should either be empty or of shape (num_verts, 3)."
         raise ValueError(message)
 
-    normals, faces_normals_idx = normals.cpu(), faces_normals_idx.cpu()
+    normal_rows = normals.detach().cpu().numpy()
 
     lines = []
-    V, D = normals.shape
+    for normal in normal_rows:
-    for i in range(V):
+        lines.append("vn %s\n" % " ".join(float_str % n for n in normal))
-        normal = [float_str % normals[i, j] for j in range(D)]
-        lines.append("vn %s\n" % " ".join(normal))
     return "".join(lines)
 
 
@@ -917,31 +903,37 @@ def _write_faces(
     faces_uvs: Optional[torch.Tensor],
     faces_normals_idx: Optional[torch.Tensor],
 ) -> None:
+    face_rows = faces.detach().cpu().numpy()
+    uv_rows = faces_uvs.detach().cpu().numpy() if faces_uvs is not None else None
+    normal_rows = (
+        faces_normals_idx.detach().cpu().numpy()
+        if faces_normals_idx is not None
+        else None
+    )
+
     F, P = faces.shape
     for i in range(F):
-        if faces_normals_idx is not None:
+        face_row = face_rows[i]
-            if faces_uvs is not None:
+        if normal_rows is not None:
+            normal_row = normal_rows[i]
+            if uv_rows is not None:
+                uv_row = uv_rows[i]
                 # Format faces as {verts_idx}/{verts_uvs_idx}/{verts_normals_idx}
                 face = [
-                    "%d/%d/%d"
+                    "%d/%d/%d" % (face_row[j] + 1, uv_row[j] + 1, normal_row[j] + 1)
-                    % (
-                        faces[i, j] + 1,
-                        faces_uvs[i, j] + 1,
-                        faces_normals_idx[i, j] + 1,
-                    )
                     for j in range(P)
                 ]
             else:
                 # Format faces as {verts_idx}//{verts_normals_idx}
                 face = [
-                    "%d//%d" % (faces[i, j] + 1, faces_normals_idx[i, j] + 1)
+                    "%d//%d" % (face_row[j] + 1, normal_row[j] + 1) for j in range(P)
-                    for j in range(P)
                 ]
-        elif faces_uvs is not None:
+        elif uv_rows is not None:
+            uv_row = uv_rows[i]
             # Format faces as {verts_idx}/{verts_uvs_idx}
-            face = ["%d/%d" % (faces[i, j] + 1, faces_uvs[i, j] + 1) for j in range(P)]
+            face = ["%d/%d" % (face_row[j] + 1, uv_row[j] + 1) for j in range(P)]
         else:
-            face = ["%d" % (faces[i, j] + 1) for j in range(P)]
+            face = ["%d" % (face_row[j] + 1) for j in range(P)]
 
         if i + 1 < F:
             f.write("f %s\n" % " ".join(face))