Update app.py

app.py

@@ -1,145 +1,201 @@
 import gradio as gr
 import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import matplotlib.image as mpimg
 import numpy as np
+import random
 import threading
 import time
-import random
-from io import BytesIO
-
-class SLAMSimulation:
-    def __init__(self):
-        self.reset(10)
-
-    def reset(self, obstacle_count=10):
-        self.robot_pos = [0, 0]
-        self.robot_path = [tuple(self.robot_pos)]
-        self.map = set()
-        self.obstacles = []
-        self.auto_mode = False
-
-        for _ in range(obstacle_count):
-            self.obstacles.append({
-                "x": random.randint(-10, 10),
-                "z": random.randint(-10, 10),
-                "radius": random.uniform(0.5, 1.5),
-            })
-
-        return self.render_env(), self.render_slam_map(), f"🔄 Simulation reset with {obstacle_count} obstacles"
-
-    def move(self, direction):
-        dx, dz = 0, 0
-        if direction == "up":
-            dz = 1
-        elif direction == "down":
-            dz = -1
-        elif direction == "left":
-            dx = -1
-        elif direction == "right":
-            dx = 1
-        self.robot_pos[0] += dx
-        self.robot_pos[1] += dz
-        self.robot_path.append(tuple(self.robot_pos))
-        self.update_map()
-        return self.render_env(), self.render_slam_map(), f"📍 Moved {direction}"
-
-    def auto_move(self):
-        direction = random.choice(["up", "down", "left", "right"])
-        return self.move(direction)
-
-    def update_map(self):
-        x, z = self.robot_pos
-        for dx in range(-2, 3):
-            for dz in range(-2, 3):
-                self.map.add((x + dx, z + dz))
-
-    def render_env(self):
-        fig, ax = plt.subplots()
-        ax.set_xlim(-15, 15)
-        ax.set_ylim(-15, 15)
-        ax.set_title("Environment")
-
-        for obs in self.obstacles:
-            circle = plt.Circle((obs["x"], obs["z"]), obs["radius"],
-                                color="gray", alpha=0.6, edgecolor="black")
-            ax.add_patch(circle)
-
-        path = np.array(self.robot_path)
-        ax.plot(path[:, 0], path[:, 1], linestyle='--', color='blue')
-        ax.plot(self.robot_pos[0], self.robot_pos[1], 'ro', label='Robot')
-        ax.legend()
-
-        return self.fig_to_image(fig)
-
-    def render_slam_map(self):
-        fig, ax = plt.subplots()
-        ax.set_xlim(-15, 15)
-        ax.set_ylim(-15, 15)
-        ax.set_title("SLAM Map")
-
-        for (x, z) in self.map:
-            ax.plot(x, z, 'ks', markersize=4)
-
-        ax.plot(self.robot_pos[0], self.robot_pos[1], 'ro', label='Robot')
-        ax.legend()
-
-        return self.fig_to_image(fig)
-
-    def fig_to_image(self, fig):
-        buf = BytesIO()
-        fig.savefig(buf, format='png')
-        buf.seek(0)
-        plt.close(fig)
-        return buf
-
-sim = SLAMSimulation()
-auto_mode_state = {"value": False}
-
-def manual_control(direction):
-    return sim.move(direction)
-
-def auto_toggle():
-    auto_mode_state["value"] = not auto_mode_state["value"]
-    return f"🤖 Auto mode {'ON' if auto_mode_state['value'] else 'OFF'}"
-
-def reset_sim():
-    return sim.reset(10)
-
-def background_loop(update_fn, delay=0.5):
-    def loop():
-        while True:
-            if auto_mode_state["value"]:
-                env, slam, status = sim.auto_move()
-                update_fn(env, slam, status)
-            time.sleep(delay)
-    threading.Thread(target=loop, daemon=True).start()
 
