msintui/Intelligent_PID

1import json2import networkx as nx3import numpy as np4import matplotlib.pyplot as plt5import traceback6import uuid7 8 9def create_connected_graph(input_data):10    """Create a connected graph from the input data"""11    try:12        # Validate input data structure13        if not isinstance(input_data, dict):14            raise ValueError("Invalid input data format")15 16        # Check for required keys in new format17        required_keys = ['symbols', 'texts', 'lines', 'nodes', 'edges']18        if not all(key in input_data for key in required_keys):19            raise ValueError(f"Missing required keys in input data. Expected: {required_keys}")20 21        # Create graph22        G = nx.Graph()23 24        # Track positions for layout25        pos = {}26 27        # Add symbol nodes28        for symbol in input_data['symbols']:29            bbox = symbol.get('bbox', [])30            symbol_id = symbol.get('id', str(uuid.uuid4()))31 32            if bbox:33                # Calculate center position34                center_x = (bbox['xmin'] + bbox['xmax']) / 235                center_y = (bbox['ymin'] + bbox['ymax']) / 236                pos[symbol_id] = (center_x, center_y)37 38                G.add_node(39                    symbol_id,40                    type='symbol',41                    class_name=symbol.get('class', ''),42                    bbox=bbox,43                    confidence=symbol.get('confidence', 0.0)44                )45 46        # Add text nodes47        for text in input_data['texts']:48            bbox = text.get('bbox', [])49            text_id = text.get('id', str(uuid.uuid4()))50 51            if bbox:52                center_x = (bbox['xmin'] + bbox['xmax']) / 253                center_y = (bbox['ymin'] + bbox['ymax']) / 254                pos[text_id] = (center_x, center_y)55 56                G.add_node(57                    text_id,58                    type='text',59                    text=text.get('text', ''),60                    bbox=bbox,61                    confidence=text.get('confidence', 0.0)62                )63 64        # Add edges from the edges list65        for edge in input_data['edges']:66            source = edge.get('source')67            target = edge.get('target')68            if source and target and source in G and target in G:69                G.add_edge(70                    source,71                    target,72                    type=edge.get('type', 'connection'),73                    properties=edge.get('properties', {})74                )75 76        # Create visualization77        plt.figure(figsize=(20, 20))78 79        # Draw nodes with fixed positions80        nx.draw_networkx_nodes(G, pos,81                               node_color=['lightblue' if G.nodes[node]['type'] == 'symbol' else 'lightgreen' for node82                                           in G.nodes()],83                               node_size=500)84 85        # Draw edges86        nx.draw_networkx_edges(G, pos, edge_color='gray', width=1)87 88        # Add labels89        labels = {}90        for node in G.nodes():91            node_data = G.nodes[node]92            if node_data['type'] == 'symbol':93                labels[node] = f"S:{node_data['class_name']}"94            else:95                text = node_data.get('text', '')96                labels[node] = f"T:{text[:10]}..." if len(text) > 10 else f"T:{text}"97 98        nx.draw_networkx_labels(G, pos, labels, font_size=8)99 100        plt.title("P&ID Network Graph")101        plt.axis('off')102 103        return G, pos, plt.gcf()104 105    except Exception as e:106        print(f"Error in create_connected_graph: {str(e)}")107        traceback.print_exc()108        return None, None, None109 110 111if __name__ == "__main__":112    # Test code113    with open('results/0_aggregated.json') as f:114        data = json.load(f)115 116    G, pos, fig = create_connected_graph(data)117    if fig:118        plt.show()