@ physigym now videos are saved

Author: Dante-Berth

rl/sb/stable_baselines.py

@@ -1,22 +1,22 @@
-
 import gymnasium as gym
 import numpy as np
 import os
 import wandb
 import physigym
 import stable_baselines3
 import sb3_contrib
-from stable_baselines3 import SAC, PPO
 from sb3_contrib import TQC
 from stable_baselines3.common.callbacks import BaseCallback
 from stable_baselines3.common.logger import configure
 import tyro
 import time
 from gymnasium.spaces import Box
 from dataclasses import dataclass
-from extending import physicell #from embedding import physicell
+from embedding import physicell  # from extending import physicell
 import matplotlib.pyplot as plt
 import pandas as pd
+
+
 # ----------------------
 # 🌟 Dataclass
 # ----------------------
@@ -36,6 +36,7 @@ class Args:
     seed: int = 1
     """seed"""
 
+
 args = tyro.cli(Args)
 config = vars(args)
 # ----------------------
@@ -46,10 +47,10 @@ class Args:
 ##### 📍 Choose Algorithm (SB3 or SB3-Contrib)
 ##### ----------------------
 algo_name = args.algo_name
-if algo_name  in sb3_contrib.__all__:
-    algorithm = getattr(sb3_contrib,algo_name)
+if algo_name in sb3_contrib.__all__:
+    algorithm = getattr(sb3_contrib, algo_name)
 elif algo_name in stable_baselines3.__all__:
-    algorithm = getattr(stable_baselines3,algo_name)
+    algorithm = getattr(stable_baselines3, algo_name)
 else:
     raise f"Algorith name does not exist: {algo_name}"
 
@@ -83,10 +84,12 @@ def __init__(self, verbose=0, video_frequency=50000):
     def _on_step(self) -> bool:
         # Get information from the environment
         if "reward" in self.locals:
-            self.logger.record("env/reward_value", self.locals["rewards"][0] )
+            self.logger.record("env/reward_value", self.locals["rewards"][0])
 
         if "number_cancer_cells" in self.locals["infos"][0]:
-            self.logger.record("env/cancer_cell_count", self.locals["infos"][0]["number_cancer_cells"])
+            self.logger.record(
+                "env/cancer_cell_count", self.locals["infos"][0]["number_cancer_cells"]
+            )
 
         if "actions" in self.locals:
             actions = self.locals["actions"][0]
@@ -97,6 +100,7 @@ def _on_step(self) -> bool:
         self.logger.dump(step=self.global_step)
         return True
 
+
 class PhysiCellModelWrapper(gym.Wrapper):
     def __init__(
         self,
@@ -179,42 +183,59 @@ def step(self, action: np.ndarray):
         )
         # Preprocess observation (if needed)
         o_observation = np.array(o_observation, dtype=float)
-        info["action"]  = d_action
+        info["action"] = d_action
         self.info = info
         return o_observation, r_reward, b_terminated, b_truncated, info
-    
-    def render(self, path="./output/image"):
-        os.makedirs(path,exist_ok=True)
-        df_cell = pd.DataFrame(physicell.get_cell(), columns=['ID','x','y','z','dead','cell_type'])
-        fig, ax = plt.subplots(1, 3, figsize=(10, 6), gridspec_kw={'width_ratios': [1, 0.2, 0.2]})
-
-        for s_celltype, s_color in sorted({'cancer_cell': 'gray', 'nurse_cell': 'red'}.items()):
-            df_celltype = df_cell.loc[(df_cell.z == 0.0) & (df_cell.cell_type == s_celltype), :]
+
+    def render(
+        self,
+        path="./output/image",
+        saving_title: str = "output_simulation_image_episode",
+    ):
+        os.makedirs(path, exist_ok=True)
+        df_cell = pd.DataFrame(
+            physicell.get_cell(), columns=["ID", "x", "y", "z", "dead", "cell_type"]
+        )
+        fig, ax = plt.subplots(
+            1, 3, figsize=(10, 6), gridspec_kw={"width_ratios": [1, 0.2, 0.2]}
+        )
+        count_cancer_cell = physicell.get_parameter("count_cancer_cell")
+
+        for s_celltype, s_color in sorted(
+            {"cancer_cell": "gray", "nurse_cell": "red"}.items()
+        ):
+            df_celltype = df_cell.loc[
+                (df_cell.z == 0.0) & (df_cell.cell_type == s_celltype), :
+            ]
             df_celltype.plot(
-                kind='scatter', x='x', y='y', c=s_color,
+                kind="scatter",
+                x="x",
+                y="y",
+                c=s_color,
                 xlim=[
-                   self.x_min,
+                    self.x_min,
                     self.x_max,
                 ],
                 ylim=[
                     self.y_min,
                     self.y_max,
                 ],
                 grid=True,
-                label = s_celltype,
+                label=s_celltype,
                 s=100,
-                title=f"episode step {str(self.unwrapped_env.step_episode).zfill(3)}",
+                title=f"episode step {str(self.unwrapped_env.step_episode).zfill(3)}, cancer cell: {count_cancer_cell}",
                 ax=ax[0],
-            ).legend(loc='lower left')
-
+            ).legend(loc="lower left")
 
         # Create a colormap for the color bars (from -1 to 1)
