feat: modify env handler has multi scene

web_server.py

@@ -9,12 +9,13 @@ import hugsim_env
 import subprocess as sp
 from collections import deque, OrderedDict
 from datetime import datetime
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
+from dataclasses import dataclass
 sys.path.append(os.getcwd())
 
 from fastapi import FastAPI, Body, Header, Depends, HTTPException, Query
 from fastapi.responses import HTMLResponse, Response
-from omegaconf import OmegaConf
+from omegaconf import OmegaConf, DictConfig
 from huggingface_hub import HfApi
 import open3d as o3d
 import numpy as np
@@ -155,29 +156,65 @@ class FifoDict:
         return self._order_dict.get(key, None)
 
 
+@dataclass
+class SceneConfig:
+    name: str
+    cfg: DictConfig
+
+
+@dataclass
+class EnvExecuteResult:
+    cur_scene_done: bool
+    done: bool
+
+
 class EnvHandler:
-    def __init__(self, cfg, output):
-        self.cfg = cfg
-        self.output = output
-        self.env = gymnasium.make('hugsim_env/HUGSim-v0', cfg=cfg, output=output)
-        self._lock = threading.Lock()
+    """A class to handle the environment for HUGSim.
+    This can include multiple scene and configurations.
+    """
+    def __init__(self, scene_list: List[SceneConfig], base_output: str):
+        self.scene_list = scene_list
+        self.base_output = base_output
+        self.env = None
         self.reset_env()
+        self._lock = threading.Lock()
+
+    def _switch_scene(self, scene_index: int):
+        """
+        Switch to a different scene based on the index.
+        Args:
+            scene_index (int): The index of the scene to switch to.
+        """
+        if scene_index < 0 or scene_index >= len(self.scene_list):
+            raise ValueError("Invalid scene index.")
+        
+        self.cur_scene_index = scene_index
+        scene_config = self.scene_list[scene_index]
+        self.env.close()
+        self.cur_otuput = os.path.join(self.base_output, scene_config.name)
+        self.env = gymnasium.make('hugsim_env/HUGSim-v0', cfg=scene_config.cfg, output=self.cur_otuput)
+        self._scene_cnt = 0
+        self._scene_done = False
+        self._save_data = {'type': 'closeloop', 'frames': []}
+        self._obs, self._info = self.env.reset()
+
+        self._log(f"Switched to scene: {scene_config.name}")
 
     def close(self):
         """
         Close the environment and release resources.
         """
-        del self.env
+        if self.env is not None:
+            del self.env
+        self.env = None
         self._log("Environment closed.")
 
     def reset_env(self):
         """
         Reset the environment and initialize variables.
         """
-        self._cnt = 0
         self._done = False
-        self._save_data = {'type': 'closeloop', 'frames': []}
-        self._obs, self._info = self.env.reset()
+        self._switch_scene(0)
         self._log_list = deque(maxlen=100)
         self._log("Environment reset complete.")
     
@@ -198,6 +235,15 @@ class EnvHandler:
             bool: True if the episode is done, False otherwise.
         """
         return self._done
+    
+    @property
+    def has_scene_done(self) -> bool:
+        """
+        Check if the current scene is done.
+        Returns:
+            bool: True if the current scene is done, False otherwise.
+        """
+        return self._scene_done
 
     @property
     def log_list(self) -> deque:
@@ -208,7 +254,7 @@ class EnvHandler:
         """
         return self._log_list
 
