msintui/Intelligent_PID

1import os2import json3import networkx as nx4import matplotlib.pyplot as plt5from pathlib import Path6import logging7import traceback8from storage import StorageFactory9 10logger = logging.getLogger(__name__)11 12 13def construct_graph_network(data: dict, validation_results_path: str, results_dir: str, storage=None):14    """Construct network graph from aggregated detection data"""15    try:16        # Use provided storage or get a new one17        if storage is None:18            storage = StorageFactory.get_storage()19 20        # Create graph21        G = nx.Graph()22        pos = {}  # For node positions23 24        # Add nodes from the aggregated data25        for node in data.get('nodes', []):26            node_id = node['id']27            node_type = node['type']28 29            # Calculate position based on node type30            if node_type == 'connection_point':31                pos[node_id] = (node['coords']['x'], node['coords']['y'])32            else:  # symbol or text33                bbox = node['bbox']34                pos[node_id] = (35                    (bbox['xmin'] + bbox['xmax']) / 2,36                    (bbox['ymin'] + bbox['ymax']) / 237                )38 39            # Add node with all its properties40            G.add_node(node_id, **node)41 42        # Add edges from the aggregated data43        for edge in data.get('edges', []):44            G.add_edge(45                edge['source'],46                edge['target'],47                **edge.get('properties', {})48            )49 50        # Create visualization51        plt.figure(figsize=(20, 20))52 53        # Draw nodes with different colors based on type54        node_colors = []55        for node in G.nodes():56            node_type = G.nodes[node]['type']57            if node_type == 'symbol':58                node_colors.append('lightblue')59            elif node_type == 'text':60                node_colors.append('lightgreen')61            else:  # connection_point62                node_colors.append('lightgray')63 64        nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=500)65        nx.draw_networkx_edges(G, pos, edge_color='gray', width=1)66 67        # Add labels68        labels = {}69        for node in G.nodes():70            node_data = G.nodes[node]71            if node_data['type'] == 'symbol':72                labels[node] = f"S:{node_data.get('properties', {}).get('class', '')}"73            elif node_data['type'] == 'text':74                content = node_data.get('content', '')75                labels[node] = f"T:{content[:10]}..." if len(content) > 10 else f"T:{content}"76            else:77                labels[node] = f"C:{node_data['properties'].get('point_type', '')}"78 79        nx.draw_networkx_labels(G, pos, labels, font_size=8)80 81        plt.title("P&ID Knowledge Graph")82        plt.axis('off')83 84        # Save the visualization85        graph_image_path = os.path.join(results_dir, f"{Path(data.get('image_path', 'graph')).stem}_graph.png")86        plt.savefig(graph_image_path, bbox_inches='tight', dpi=300)87        plt.close()88 89        # Save graph data as JSON for future use90        graph_json_path = os.path.join(results_dir, f"{Path(data.get('image_path', 'graph')).stem}_graph_data.json")91        with open(graph_json_path, 'w') as f:92            json.dump(nx.node_link_data(G), f, indent=2)93 94        return G, pos, plt.gcf()95 96    except Exception as e:97        logger.error(f"Error in construct_graph_network: {str(e)}")98        traceback.print_exc()99        return None, None, None100 101 102if __name__ == "__main__":103    # Test code104    test_data_path = "results/test_aggregated.json"105    if os.path.exists(test_data_path):106        with open(test_data_path, 'r') as f:107            test_data = json.load(f)108 109        G, pos, fig = construct_graph_network(110            test_data,111            "results/validation.json",112            "results"113        )114        if fig:115            plt.show()