Create app.py

app.py

@@ -0,0 +1,448 @@
+import os
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+import networkx as nx
+import gradio as gr
+from matplotlib.colors import LinearSegmentedColormap
+import matplotlib.patches as mpatches
+
+# Check if GPU is available
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+class EnhancedMindMapGenerator:
+    def __init__(self):
+        self.graph = nx.DiGraph()  # Using DiGraph for directed edges
+        self.node_positions = {}
+        self.node_colors = {}
+        self.edge_colors = {}
+        self.node_sizes = {}
+        self.node_depth = {}
+        self.levels = {}
+
+    def reset(self):
+        self.graph = nx.DiGraph()
+        self.node_positions = {}
+        self.node_colors = {}
+        self.edge_colors = {}
+        self.node_sizes = {}
+        self.node_depth = {}
+        self.levels = {}
+        return "Mind map reset successfully"
+
+    def parse_input(self, text):
+        """Parse the input text into nodes and relationships"""
+        lines = text.strip().split('\n')
+        root_node = None
+        parent_map = {}  # Track parent nodes based on indent level
+        current_indent_level = -1
+        current_parent = None
+
+        # First pass: Build hierarchy based on indentation
+        for line in lines:
+            original_line = line
+            line = line.strip()
+            if not line or '->' in line:
+                continue  # Skip empty lines and relationship lines for now
+
+            # Calculate indent level
+            indent_level = len(original_line) - len(original_line.lstrip())
+
+            if root_node is None:
+                # This is the root node
+                root_node = line
+                self.add_node(root_node, is_root=True, depth=0)
+                parent_map[0] = root_node
+                current_indent_level = indent_level
+                current_parent = root_node
+                self.levels[0] = [root_node]
+            else:
+                # Handle indentation to determine parent-child relationships
+                if indent_level > current_indent_level:
+                    # This is a child of the previous node
+                    parent_map[indent_level] = current_parent
+
+                parent = None
+                if indent_level in parent_map:
+                    parent = parent_map[indent_level]
+                    # If this is a new indent level, set the parent to the previous node
+                    if indent_level > current_indent_level:
+                        parent = current_parent
+                else:
+                    # Find the closest parent based on indent
+                    closest_indent = max([i for i in parent_map.keys() if i < indent_level], default=0)
+                    parent = parent_map[closest_indent]
+
+                # Calculate depth based on parent's depth
+                parent_depth = self.node_depth.get(parent, 0)
+                current_depth = parent_depth + 1
+
+                # Add node and edge
+                self.add_node(line, depth=current_depth)
+                self.add_edge(parent, line, "hierarchy")
+
+                # Add to level structure
+                if current_depth not in self.levels:
+                    self.levels[current_depth] = []
+                self.levels[current_depth].append(line)
+
+                # Update tracking variables
+                current_indent_level = indent_level
+                current_parent = line
+                parent_map[indent_level] = line
+
+        # Second pass: Process explicit relationships (->)
+        for line in lines:
+            line = line.strip()
+            if '->' in line:
+                parts = line.split('->')
+                if len(parts) == 2:
+                    source = parts[0].strip()
+                    target = parts[1].strip()
+                    self.add_edge(source, target, "relationship")
+
+        return f"Parsed mind map with root: {root_node}"
+
+    def add_node(self, node_name, is_root=False, depth=0):
+        """Add a node to the graph"""
+        if node_name not in self.graph.nodes:
+            self.graph.add_node(node_name)
+            self.node_depth[node_name] = depth
+
+            # Set color based on depth
+            if is_root:
+                self.node_colors[node_name] = '#FF5733'  # Root is red
+                self.node_sizes[node_name] = 2500
+            else:
+                # Use a color scheme based on depth
+                color_map = {
+                    1: '#3498DB',  # Blue
+                    2: '#F39C12',  # Orange
+                    3: '#2ECC71',  # Green
+                    4: '#9B59B6',  # Purple
+                    5: '#E74C3C',  # Red
+                }
+                self.node_colors[node_name] = color_map.get(depth % len(color_map), '#95A5A6')  # Gray as default
+                self.node_sizes[node_name] = 2000 - (depth * 200)  # Size decreases with depth
+
+    def add_edge(self, source, target, edge_type="hierarchy"):
+        """Add an edge between two nodes"""
+        if source not in self.graph.nodes:
+            self.add_node(source)
+        if target not in self.graph.nodes:
+            self.add_node(target)
+
+        if not self.graph.has_edge(source, target):
+            self.graph.add_edge(source, target)
+
+            # Color edges based on type
+            if edge_type == "relationship":
+                self.edge_colors[(source, target)] = 'green'
+            else:
+                self.edge_colors[(source, target)] = 'gray'
+
+    def calculate_hierarchical_layout(self):
+        """Calculate a hierarchical layout based on node depth"""
+        # Use hierarchical layout with depth levels
+        pos = {}
+        max_nodes_per_level = max([len(nodes) for nodes in self.levels.values()])
+
+        for level, nodes in self.levels.items():
+            y = -level * 2  # Vertical position based on level
+
+            # Center the nodes at each level
+            width = max(max_nodes_per_level, len(nodes))
+            for i, node in enumerate(nodes):
+                x = (i - (len(nodes) - 1) / 2) * 3  # Horizontal spacing
+                pos[node] = np.array([x, y])
+
+        return pos
+
+    def optimize_layout(self):
+        """Use GPU-accelerated optimization for node layout (if available)"""
+        # First set initial positions using hierarchical layout
+        initial_pos = self.calculate_hierarchical_layout()
+        self.node_positions = initial_pos
+
+        if device.type == "cuda":
+            print("Optimizing layout using GPU...")
+            # Implement GPU optimization if needed
+            nodes = list(self.graph.nodes)
+            positions = torch.tensor([self.node_positions[node] for node in nodes], device=device)
+
+            # Simple force-directed algorithm using PyTorch (maintains hierarchical structure)
+            for _ in range(50):
+                # Calculate attractive forces (edges)
+                attractive_force = torch.zeros_like(positions)
+                for u, v in self.graph.edges:
+                    u_idx = nodes.index(u)
+                    v_idx = nodes.index(v)
+                    direction = positions[v_idx] - positions[u_idx]
+                    distance = torch.norm(direction) + 1e-5
+                    force = direction * torch.log(distance / 2) * 0.1
+                    attractive_force[u_idx] += force
+                    attractive_force[v_idx] -= force
+
+                # Calculate repulsive forces (nodes at same level)
+                repulsive_force = torch.zeros_like(positions)
+                for level_nodes in self.levels.values():
+                    level_indices = [nodes.index(node) for node in level_nodes if node in nodes]
+                    for i_idx, i in enumerate(level_indices):
+                        for j in level_indices[i_idx+1:]:
+                            direction = positions[j] - positions[i]
+                            distance = torch.norm(direction) + 1e-5
+                            if distance < 3.0:  # Only apply repulsion when nodes are close
+                                force = direction / (distance ** 2) * 0.5
+                                repulsive_force[i] -= force
+                                repulsive_force[j] += force
+
+                # Update positions but maintain y-coordinate (level)
+                new_pos = positions + (attractive_force + repulsive_force) * 0.1
+
+                # Preserve y-coordinates to maintain hierarchical layout
+                for i, node in enumerate(nodes):
+                    level = self.node_depth[node]
+                    new_pos[i, 1] = positions[i, 1]  # Keep original y-coordinate
+
+                positions = new_pos
+
+            # Copy back to CPU and update positions
+            positions_cpu = positions.cpu().numpy()
+            for i, node in enumerate(nodes):
+                self.node_positions[node] = positions_cpu[i]
+
+            return "Layout optimized using GPU acceleration while preserving hierarchy"
+        else:
+            # CPU-based optimization
+            # Adjust positions to prevent overlaps while maintaining hierarchy
+            pos = nx.spring_layout(
+                self.graph,
+                pos=self.node_positions,
+                fixed=None,  # Don't fix positions
+                k=1.5,  # Increase node separation
+                iterations=50,
+                weight=None
+            )
+
+            # Preserve y-coordinates to maintain hierarchical layout
+            for node in self.graph.nodes:
+                pos[node][1] = self.node_positions[node][1]  # Keep original y-coordinate
+
+            self.node_positions = pos
+            return "Layout optimized using CPU while preserving hierarchy"
+
+    def visualize(self):
+        """Generate a visualization of the mind map"""
+        if not self.graph.nodes:
+            return None
+
+        plt.figure(figsize=(16, 12), dpi=100)
+
+        # Use calculated positions from hierarchical layout or optimization
+        pos = self.node_positions
+
+        # Create a legend for depth levels
+        depth_colors = {}
+        for node, depth in self.node_depth.items():
+            if depth not in depth_colors:
+                depth_colors[depth] = self.node_colors[node]
+
+        # Draw edges with curved arrows for relationships
+        for edge in self.graph.edges:
+            edge_color = self.edge_colors.get(edge, 'gray')
+
+            # Use curved edges for explicit relationships, straight for hierarchy
+            if edge_color == 'green':
+                nx.draw_networkx_edges(
+                    self.graph,
+                    pos,
+                    edgelist=[edge],
+                    width=2.5,
+                    edge_color=edge_color,
+                    alpha=0.8,
+                    arrows=True,
+                    arrowsize=15,
+                    connectionstyle="arc3,rad=0.3"
+                )
+            else:
+                nx.draw_networkx_edges(
+                    self.graph,
+                    pos,
+                    edgelist=[edge],
+                    width=1.5,
+                    edge_color=edge_color,
+                    alpha=0.7,
+                    arrows=True,
+                    arrowsize=12
+                )
+
+        # Draw nodes with depth-based colors
+        for node in self.graph.nodes:
+            nx.draw_networkx_nodes(
+                self.graph,
+                pos,
+                nodelist=[node],
+                node_color=self.node_colors.get(node, 'blue'),
+                node_size=self.node_sizes.get(node, 1000),
+                alpha=0.9,
+                edgecolors='black',
+                linewidths=1
+            )
+
+        # Draw labels with white background for better readability
+        label_pos = {node: (pos[node][0], pos[node][1]) for node in self.graph.nodes}
+        nx.draw_networkx_labels(
+            self.graph,
+            label_pos,
+            font_size=10,
+            font_family='sans-serif',
+            font_weight='bold',
+            bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', boxstyle='round,pad=0.3')
+        )
+
+        # Add a legend
+        legend_elements = [
+            mpatches.Patch(color='#FF5733', label='Root'),
+            mpatches.Patch(color='#3498DB', label='Level 1'),
+            mpatches.Patch(color='#F39C12', label='Level 2'),
+            mpatches.Patch(color='#2ECC71', label='Level 3'),
+            mpatches.Patch(color='#9B59B6', label='Level 4+'),
+            mpatches.Patch(color='green', label='Explicit Relationship'),
+            mpatches.Patch(color='gray', label='Hierarchical Relationship')
+        ]
+        plt.legend(handles=legend_elements, loc='upper right')
+
+        plt.title("Mind Map Visualization", fontsize=16, fontweight='bold')
+        plt.axis('off')
+        plt.tight_layout()
+
+        # Save to a temporary file
+        temp_path = "mindmap_output.png"
+        plt.savefig(temp_path, format="png", dpi=300, bbox_inches='tight', facecolor='white')
+        plt.close()
+
+        return temp_path
+
+# Create the Gradio interface
+def create_mind_map(input_text, optimization):
+    """Create a mind map from input text"""
+    generator = EnhancedMindMapGenerator()
+    message = generator.parse_input(input_text)
+    print(message)
+
+    if optimization:
+        message = generator.optimize_layout()
+        print(message)
+
+    image_path = generator.visualize()
+    return image_path
+
+# For Colab, use this function to create and launch the demo
+def create_and_launch():
+    """Create and launch the Gradio demo"""
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# Enhanced Mind Map Generator")
+        gr.Markdown("Enter your mind map structure below. Use indentation to represent hierarchy or use -> for direct relationships.")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                input_text = gr.Textbox(
+                    placeholder="Project Name\n  Task 1\n    Subtask 1.1\n    Subtask 1.2\n  Task 2\nTask 1 -> Task 2",
+                    label="Mind Map Structure",
+                    lines=15
+                )
+
+                with gr.Row():
+                    optimization = gr.Checkbox(label="Use Layout Optimization", value=True)
+                    generate_btn = gr.Button("Generate Mind Map", variant="primary")
+
+                gr.Markdown("### Format Guide:")
+                gr.Markdown("""
+                - Use indentation (spaces) to create parent-child relationships
+                - Each level of indentation creates a new depth level
+                - Use '-> ' to create explicit connections (e.g., 'NodeA -> NodeB')
+                - The first non-indented line becomes the root node
+                """)
+
+            with gr.Column(scale=3):
+                output_image = gr.Image(label="Generated Mind Map", type="filepath")
+
+        generate_btn.click(fn=create_mind_map, inputs=[input_text, optimization], outputs=output_image)
+
+        # Add examples
+        example_input1 = """Software Project
+  Planning
+    Requirements Gathering
+    Project Timeline
+    Resource Allocation
+  Development
+    Frontend
+      UI Design
+      React Components
+    Backend
+      API Development
+      Database Setup
+    Testing
+      Unit Tests
+      Integration Tests
+  Deployment
+    CI/CD Pipeline
+    Production Release
+Planning -> Development
+Development -> Testing
+Testing -> Deployment"""
+
+        example_input2 = """Business Strategy
+  Market Analysis
+    Customer Demographics
+    Competitor Research
+    Market Trends
+  Internal Assessment
+    SWOT Analysis
+    Resource Inventory
+  Strategic Goals
+    Short-term Objectives
+    Long-term Vision
+  Implementation
+    Action Plans
+Market Analysis -> Strategic Goals
+Internal Assessment -> Strategic Goals
+Strategic Goals -> Implementation"""
+
+        gr.Examples(
+            examples=[[example_input1, True], [example_input2, True]],
+            inputs=[input_text, optimization],
+            outputs=output_image,
+            fn=create_mind_map,
+            cache_examples=True,
+        )
+
+    # Launch with sharing enabled for Colab
+    demo.launch(share=True, debug=True)
+    return demo
+
+# Main execution
+def run_in_colab():
+    # Install necessary packages
+    print("Installing required packages...")
+    try:
+        import gradio
+        import networkx
+    except ImportError:
+        !pip install gradio networkx matplotlib
+        print("Packages installed!")
+
+    # Create and launch the demo
+    print("Launching the Enhanced Mind Map Generator...")
+    create_and_launch()
+
+# For Google Colab, use this
+try:
+    import google.colab
+    print("Running in Google Colab environment")
+    run_in_colab()
+except:
+    print("Running in local environment")
+    # If not in Colab, just create and launch
+    create_and_launch()