-        list_colors = ["royalblue","darkorange"]
+        list_colors = ["royalblue", "darkorange"]
 
         # Function to create fluid-like color bars
         def create_fluid_bar(ax_bar, drug_amount, title, max_amount=30, color="cyan"):
             ax_bar.set_xlim(0, 1)
-            ax_bar.set_ylim(0, 1)  # Set y-axis from 0 to 1 for percentage representation
+            ax_bar.set_ylim(
+                0, 1
+            )  # Set y-axis from 0 to 1 for percentage representation
             ax_bar.set_title(title, fontsize=10)
             ax_bar.set_xticks([])
             ax_bar.set_yticks(np.linspace(0, 1, 5))  # 0% to 100% scale
@@ -226,35 +247,103 @@ def create_fluid_bar(ax_bar, drug_amount, title, max_amount=30, color="cyan"):
             ax_bar.fill_betweenx(np.linspace(0, fill_level, 100), 0, 1, color=color)
 
             # Draw container border
-            ax_bar.spines['left'].set_visible(False)
-            ax_bar.spines['right'].set_visible(False)
-            ax_bar.spines['top'].set_visible(True)
-            ax_bar.spines['bottom'].set_visible(True)
-
+            ax_bar.spines["left"].set_visible(False)
+            ax_bar.spines["right"].set_visible(False)
+            ax_bar.spines["top"].set_visible(True)
+            ax_bar.spines["bottom"].set_visible(True)
 
         action = self.info["action"]
         for i, (key, value) in enumerate(action.items(), start=1):  # Start index from 1
-            create_fluid_bar(ax[i], value[0], f"drug_{i}", color=list_colors[i-1])
+            create_fluid_bar(ax[i], value[0], f"drug_{i}", color=list_colors[i - 1])
 
-        # fig.savefig(f"output_image_{self.unwrapped_env.step_episode}.png", bbox_inches='tight')
-        # Convert figure to NumPy array (store frame)
-        plt.savefig(path+f"/output_simulation_image_episode step {str(self.unwrapped_env.step_episode).zfill(3)}")
+        plt.savefig(
+            path
+            + f"/{saving_title} step {str(self.unwrapped_env.step_episode).zfill(3)}"
+        )
         plt.close(fig)
 
+
 import subprocess
+
+
 def png_to_video_ffmpeg(image_folder, output_video, fps=10):
     command = [
-        "ffmpeg", "-framerate", str(fps),
-        "-pattern_type", "glob", "-i", f"{image_folder}/*.png",
-        "-c:v", "libx264", "-pix_fmt", "yuv420p",
-        output_video
+        "ffmpeg",
+        "-framerate",
+        str(fps),
+        "-pattern_type",
+        "glob",
+        "-i",
+        f"{image_folder}/*.png",
+        "-c:v",
+        "libx264",
+        "-pix_fmt",
+        "yuv420p",
+        output_video,
     ]
     subprocess.run(command, check=True)
     print(f"✅ Video saved as {output_video}")