+# Global State
+pose = {"x": 0, "z": 0, "angle": 0}
+trajectory = [(0, 0)]
+obstacle_hits = []
+color_index = 0
+rgb_colors = ['#FF5733', '#33FF57', '#3357FF']  # More distinguishable colors
+noise_enabled = True
+auto_mode = False
+obstacles = []
+
+# Thread lock for safety in auto mode
+lock = threading.Lock()
+
+# --- Obstacle Generation ---
+def generate_obstacles(count=10):
+    return [{
+        "x": random.uniform(-8, 8),
+        "z": random.uniform(-8, 8),
+        "radius": random.uniform(0.5, 1.2)
+    } for _ in range(count)]
+
+obstacles = generate_obstacles(10)
+
+# --- Collision Check ---
+def check_collision(x, z):
+    for obs in obstacles:
+        dist = np.sqrt((obs["x"] - x)**2 + (obs["z"] - z)**2)
+        if dist <= obs["radius"] + 0.2:
+            return True
+    return False
+
+# --- Movement ---
+def move_robot(direction):
+    global pose, trajectory
+    step = 1
+    direction = direction.upper()
+    dx, dz, angle = 0, 0, 0
+
+    if direction == "W":
+        dz, angle = step, 90
+    elif direction == "S":
+        dz, angle = -step, -90
+    elif direction == "A":
+        dx, angle = -step, 180
+    elif direction == "D":
+        dx, angle = step, 0
+    else:
+        return render_env(), render_slam_map(), "❌ Invalid Key"
+
+    new_x, new_z = pose["x"] + dx, pose["z"] + dz
+
+    if check_collision(new_x, new_z):
+        return render_env(), render_slam_map(), "🚫 Collision detected!"
+
+    pose["x"], pose["z"], pose["angle"] = new_x, new_z, angle
+    noisy_x = new_x + random.uniform(-0.1, 0.1) if noise_enabled else new_x
+    noisy_z = new_z + random.uniform(-0.1, 0.1) if noise_enabled else new_z
+    trajectory.append((noisy_x, noisy_z))
+
+    return render_env(), render_slam_map(), f"Moved {direction}"
+
+# --- Rendering ---
+def render_env():
+    global obstacle_hits
+    fig, ax = plt.subplots(figsize=(5, 5))
+    ax.set_xlim(-10, 10)
+    ax.set_ylim(-10, 10)
+    ax.set_title("SLAM Environment View")
+
+    # Background
+    try:
+        bg = mpimg.imread("map.png")
+        ax.imshow(bg, extent=(-10, 10, -10, 10), alpha=0.2)
+    except:
+        pass
+
+    # Obstacles
+    for obs in obstacles:
+        circ = plt.Circle((obs["x"], obs["z"]), obs["radius"], color="gray", alpha=0.6)
+        ax.add_patch(circ)
+
+    # Robot
+    ax.plot(pose["x"], pose["z"], 'ro', markersize=8)
+
+    # LIDAR Simulation
+    obstacle_hits.clear()
+    for ang in np.linspace(0, 2*np.pi, 24):
+        for r in np.linspace(0, 3, 30):
+            scan_x = pose["x"] + r * np.cos(ang)
+            scan_z = pose["z"] + r * np.sin(ang)
+            if check_collision(scan_x, scan_z):
+                ax.plot([pose["x"], scan_x], [pose["z"], scan_z], 'g-', linewidth=0.5)
+                obstacle_hits.append((scan_x, scan_z))
+                break
+
+    plt.close(fig)
+    return fig
+
+def render_slam_map():
+    global color_index
+    fig, ax = plt.subplots(figsize=(5, 5))
+    ax.set_title("SLAM Trajectory Map")
+    ax.grid(True)
+
+    x_vals, z_vals = zip(*trajectory)
+    ax.plot(x_vals, z_vals, 'bo-', markersize=3)
+
+    if obstacle_hits:
+        current_color = rgb_colors[color_index % len(rgb_colors)]
+        for hit in obstacle_hits[-20:]:
+            ax.plot(hit[0], hit[1], 'o', color=current_color, markersize=6)
+        color_index += 1
+
+    plt.close(fig)
+    return fig
+
+# --- Handlers ---
+def toggle_noise():
+    global noise_enabled
+    noise_enabled = not noise_enabled
+    return "Noise: ON" if noise_enabled else "Noise: OFF"
+
+def reset_sim(count):
+    global pose, trajectory, obstacles, obstacle_hits, color_index
+    pose = {"x": 0, "z": 0, "angle": 0}
+    trajectory = [(0, 0)]
+    obstacle_hits.clear()
+    color_index = 0
+    obstacles = generate_obstacles(int(count))
+    return render_env(), render_slam_map(), f"Simulation Reset with {count} obstacles"
+
+def handle_text_input(direction):
+    return move_robot(direction.strip().upper())
+
+def auto_movement(update_callback):
+    global auto_mode
+    directions = ['W', 'A', 'S', 'D']
+    while auto_mode:
+        with lock:
+            direction = random.choice(directions)
+            env, slam, msg = move_robot(direction)
+            update_callback(env, slam, msg)
+        time.sleep(1)
+
+def toggle_auto_mode(env_plot, slam_plot, status_text):
+    global auto_mode
+    auto_mode = not auto_mode
+
+    if auto_mode:
+        def update_ui(e, s, t):
+            env_plot.update(value=e)
+            slam_plot.update(value=s)
+            status_text.update(value=t)
+
+        thread = threading.Thread(target=auto_movement, args=(update_ui,), daemon=True)
+        thread.start()
+        return "🟢 Auto Mode: ON"
+    else:
+        return "⚪ Auto Mode: OFF"
+
+# --- Gradio Interface ---
 with gr.Blocks() as demo:
