Update app.py

app.py

@@ -1,105 +1,196 @@
+import gradio as gr
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import numpy as np
 import random
 import threading
 import time
-import matplotlib.pyplot as plt
-import gradio as gr
 
+# ---------- GLOBAL STATE ----------
 class SLAMSimulation:
     def __init__(self):
-        self.robot_pos = [5, 5]  # start position
-        self.grid_size = 10
+        self.pose = {"x": 0, "z": 0, "angle": 0}
+        self.trajectory = [(0, 0)]
+        self.obstacle_hits = []
+        self.color_index = 0
+        self.rgb_colors = ['red', 'green', 'blue']
+        self.noise_enabled = True
+        self.obstacles = []
         self.auto_mode = False
-        self.lock = threading.Lock()
+        self.reset(10)
+
+    def generate_obstacles(self, count=10):
+        return [{
+            "x": random.uniform(-8, 8),
+            "z": random.uniform(-8, 8),
+            "radius": random.uniform(0.5, 1.2)
+        } for _ in range(count)]
+
+    def toggle_noise(self):
+        self.noise_enabled = not self.noise_enabled
+        return f"Noise: {'ON' if self.noise_enabled else 'OFF'}"
+
+    def reset(self, obstacle_count):
+        self.pose = {"x": 0, "z": 0, "angle": 0}
+        self.trajectory = [(0, 0)]
+        self.obstacle_hits.clear()
+        self.color_index = 0
+        self.obstacles = self.generate_obstacles(int(obstacle_count))
+        return self.render_env(), self.render_slam_map(), f"🔄 Simulation reset with {obstacle_count} obstacles"
+
+    def check_collision(self, x, z):
+        return any(np.hypot(obs["x"] - x, obs["z"] - z) <= obs["radius"] + 0.2 for obs in self.obstacles)
+
+    def move_robot(self, direction):
+        step = 1
+        direction = direction.upper()
+        dx, dz, angle = {
+            "W": (0, step, 90),
+            "S": (0, -step, -90),
+            "A": (-step, 0, 180),
+            "D": (step, 0, 0)
+        }.get(direction, (0, 0, self.pose["angle"]))
+
+        if direction not in "WASD":
+            return self.render_env(), self.render_slam_map(), "❌ Invalid direction"
+
+        new_x, new_z = self.pose["x"] + dx, self.pose["z"] + dz
+
+        if self.check_collision(new_x, new_z):
+            return self.render_env(), self.render_slam_map(), "🚫 Collision detected!"
+
+        self.pose.update({"x": new_x, "z": new_z, "angle": angle})
+
+        noisy_x = new_x + random.uniform(-0.1, 0.1) if self.noise_enabled else new_x
+        noisy_z = new_z + random.uniform(-0.1, 0.1) if self.noise_enabled else new_z
+        self.trajectory.append((noisy_x, noisy_z))
+
+        return self.render_env(), self.render_slam_map(), f"✅ Moved {direction}"
 
     def render_env(self):
-        fig, ax = plt.subplots()
-        ax.set_xlim(0, self.grid_size)
-        ax.set_ylim(0, self.grid_size)
-        ax.grid(True)
-        ax.plot(self.robot_pos[0], self.robot_pos[1], 'ro', markersize=15)
-        ax.set_title("Environment")
+        fig, ax = plt.subplots(figsize=(5, 5))
+        ax.set_xlim(-10, 10)
+        ax.set_ylim(-10, 10)
+        ax.set_title("SLAM Environment View")
+
+        try:
+            bg = mpimg.imread("map.png")
+            ax.imshow(bg, extent=(-10, 10, -10, 10), alpha=0.2)
+        except FileNotFoundError:
+            pass
+
+        for obs in self.obstacles:
+            circle = plt.Circle((obs["x"], obs["z"]), obs["radius"], facecolor="gray", alpha=0.6, edgecolor="black")
+            ax.add_patch(circle)
+
+        ax.plot(self.pose["x"], self.pose["z"], 'ro', markersize=8)
+        self.obstacle_hits.clear()
+
+        angles = np.linspace(0, 2 * np.pi, 24)
+        for ang in angles:
+            for r in np.linspace(0, 3, 30):
+                scan_x = self.pose["x"] + r * np.cos(ang)
+                scan_z = self.pose["z"] + r * np.sin(ang)
+                if self.check_collision(scan_x, scan_z):
+                    ax.plot([self.pose["x"], scan_x], [self.pose["z"], scan_z], 'g-', linewidth=0.5)
+                    self.obstacle_hits.append((scan_x, scan_z))
+                    break
+
         plt.close(fig)
         return fig
 
     def render_slam_map(self):