+
+
+import os
+import glob
+import imageio
+import imageio.v3 as iio  # Newer version of imageio
+import imageio_ffmpeg  # Ensure ffmpeg support
+
+
+def png_to_video_imageio(image_folder, output_video, fps=10):
+    images = sorted(glob.glob(os.path.join(image_folder, "*.png")))
+
+    if not images:
+        print("❌ No images found in the directory:", image_folder)
+        return
+
+    print(f"🖼️ Found {len(images)} images. First image: {images[0]}")
+
+    # Read first image to get size
+    frame = iio.imread(images[0])
+    height, width, _ = frame.shape
+    print(f"📏 Image size: {width}x{height}")
+
+    writer = imageio.get_writer(
+        output_video, fps=fps, codec="libx264", format="FFMPEG", pixelformat="yuv420p"
+    )
+
+    for img in images:
+        frame = iio.imread(img)
+        writer.append_data(frame)
+
+    writer.close()
+    print(f"✅ Video saved as {output_video}")
+
+
+def _video_save(
+    env,
+    seed,
+    step,
+    image_folder="./output/image",
+    deterministic=False,
+    wandb_path="test/simulation_video",
+    wandb=wandb,
+):
+    output_video = f"seed_{seed}_step_{step}.mp4"
+    obs, info = env.reset(seed=seed)
+    done = False
+    while not done:
+        action, _states = model.predict(obs, deterministic=deterministic)
+        obs, reward, terminated, truncated, info = env.step(action)
+        env.render()
+        if terminated or truncated:
+            png_to_video_imageio(image_folder, output_video, fps=10)
+            wandb.log({wandb_path: wandb.Video(output_video, fps=10, format="mp4")})
+            obs, info = env.reset(seed=args.seed)
+
+
 # ----------------------
 # 🏗️  Environment Setup
 # ----------------------
-env = gym.make(args.env_id,observation_type=args.observation_type)
+env = gym.make(args.env_id, observation_type=args.observation_type)
 env = PhysiCellModelWrapper(env)
 env = gym.wrappers.RescaleAction(env, min_action=-1, max_action=1)
 env = gym.wrappers.GrayscaleObservation(env)
@@ -264,7 +353,7 @@ def png_to_video_ffmpeg(image_folder, output_video, fps=10):
 # ----------------------
 # 📂 Logging Setup
 # ----------------------
-log_dir = f"./tensorboard_logs/{algo_name}" 
+log_dir = f"./tensorboard_logs/{algo_name}"
 os.makedirs(log_dir, exist_ok=True)
 
 # ----------------------
@@ -273,28 +362,24 @@ def png_to_video_ffmpeg(image_folder, output_video, fps=10):
 model = algorithm("CnnPolicy", env, verbose=1, tensorboard_log=log_dir, seed=args.seed)
 new_logger = configure(log_dir, ["tensorboard"])
 model.set_logger(new_logger)
-model.learn(total_timesteps=int(1e6), log_interval=1, progress_bar=False, callback=TensorboardCallback())
-path_saving_model = run_name+"/model" 
-model.save(path_saving_model)
 # ✅ Finish WandB run
-del model # remove to demonstrate saving and loading
-wandb.finish()  # ✅ Finish WandB run
-
+# del model # remove to demonstrate saving and loading
 # ----------------------
 # 🎮 Run the Trained Agent
 # ----------------------
-model = algorithm.load(path_saving_model) # load model
-obs, info = env.reset()
-dictionnary = {}
-for i in range(5):
-    step = 0
-    while True:
-        action, _states = model.predict(obs, deterministic=True)
-        obs, reward, terminated, truncated, info = env.step(action)
-        step +=1
-        if terminated or truncated:
-            png_to_video_ffmpeg("./output/image", f"output_video_{i}.mp4", fps=10)
-            obs, info = env.reset()
-print("Finished")
+# model = algorithm.load(path_saving_model) # load model
+for i in range(10):
+    _video_save(env=env, seed=args.seed, step=(i) * 25000, wandb=wandb)
+    model.learn(
+        total_timesteps=int(25000),
+        log_interval=1,
+        progress_bar=False,
+        callback=TensorboardCallback(),
+    )
+    # _video_save(env=env,seed=args.seed, step=(i+1)*25000,wandb=wandb)
+
+path_saving_model = run_name + "/model"
+model.save(path_saving_model)
 
-    
+print("Finished")
+wandb.finish()  # ✅ Finish WandB run