Update app.py

app.py

@@ -9,6 +9,8 @@ from gymnasium import spaces
 import numpy as np
 from PIL import Image
 import gradio as gr
+import io
+import base64
 
 # Environment Variables
 HOSTED_API_URL = os.getenv("HOSTED_API_URL")  # FastAPI backend URL
@@ -51,7 +53,7 @@ def analyze_traffic(image_path):
             if processed_image_url:
                 img_response = requests.get(processed_image_url)
                 if img_response.status_code == 200:
-                    processed_image = Image.open(BytesIO(img_response.content))
+                    processed_image = Image.open(io.BytesIO(img_response.content))
             
             return vehicle_count, congestion_level, flow_rate, processed_image
         else:
@@ -85,7 +87,7 @@ class TrafficSimEnv(gym.Env):
         super().reset(seed=seed)
         self.state = np.array([self.congestion_level, 30], dtype=np.float32)
         self.done = False
-        return self.state
+        return self.state, {}
 
     def step(self, action):
         if self.done:
@@ -144,7 +146,7 @@ def optimize_signal_rl(congestion_level):
         model.learn(total_timesteps=1000)
 
         # Reset the environment and get the initial observation
-        obs = env.reset()
+        obs, _ = env.reset()  # Ensure this is a tuple (obs, info)
         logging.debug(f"Reset observation: {obs}, type: {type(obs)}, shape: {np.shape(obs)}")
 
         # Ensure obs is a valid array
@@ -153,7 +155,7 @@ def optimize_signal_rl(congestion_level):
         # RL Optimization loop
         for _ in range(10):
             action, _ = model.predict(obs, deterministic=True)
-            obs, rewards, dones, infos = env.step(action)
+            obs, rewards, dones, _, infos = env.step(action)
             
             # Debug observation structure
             logging.debug(f"Step observation: {obs}, type: {type(obs)}, shape: {np.shape(obs)}")
@@ -253,4 +255,4 @@ if __name__ == "__main__":
     )
 
     # Launch Gradio app
-    interface.launch()
+    interface.launch()