Update app.py

app.py

@@ -1,7 +1,9 @@
 import streamlit as st
 import os
-from PIL import Image
 import random
+from PIL import Image
+import networkx as nx
+import matplotlib.pyplot as plt
 
 # Function to calculate aspect ratio
 def calculate_aspect_ratio(image_path):
@@ -36,20 +38,43 @@ def generate_dungeon_map(image_files):
         else:
             dungeon["Rooms"].append(img_file)
 
-    # Add logic to connect rooms, hallways, and doors
-    layout = {}
-    if dungeon["Hallways"]:
-        for i, hallway in enumerate(dungeon["Hallways"]):
-            layout[f"Hallway {i+1}"] = {"type": "Hallway", "connected_to": []}
-            if dungeon["Rooms"]:
-                room = dungeon["Rooms"].pop(0)
-                layout[f"Hallway {i+1}"]["connected_to"].append(room)
-            if dungeon["Doors"]:
-                door = dungeon["Doors"].pop(0)
-                layout[f"Hallway {i+1}"]["connected_to"].append(door)
+    # Generate connections between elements
+    layout = nx.Graph()
+    for i, hallway in enumerate(dungeon["Hallways"]):
+        layout.add_node(f"Hallway {i+1}", type="Hallway")
+        if dungeon["Rooms"]:
+            room = dungeon["Rooms"].pop(0)
+            layout.add_node(room, type="Room")
+            layout.add_edge(f"Hallway {i+1}", room)
+
+        if dungeon["Doors"]:
+            door = dungeon["Doors"].pop(0)
+            layout.add_node(door, type="Door")
+            layout.add_edge(f"Hallway {i+1}", door)
 
     return layout, dungeon
 
+# Function to visualize the dungeon layout
+def visualize_dungeon(layout, output_path="dungeon_tree.png"):
+    plt.figure(figsize=(12, 8))
+    pos = nx.spring_layout(layout)  # Positions for nodes
+    types = nx.get_node_attributes(layout, 'type')
+
+    # Draw nodes with different colors
+    color_map = {
+        "Hallway": "blue",
+        "Door": "green",
+        "Room": "orange"
+    }
+    colors = [color_map[types[node]] for node in layout.nodes]
+
+    nx.draw(
+        layout, pos, with_labels=True, node_color=colors, node_size=2000,
+        font_size=10, font_color="white", font_weight="bold", edge_color="gray"
+    )
+    plt.savefig(output_path)
+    plt.close()
+
 # Streamlit App
 st.title("Dynamic Dungeon Map Generator")
 st.write("Automatically generates dungeon maps by analyzing PNG images in a directory.")
@@ -62,28 +87,33 @@ if os.path.exists(map_dir):
     image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
 
     if image_files:
-        st.subheader("Classified Rooms")
-        layout, dungeon = generate_dungeon_map(image_files)
+        st.subheader("Generate Dungeon Map")
 
-        st.write("### Hallways")
-        st.write(", ".join(dungeon["Hallways"]) if dungeon["Hallways"] else "No hallways found.")
+        # Add emoji buttons
+        if st.button("🔄 Regenerate Dungeon"):
+            layout, dungeon = generate_dungeon_map(image_files)
+            visualize_dungeon(layout)
 
-        st.write("### Doors")
-        st.write(", ".join(dungeon["Doors"]) if dungeon["Doors"] else "No doors found.")
+            st.image("dungeon_tree.png", caption="Dungeon Layout Tree")
 
-        st.write("### Rooms")
-        st.write(", ".join(dungeon["Rooms"]) if dungeon["Rooms"] else "No rooms found.")
+            # Display classified rooms
+            st.subheader("Classified Rooms")
+            st.write("### Hallways")
+            st.write(", ".join(dungeon["Hallways"]) if dungeon["Hallways"] else "No hallways found.")
 
-        st.subheader("Dungeon Layout")
-        for key, value in layout.items():
-            st.write(f"{key}: {value}")
+            st.write("### Doors")
+            st.write(", ".join(dungeon["Doors"]) if dungeon["Doors"] else "No doors found.")
 
+            st.write("### Rooms")
+            st.write(", ".join(dungeon["Rooms"]) if dungeon["Rooms"] else "No rooms found.")
+        else:
+            st.write("Click the button above to generate a dungeon map.")
     else:
         st.write("No PNG files found in the specified directory.")
 else:
     st.write("The specified directory does not exist. Please check the path.")
 
-# File upload option
+# Sidebar for file upload
 st.sidebar.title("Options")
 uploaded_file = st.sidebar.file_uploader("Upload a PNG file", type="png")
 
@@ -92,7 +122,3 @@ if uploaded_file:
     with open(os.path.join(map_dir, uploaded_file.name), "wb") as f:
         f.write(uploaded_file.getbuffer())
     st.sidebar.success("File uploaded successfully! Refresh the app to include it in the dungeon map.")
-
-# Placeholder for visualization
-st.subheader("Visual Representation")
-st.write("A feature to visualize the map will be added soon.")