-    env_plot = gr.Image(label="Environment")
-    slam_plot = gr.Image(label="SLAM Map")
-    status_text = gr.Textbox(label="Status")
+    gr.Markdown("## 🤖 SLAM Simulation: Noise, Collision, and Auto Mode")
+
+    obstacle_slider = gr.Slider(1, 20, value=10, step=1, label="Obstacle Count")
+    direction_input = gr.Textbox(label="Manual Control (W/A/S/D)", placeholder="Type a direction...")
+    status_text = gr.Textbox(label="Status", interactive=False)
 
     with gr.Row():
-        btn_up = gr.Button("⬆️ Up")
-        btn_down = gr.Button("⬇️ Down")
-        btn_left = gr.Button("⬅️ Left")
-        btn_right = gr.Button("➡️ Right")
+        env_plot = gr.Plot(label="Environment View")
+        slam_plot = gr.Plot(label="SLAM Map")
+
     with gr.Row():
-        btn_auto = gr.Button("🤖 Toggle Auto")
-        btn_reset = gr.Button("🔄 Reset")
-
-    def update_ui(env, slam, status):
-        env_plot.update(env)
-        slam_plot.update(slam)
-        status_text.update(status)
-
-    btn_up.click(fn=lambda: sim.move("up"), outputs=[env_plot, slam_plot, status_text])
-    btn_down.click(fn=lambda: sim.move("down"), outputs=[env_plot, slam_plot, status_text])
-    btn_left.click(fn=lambda: sim.move("left"), outputs=[env_plot, slam_plot, status_text])
-    btn_right.click(fn=lambda: sim.move("right"), outputs=[env_plot, slam_plot, status_text])
-    btn_reset.click(fn=reset_sim, outputs=[env_plot, slam_plot, status_text])
-    btn_auto.click(fn=auto_toggle, outputs=[status_text])
-
-    # Start background loop for auto mode
-    background_loop(update_ui)
+        w = gr.Button("⬆️ W")
+        a = gr.Button("⬅️ A")
+        s = gr.Button("⬇️ S")
+        d = gr.Button("➡️ D")
+        reset = gr.Button("🔄 Reset")
+        toggle = gr.Button("🔀 Toggle Noise")
+        auto = gr.Button("🤖 Toggle Auto")
+
+    w.click(fn=lambda: move_robot("W"), outputs=[env_plot, slam_plot, status_text])
+    a.click(fn=lambda: move_robot("A"), outputs=[env_plot, slam_plot, status_text])
+    s.click(fn=lambda: move_robot("S"), outputs=[env_plot, slam_plot, status_text])
+    d.click(fn=lambda: move_robot("D"), outputs=[env_plot, slam_plot, status_text])
+
+    reset.click(fn=reset_sim, inputs=[obstacle_slider], outputs=[env_plot, slam_plot, status_text])
+    toggle.click(fn=lambda: (None, None, toggle_noise()), outputs=[env_plot, slam_plot, status_text])
+    auto.click(fn=toggle_auto_mode, inputs=[env_plot, slam_plot, status_text], outputs=auto)
+    direction_input.submit(fn=handle_text_input, inputs=direction_input, outputs=[env_plot, slam_plot, status_text])
 
 demo.launch()