feat: update generate video

Dockerfile

@@ -30,6 +30,8 @@ RUN ./.pixi/envs/default/bin/python3 -m ensurepip
 
 RUN ./.pixi/envs/default/bin/python3 -m pip install psutil
 
+RUN ./.pixi/envs/default/bin/python3 -m pip install moviepy
+
 RUN ./.pixi/envs/default/bin/python /app/download_pre_datas.py
 
 CMD ["./.pixi/envs/default/bin/python", "web_server.py"]

web_server.py

@@ -14,6 +14,7 @@ from typing import Any, Dict, Optional, List
 from dataclasses import dataclass
 sys.path.append(os.getcwd())
 
+from moviepy import ImageSequenceClip
 from fastapi import FastAPI, Body, Header, Depends, HTTPException, Query
 from fastapi.responses import HTMLResponse, Response
 from omegaconf import OmegaConf, DictConfig
@@ -43,6 +44,17 @@ class SubmissionStatus(enum.Enum):
     FAILED = 4
 
 
+def to_video(observations: List[Any], output_path: str):
+    frames = []
+    for obs in observations:
+        row1 = np.concatenate([obs['CAM_FRONT_LEFT'], obs['CAM_FRONT'], obs['CAM_FRONT_RIGHT']], axis=1)
+        row2 = np.concatenate([obs['CAM_BACK_RIGHT'], obs['CAM_BACK'], obs['CAM_BACK_LEFT']], axis=1)
+        frame = np.concatenate([row1, row2], axis=0)
+        frames.append(frame)
+    clip = ImageSequenceClip(frames, fps=4)
+    clip.write_videofile(output_path)
+
+
 def get_gpu_memory():
     output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
     COMMAND = "nvidia-smi --query-gpu=memory.used --format=csv"
@@ -198,6 +210,7 @@ class EnvHandler:
         self._scene_cnt = 0
         self._scene_done = False
         self._save_data = {'type': 'closeloop', 'frames': []}
+        self._observations_save = []
         self._obs, self._info = self.env.reset()
 
         self._log(f"Switched to scene: {scene_config.name}")
@@ -289,6 +302,7 @@ class EnvHandler:
             'collision': self._info['collision'],
             'rc': self._info['rc']
         })
+        self._observations_save.append(self._obs['rgb'])
         
         if not self._scene_done:
             return EnvExecuteResult(cur_scene_done=False, done=False)
@@ -302,6 +316,7 @@ class EnvHandler:
         with open(os.path.join(self.cur_otuput, 'eval.json'), 'w') as f:
             json.dump(results, f)
         self._score_list.append(results.copy())
+        to_video(self._observations_save, os.path.join(self.cur_otuput, 'video.mp4'))
         
         self._log(f"Scene {self.cur_scene_index} completed. Evaluation results saved.")