-        # For demo, same as env
-        return self.render_env()
+        fig, ax = plt.subplots(figsize=(5, 5))
+        ax.set_title("SLAM Trajectory Map")
+        if self.trajectory:
+            x_vals, z_vals = zip(*self.trajectory)
+            ax.plot(x_vals, z_vals, 'bo-', markersize=3)
+        ax.grid(True)
+
+        if self.obstacle_hits:
+            hit_color = self.rgb_colors[self.color_index % len(self.rgb_colors)]
+            for hit in self.obstacle_hits[-20:]:
+                ax.plot(hit[0], hit[1], 'o', color=hit_color, markersize=6)
+            self.color_index += 1
+
+        plt.close(fig)
+        return fig
 
     def get_valid_moves(self):
-        x, z = self.robot_pos
-        moves = []
-        if x > 0:
-            moves.append("left")
-        if x < self.grid_size - 1:
-            moves.append("right")
-        if z > 0:
-            moves.append("down")
-        if z < self.grid_size - 1:
-            moves.append("up")
-        return moves
+        deltas = {
+            "W": (0, 1),
+            "S": (0, -1),
+            "A": (-1, 0),
+            "D": (1, 0)
+        }
+        return [
+            dir for dir, (dx, dz) in deltas.items()
+            if not self.check_collision(self.pose["x"] + dx, self.pose["z"] + dz)
+        ]
+
+    def auto_move(self):
+        valid_moves = self.get_valid_moves()
+        if not valid_moves:
+            return self.render_env(), self.render_slam_map(), "⚠️ No valid moves"
+        direction = random.choice(valid_moves)
+        return self.move_robot(direction)
 
-    def move_robot(self, direction):
-        with self.lock:
-            x, z = self.robot_pos
-            if direction == "left" and x > 0:
-                self.robot_pos[0] -= 1
-            elif direction == "right" and x < self.grid_size - 1:
-                self.robot_pos[0] += 1
-            elif direction == "up" and z < self.grid_size - 1:
-                self.robot_pos[1] += 1
-            elif direction == "down" and z > 0:
-                self.robot_pos[1] -= 1
-
-    def auto_move_step(self):
-        if self.auto_mode:
-            moves = self.get_valid_moves()
-            if moves:
-                direction = random.choice(moves)
-                self.move_robot(direction)
-                return f"Auto moved {direction}"
-            else:
-                return "No valid moves"
-        return "Auto mode off"
 
+# ---------- GRADIO UI ----------
 sim = SLAMSimulation()
 
-def toggle_auto_mode():
-    sim.auto_mode = not sim.auto_mode
-    if sim.auto_mode:
-        return "🟢 Auto Mode ON"
-    else:
-        return "⚪ Auto Mode OFF"
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 SLAM Simulation with Auto Mode + Obstacle Avoidance")
 
-def manual_move(direction):
-    sim.move_robot(direction)
-    return sim.render_env(), sim.render_slam_map(), f"Moved {direction}"
+    auto_mode_state = gr.State(value=False)
 
