Update scripts/utils.py

scripts/utils.py

@@ -1,11 +1,22 @@
+import os
 import torch
 import numpy as np
 from PIL import Image
 import pymeshlab
 import pymeshlab as ml
 from pymeshlab import PercentageValue
-from pytorch3d.renderer import TexturesVertex
+from pytorch3d.renderer import (
+    look_at_view_transform,
+    FoVPerspectiveCameras,
+    PointLights,
+    RasterizationSettings,
+    MeshRenderer,
+    MeshRasterizer,
+    SoftPhongShader,
+    TexturesVertex,
+)
 from pytorch3d.structures import Meshes
+import imageio
 from rembg import new_session, remove
 import torch
 import torch.nn.functional as F
@@ -239,17 +250,66 @@ def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to
         fix_vert_color_glb(save_glb_path)
     print(f"saving to {save_glb_path}")
 
-
-def save_glb_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=True, dist=3.5, azim_offset=180, resolution=512, fov_in_degrees=1 / 1.15, cam_type="ortho", view_padding=60, export_video=True) -> Tuple[str, str]:
+def save_glb_and_video(
+    save_mesh_prefix: str, 
+    meshes: Meshes, 
+    with_timestamp=True, 
+    dist=3.5, 
+    resolution=512, 
+    fov_in_degrees=45, 
+    cam_type="persp", 
+    export_video=True
+) -> Tuple[str, str]:
     import time
+
+    # Generar el nombre del archivo GLB
     if '.' in save_mesh_prefix:
         save_mesh_prefix = ".".join(save_mesh_prefix.split('.')[:-1])
     if with_timestamp:
         save_mesh_prefix = save_mesh_prefix + f"_{int(time.time())}"
     ret_mesh = save_mesh_prefix + ".glb"
-    # optimizied version
+
+    # Guardar el modelo GLB
     save_py3dmesh_with_trimesh_fast(meshes, ret_mesh)
-    return ret_mesh, None
+
+    # Si no se solicita el video, retornar solo el GLB
+    if not export_video:
+        return ret_mesh, None
+
+    # Configurar el directorio para guardar el video
+    output_dir = os.path.dirname(save_mesh_prefix)
+    os.makedirs(output_dir, exist_ok=True)
+    video_path = save_mesh_prefix + "_video.mp4"
+
+    # Configurar la cámara
+    device = meshes.device
+    R, T = look_at_view_transform(dist=dist, elev=0, azim=np.linspace(0, 360, 72, endpoint=False))
+    cameras = FoVPerspectiveCameras(device=device, R=R, T=T, fov=fov_in_degrees)
+
+    # Configurar el renderizador
+    raster_settings = RasterizationSettings(
+        image_size=resolution,
+        blur_radius=0.0,
+        faces_per_pixel=1,
+    )
+    lights = PointLights(device=device, location=[[0.0, 0.0, -3.0]])
+    renderer = MeshRenderer(
+        rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+        shader=SoftPhongShader(device=device, cameras=cameras, lights=lights),
+    )
+
+    # Renderizar fotogramas
+    frames = []
+    for i in range(len(cameras)):
+        camera = FoVPerspectiveCameras(device=device, R=R[i].unsqueeze(0), T=T[i].unsqueeze(0), fov=fov_in_degrees)
+        image = renderer(meshes, cameras=camera)
+        image = (image[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)  # Convertir a imagen RGB
+        frames.append(image)
+
+    # Guardar los fotogramas como un video
+    imageio.mimsave(video_path, frames, fps=30)
+
+    return ret_mesh, video_path
 
 
 def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, apply_sub_divide=False, sub_divide_threshold=0.25):