from PySide6.QtWidgets import QGraphicsScene, QGraphicsEllipseItem, QGraphicsLineItem from PySide6.QtGui import QBrush, QColor, QPen from PySide6.QtCore import QPointF from app.gui.node_widget import NodeWidget class Port(QGraphicsEllipseItem): """An input/output port that is placed next to each node.""" def __init__(self, parent_node: NodeWidget, kind: str, offset_x: float, offset_y: float): super().__init__(-5, -5, 10, 10, parent_node) self.parent_node = parent_node self.kind = kind # "input" Or "output" self.setBrush(QBrush(QColor("#2e7d32") if kind == "input" else QColor("#1565c0"))) self.setPos(parent_node.rect().x() + offset_x, parent_node.rect().y() + offset_y) self.setZValue(1.0) class ConnectionLine(QGraphicsLineItem): """Connection between two ports.""" def __init__(self, source_port: Port, target_port: Port): super().__init__() self.source_port = source_port self.target_port = target_port pen = QPen(QColor("#555")) pen.setWidth(2) self.setPen(pen) self.update_positions() def update_positions(self): src = self.source_port.mapToScene(QPointF(0, 0)) tgt = self.target_port.mapToScene(QPointF(0, 0)) self.setLine(src.x(), src.y(), tgt.x(), tgt.y()) class GraphScene(QGraphicsScene): """Graph scene with nodes, ports, and connections. Only connected nodes are exported.""" def __init__(self): super().__init__() self.connections: list[ConnectionLine] = [] # Manual connection mode: You click on an outgoing port, then click on the destination incoming port. self._pending_source_port: Port | None = None # ---------- Helper tool for adding default nodes ---------- def add_default_medical_nodes(self): """Adds five medical nodes and installs their ports.""" # Create nodes reader = NodeWidget("Image Reader", [], ["output"]) registration = NodeWidget("Image Registration", ["input"], ["output"]) fusion = NodeWidget("Image Fusion", ["input1", "input2"], ["output"]) extraction = NodeWidget("Image Extraction", ["input"], ["output"]) writer = NodeWidget("Image Writer", ["input"], []) # Initial placement reader.setPos(50, 50) registration.setPos(250, 50) fusion.setPos(450, 50) extraction.setPos(650, 50) writer.setPos(850, 50) # Add to scene for n in (reader, registration, fusion, extraction, writer): self.addItem(n) # Install ports for each node (a simple I/O is enough) # Reader: Output only reader.output_port = Port(reader, "output", 150, 45) # Registration: Input + Output registration.input_port = Port(registration, "input", 0, 45) registration.output_port = Port(registration, "output", 150, 45) # Fusion: Two inputs + one output fusion.input_port_1 = Port(fusion, "input", 0, 30) fusion.input_port_2 = Port(fusion, "input", 0, 60) fusion.output_port = Port(fusion, "output", 150, 45) # Extraction: Input + Output extraction.input_port = Port(extraction, "input", 0, 45) extraction.output_port = Port(extraction, "output", 150, 45) # Writer: Input only writer.input_port = Port(writer, "input", 0, 45) # Register the port click handle for manual connection self._install_port_handlers([reader, registration, fusion, extraction, writer]) def _install_port_handlers(self, nodes: list[NodeWidget]): """Mouse handles for ports so the user can make connections.""" all_ports: list[Port] = [] for n in nodes: for attr in dir(n): if attr.endswith("port"): p = getattr(n, attr) if isinstance(p, Port): all_ports.append(p) for port in all_ports: port.mousePressEvent = lambda event, p=port: self._on_port_clicked(p) def _on_port_clicked(self, port: Port): """Ports click logic: output first, then input; connection is made.""" if port.kind == "output": # Start connection self._pending_source_port = port elif port.kind == "input" and self._pending_source_port is not None: # Connection completion line = ConnectionLine(self._pending_source_port, port) self.addItem(line) self.connections.append(line) # Update positions when nodes move self._pending_source_port = None # ----------Export/Load Graph ---------- def export_graph(self): """Exports only nodes involved in connections.""" data = {"nodes": [], "connections": []} used_nodes = set() for conn in self.connections: src_node = conn.source_port.parent_node tgt_node = conn.target_port.parent_node used_nodes.add(src_node) used_nodes.add(tgt_node) data["connections"].append({ "source": src_node.title.toPlainText(), "target": tgt_node.title.toPlainText() }) for node in used_nodes: data["nodes"].append({ "id": node.title.toPlainText(), # Persistent ID based on title "type": node.title.toPlainText(), "pos": [node.pos().x(), node.pos().y()] }) return data def save_graph(self, filename: str): import json from pathlib import Path Path(filename).write_text(json.dumps(self.export_graph(), indent=2)) def load_graph(self, filename: str): """Rebuild ports and connections from saved file. Note: Here we assume that the default nodes are present on the scene; we only draw the connections. """ import json from pathlib import Path if not Path(filename).exists(): return data = json.loads(Path(filename).read_text()) # Create a mapping from title → node in the scene title_to_node = {} for item in self.items(): if isinstance(item, NodeWidget): title_to_node[item.title.toPlainText()] = item # Rebuild connections self._rebuild_ports_if_missing(title_to_node) for conn in data.get("connections", []): src = title_to_node.get(conn["source"]) tgt = title_to_node.get(conn["target"]) if not src or not tgt: continue # Select default input/output port based on node name src_port = getattr(src, "output_port", None) tgt_port = getattr(tgt, "input_port", None) # Fusion has two inputs; if the destination is Fusion and the first input is busy, take the second input if tgt.title.toPlainText() == "Image Fusion": # If there is no previous connection to input_port_1, connect to it; otherwise connect to input_port_2 candidates = [tgt.input_port_1, tgt.input_port_2] tgt_port = candidates[0] for c in self.connections: if c.target_port is candidates[0]: tgt_port = candidates[1] break if src_port and tgt_port: line = ConnectionLine(src_port, tgt_port) self.addItem(line) self.connections.append(line) def _rebuild_ports_if_missing(self, title_to_node: dict): """If the node ports haven't been created yet for some reason, we'll create them here.""" for title, n in title_to_node.items(): # Reader if title == "Image Reader" and not hasattr(n, "output_port"): n.output_port = Port(n, "output", 150, 45) # Registration if title == "Image Registration": if not hasattr(n, "input_port"): n.input_port = Port(n, "input", 0, 45) if not hasattr(n, "output_port"): n.output_port = Port(n, "output", 150, 45) # Fusion if title == "Image Fusion": if not hasattr(n, "input_port_1"): n.input_port_1 = Port(n, "input", 0, 30) if not hasattr(n, "input_port_2"): n.input_port_2 = Port(n, "input", 0, 60) if not hasattr(n, "output_port"): n.output_port = Port(n, "output", 150, 45) # Extraction if title == "Image Extraction": if not hasattr(n, "input_port"): n.input_port = Port(n, "input", 0, 45) if not hasattr(n, "output_port"): n.output_port = Port(n, "output", 150, 45) # Writer if title == "Image Writer" and not hasattr(n, "input_port"): n.input_port = Port(n, "input", 0, 45)