-def update_plots():
-    # This will be called repeatedly by gr.Timer
-    msg = sim.auto_move_step()
-    return sim.render_env(), sim.render_slam_map(), msg
+    obstacle_slider = gr.Slider(1, 20, value=10, step=1, label="Number of Obstacles")
+    direction_input = gr.Textbox(label="Type W / A / S / D and press Enter", placeholder="e.g., W")
+    status_text = gr.Textbox(label="Status", interactive=False)
 
-with gr.Blocks() as demo:
-    env_plot = gr.Plot(value=sim.render_env())
-    slam_plot = gr.Plot(value=sim.render_slam_map())
-    status = gr.Textbox(value="Ready", interactive=False)
-
-    btn_up = gr.Button("Up")
-    btn_down = gr.Button("Down")
-    btn_left = gr.Button("Left")
-    btn_right = gr.Button("Right")
-    btn_auto = gr.Button("Toggle Auto")
-
-    btn_up.click(fn=lambda: manual_move("up"), inputs=None, outputs=[env_plot, slam_plot, status])
-    btn_down.click(fn=lambda: manual_move("down"), inputs=None, outputs=[env_plot, slam_plot, status])
-    btn_left.click(fn=lambda: manual_move("left"), inputs=None, outputs=[env_plot, slam_plot, status])
-    btn_right.click(fn=lambda: manual_move("right"), inputs=None, outputs=[env_plot, slam_plot, status])
-    btn_auto.click(fn=toggle_auto_mode, inputs=None, outputs=status)
-
-    # Timer updates every 0.5 seconds
-    timer = gr.Timer(interval=0.5, 
-                     fn=update_plots, 
-                     inputs=None, 
-                     outputs=[env_plot, slam_plot, status])
+    with gr.Row():
+        with gr.Column():
+            env_plot = gr.Plot(label="Robot View")
+        with gr.Column():
+            slam_plot = gr.Plot(label="SLAM Map")
+
+    with gr.Row():
+        w_btn = gr.Button("⬆️ W")
+        a_btn = gr.Button("⬅️ A")
+        s_btn = gr.Button("⬇️ S")
+        d_btn = gr.Button("➡️ D")
+        reset_btn = gr.Button("🔄 Reset")
+        toggle_noise_btn = gr.Button("🎲 Toggle Noise")
+        auto_btn = gr.Button("🤖 Start/Stop Auto")
+
+    # Button click handlers
+    w_btn.click(fn=lambda: sim.move_robot("W"), outputs=[env_plot, slam_plot, status_text])
+    a_btn.click(fn=lambda: sim.move_robot("A"), outputs=[env_plot, slam_plot, status_text])
+    s_btn.click(fn=lambda: sim.move_robot("S"), outputs=[env_plot, slam_plot, status_text])
+    d_btn.click(fn=lambda: sim.move_robot("D"), outputs=[env_plot, slam_plot, status_text])
+
+    reset_btn.click(fn=sim.reset, inputs=[obstacle_slider], outputs=[env_plot, slam_plot, status_text])
+    toggle_noise_btn.click(fn=sim.toggle_noise, outputs=status_text)
+    direction_input.submit(fn=lambda text: sim.move_robot(text.strip().upper()), inputs=direction_input, outputs=[env_plot, slam_plot, status_text])
+
+    # Auto mode toggle logic
+    def toggle_auto_mode(state):
+        state.value = not state.value
+        if state.value:
+            return "🟢 Auto Mode: ON", state
+        else:
+            return "⚪ Auto Mode: OFF", state
+
+    auto_btn.click(fn=toggle_auto_mode, inputs=auto_mode_state, outputs=[status_text, auto_mode_state])
+
+    # Timer to do auto moves every 0.5s if enabled
+    def auto_update():
+        if auto_mode_state.value:
+            return sim.auto_move()
+        else:
+            # just update current plots and status without moving
+            return sim.render_env(), sim.render_slam_map(), status_text.value
+
+    timer = gr.Timer(every=0.5, fn=auto_update, outputs=[env_plot, slam_plot, status_text])
 
 demo.launch()