-    def execute_action(self, plan_traj: np.ndarray) -> bool:
+    def execute_action(self, plan_traj: np.ndarray) -> EnvExecuteResult:
         """
         Execute the action based on the planned trajectory.
         Args:
@@ -221,8 +267,8 @@ class EnvHandler:
         self._log("Executing action:", action)
     
         self._obs, _, terminated, truncated, self._info = self.env.step(action)
-        self._cnt += 1
-        self._done = terminated or truncated or self._cnt > 400
+        self._scene_cnt += 1
+        self._scene_done = terminated or truncated or self._scene_cnt > 400
 
         imu_plan_traj = plan_traj[:, [1, 0]]
         imu_plan_traj[:, 1] *= -1
@@ -241,20 +287,26 @@ class EnvHandler:
             'rc': self._info['rc']
         })
         
-        if not self._done:
-            return False
+        if not self._scene_done:
+            return EnvExecuteResult(cur_scene_done=False, done=False)
 
-        with open(os.path.join(self.output, 'data.pkl'), 'wb') as wf:
+        with open(os.path.join(self.cur_otuput, 'data.pkl'), 'wb') as wf:
             pickle.dump([self._save_data], wf)
         
-        ground_xyz = np.asarray(o3d.io.read_point_cloud(os.path.join(self.output, 'ground.ply')).points)
-        scene_xyz = np.asarray(o3d.io.read_point_cloud(os.path.join(self.output, 'scene.ply')).points)
+        ground_xyz = np.asarray(o3d.io.read_point_cloud(os.path.join(self.cur_otuput, 'ground.ply')).points)
+        scene_xyz = np.asarray(o3d.io.read_point_cloud(os.path.join(self.cur_otuput, 'scene.ply')).points)
         results = hugsim_evaluate([self._save_data], ground_xyz, scene_xyz)
-        with open(os.path.join(self.output, 'eval.json'), 'w') as f:
+        with open(os.path.join(self.cur_otuput, 'eval.json'), 'w') as f:
             json.dump(results, f)
         
-        self._log("Evaluation results saved.")
-        return True
+        self._log(f"Scene {self.cur_scene_index} completed. Evaluation results saved.")
+    
+        if self.cur_scene_index < len(self.scene_list) - 1:
+            self._switch_scene(self.cur_scene_index + 1)
+            return EnvExecuteResult(cur_scene_done=True, done=False)
+
+        self._done = True
+        return EnvExecuteResult(cur_scene_done=True, done=True)
 
     def _log(self, *messages):
         log_message = f"[{str(datetime.now())}]" + " ".join([str(msg) for msg in messages]) + "\n"
@@ -267,7 +319,12 @@ class EnvHandlerManager:
         self._env_handlers = {}
         self._lock = threading.Lock()
 
-    def _generate_env_handler(self, env_id: str):
+    def _get_scene_list(self, env_id: str, base_output: str) -> List[SceneConfig]:
+        """
+        Load the scene configurations from the YAML files.
+        Returns:
+            List[SceneConfig]: A list of scene configurations.
+        """
         base_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_base.yaml')
         scenario_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'scene-0383-medium-00.yaml')
         camera_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_camera.yaml')
@@ -288,11 +345,17 @@ class EnvHandlerManager:
         model_config = OmegaConf.load(os.path.join(model_path, 'cfg.yaml'))
         model_config.update({"model_path": "/app/app_datas/PAMI2024/release/ss/scenes/nuscenes/scene-0383"})
         cfg.update(model_config)
-        cfg.base.output_dir = "/app/app_datas/env_output"
-        
-        output = os.path.join(cfg.base.output_dir, f"{env_id}_hugsim_env")
+        cfg.base.output_dir = base_output
+        return [
+            SceneConfig(name=cfg.scenario.scene_name, cfg=cfg)
+        ]
+
+    def _generate_env_handler(self, env_id: str):
+        base_output = "/app/app_datas/env_output"
+        scene_list = self._get_scene_list(env_id, base_output)
+        output = os.path.join(base_output, f"{env_id}_hugsim_env")
         os.makedirs(output, exist_ok=True)
-        return EnvHandler(cfg, output)
+        return EnvHandler(scene_list, base_output=output)
 
     def exists_env_handler(self, env_id: str) -> bool:
         """
@@ -383,7 +446,12 @@ def get_current_state_endpoint(env_handler: EnvHandler = Depends(_get_env_handle
     Get the current state of the environment.
     """
     state = env_handler.get_current_state()
-    return Response(content=pickle.dumps({"done": env_handler.has_done, "state": state}), media_type="application/octet-stream")
+    data = {
+        "done": env_handler.has_done,
+        "cur_scene_done": env_handler.has_scene_done,
+        "state": state,
+    }
+    return Response(content=pickle.dumps(data), media_type="application/octet-stream")
 
 
 @app.post("/execute_action")
@@ -407,29 +475,29 @@ def execute_action_endpoint(
         return Response(content=cache_result, media_type="application/octet-stream")
 
     if env_handler.has_done:
-        result = pickle.dumps({"done": True, "state": env_handler.get_current_state()})
+        result = pickle.dumps({"done": True, "cur_scene_done": True, "state": env_handler.get_current_state()})
         _result_dict.push(transaction_id, result)
         return Response(content=result, media_type="application/octet-stream")
 
     plan_traj = _load_numpy_ndarray_json_str(plan_traj)
-    done = env_handler.execute_action(plan_traj)
-    if done:
+    execute_result = env_handler.execute_action(plan_traj)
+    if execute_result.done:
         token_info = get_token_info(auth_token)
         env_manager.close_env_handler(token_info["submission_id"])
         delete_client_space(token_info["client_space_id"])
         update_submission_status(token_info["team_id"], token_info["submission_id"], SubmissionStatus.SUCCESS.value)
         hf_api.upload_folder(
             repo_id=COMPETITION_ID,
-            folder_path=env_handler.output,
+            folder_path=env_handler.base_output,
             repo_type="dataset",
             path_in_repo=f"eval_results/{token_info['submission_id']}",
         )
-        result = pickle.dumps({"done": done, "state": env_handler.get_current_state()})
+        result = pickle.dumps({"done": execute_result.done, "cur_scene_done": execute_result.cur_scene_done, "state": env_handler.get_current_state()})
         _result_dict.push(transaction_id, result)
         return Response(content=result, media_type="application/octet-stream")
     
     state = env_handler.get_current_state()
-    result = pickle.dumps({"done": done, "state": state})
+    result = pickle.dumps({"done": execute_result.done, "cur_scene_done": execute_result.cur_scene_done, "state": state})
     _result_dict.push(transaction_id, result)
     return Response(content=result, media_type="application/octet-stream")