mirror of
https://github.com/wassname/curl.git
synced 2026-08-22 12:00:28 +08:00
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import imageio
|
|
import os
|
|
import numpy as np
|
|
import logging
|
|
|
|
class VideoRecorder(object):
|
|
def __init__(self, dir_name, height=256, width=256, camera_id=0, fps=30):
|
|
self.dir_name = dir_name
|
|
self.height = height
|
|
self.width = width
|
|
self.camera_id = camera_id
|
|
self.fps = fps
|
|
self.frames = []
|
|
|
|
def init(self, enabled=True):
|
|
self.frames = []
|
|
self.enabled = self.dir_name is not None and enabled
|
|
|
|
def record(self, env, yaw=0):
|
|
if self.enabled:
|
|
try:
|
|
frame = env.render(
|
|
mode='rgb_array',
|
|
height=self.height,
|
|
width=self.width,
|
|
yaw=yaw
|
|
)
|
|
except:
|
|
frame = env.render(
|
|
mode='rgb_array',
|
|
)
|
|
|
|
self.frames.append(frame)
|
|
|
|
def save(self, file_name):
|
|
if self.enabled:
|
|
path = os.path.join(self.dir_name, file_name)
|
|
imageio.mimsave(path, self.frames, fps=self.fps)
|
|
|
|
logging.info(f"saved {len(self.frames)} frames to video: '{path}'")
|