Spaces:
Running
on
Zero
Running
on
Zero
import io | |
import random | |
import struct | |
import tempfile | |
from contextlib import contextmanager | |
from typing import List, Union | |
import numpy as np | |
import PIL.Image | |
import PIL.ImageOps | |
from .import_utils import ( | |
BACKENDS_MAPPING, | |
is_opencv_available, | |
) | |
from .logging import get_logger | |
global_rng = random.Random() | |
logger = get_logger(__name__) | |
def buffered_writer(raw_f): | |
f = io.BufferedWriter(raw_f) | |
yield f | |
f.flush() | |
def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None) -> str: | |
if output_gif_path is None: | |
output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name | |
image[0].save( | |
output_gif_path, | |
save_all=True, | |
append_images=image[1:], | |
optimize=False, | |
duration=100, | |
loop=0, | |
) | |
return output_gif_path | |
def export_to_ply(mesh, output_ply_path: str = None): | |
""" | |
Write a PLY file for a mesh. | |
""" | |
if output_ply_path is None: | |
output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply").name | |
coords = mesh.verts.detach().cpu().numpy() | |
faces = mesh.faces.cpu().numpy() | |
rgb = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1) | |
with buffered_writer(open(output_ply_path, "wb")) as f: | |
f.write(b"ply\n") | |
f.write(b"format binary_little_endian 1.0\n") | |
f.write(bytes(f"element vertex {len(coords)}\n", "ascii")) | |
f.write(b"property float x\n") | |
f.write(b"property float y\n") | |
f.write(b"property float z\n") | |
if rgb is not None: | |
f.write(b"property uchar red\n") | |
f.write(b"property uchar green\n") | |
f.write(b"property uchar blue\n") | |
if faces is not None: | |
f.write(bytes(f"element face {len(faces)}\n", "ascii")) | |
f.write(b"property list uchar int vertex_index\n") | |
f.write(b"end_header\n") | |
if rgb is not None: | |
rgb = (rgb * 255.499).round().astype(int) | |
vertices = [ | |
(*coord, *rgb) | |
for coord, rgb in zip( | |
coords.tolist(), | |
rgb.tolist(), | |
) | |
] | |
format = struct.Struct("<3f3B") | |
for item in vertices: | |
f.write(format.pack(*item)) | |
else: | |
format = struct.Struct("<3f") | |
for vertex in coords.tolist(): | |
f.write(format.pack(*vertex)) | |
if faces is not None: | |
format = struct.Struct("<B3I") | |
for tri in faces.tolist(): | |
f.write(format.pack(len(tri), *tri)) | |
return output_ply_path | |
def export_to_obj(mesh, output_obj_path: str = None): | |
if output_obj_path is None: | |
output_obj_path = tempfile.NamedTemporaryFile(suffix=".obj").name | |
verts = mesh.verts.detach().cpu().numpy() | |
faces = mesh.faces.cpu().numpy() | |
vertex_colors = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1) | |
vertices = [ | |
"{} {} {} {} {} {}".format(*coord, *color) for coord, color in zip(verts.tolist(), vertex_colors.tolist()) | |
] | |
faces = ["f {} {} {}".format(str(tri[0] + 1), str(tri[1] + 1), str(tri[2] + 1)) for tri in faces.tolist()] | |
combined_data = ["v " + vertex for vertex in vertices] + faces | |
with open(output_obj_path, "w") as f: | |
f.writelines("\n".join(combined_data)) | |
def export_to_video( | |
video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 8 | |
) -> str: | |
if is_opencv_available(): | |
import cv2 | |
else: | |
raise ImportError(BACKENDS_MAPPING["opencv"][1].format("export_to_video")) | |
if output_video_path is None: | |
output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name | |
if isinstance(video_frames[0], PIL.Image.Image): | |
video_frames = [np.array(frame) for frame in video_frames] | |
fourcc = cv2.VideoWriter_fourcc(*"mp4v") | |
h, w, c = video_frames[0].shape | |
video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=fps, frameSize=(w, h)) | |
for i in range(len(video_frames)): | |
img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR) | |
video_writer.write(img) | |
return output_video_path | |