import torch import torch.nn as nn from torch_geometric.nn import GCNConv from torch_geometric.data import Data, DataLoader from torch.utils.data import Dataset import torchvision.transforms as transforms import pandas as pd from PIL import Image

Define the CNN model to reduce node pixel feature dimension

class CNN(nn.Module): def init(self): super(CNN, self).init() self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.fc1 = nn.Linear(32 * 10 * 10, 128) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(128, 8)

def forward(self, x):
    out = self.conv1(x)
    out = self.relu1(out)
    out = self.pool1(out)
    out = self.conv2(out)
    out = self.relu2(out)
    out = self.pool2(out)
    out = out.view(out.size(0), -1)
    out = self.fc1(out)
    out = self.relu3(out)
    out = self.fc2(out)
    return out

Define the GCN model for multi-label classification

class GCN(nn.Module): def init(self, num_features, hidden_size, num_labels): super(GCN, self).init() self.conv1 = GCNConv(num_features, hidden_size) self.relu1 = nn.ReLU() self.conv2 = GCNConv(hidden_size, hidden_size) self.relu2 = nn.ReLU() self.fc = nn.Linear(hidden_size, num_labels)

def forward(self, x, edge_index):
    x = self.conv1(x, edge_index)
    x = self.relu1(x)
    x = self.conv2(x, edge_index)
    x = self.relu2(x)
    x = self.fc(x)
    return x

Define the custom dataset

class GraphDataset(Dataset): def init(self, num_graphs, num_nodes, image_size): self.num_graphs = num_graphs self.num_nodes = num_nodes self.image_size = image_size self.transform = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor() ])

def __len__(self):
    return self.num_graphs

def __getitem__(self, index):
    # Load node pixel features
    nodes = []
    for i in range(self.num_nodes):
        image_path = f"C:/Users/jh/Desktop/data/input/images{index+1}.png_{i}.png"
        image = Image.open(image_path)
        image = self.transform(image)
        nodes.append(image)

    # Load node labels
    labels = []
    for j in range(self.num_nodes):
        label_path = f"C:/Users/jh/Desktop/data/input/labels{index+1}_{j}.txt"
        with open(label_path, 'r') as file:
            label = [int(x) for x in file.read().split()]
            labels.append(label)

    # Load edge relationships
    edges_path = "C:/Users/jh/Desktop/data/input/edges_L.csv"
    edges_df = pd.read_csv(edges_path, header=None)
    edges = edges_df.values.tolist()

    x = torch.stack(nodes)
    y = torch.tensor(labels)
    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()

    data = Data(x=x, y=y, edge_index=edge_index)
    return data

Create the CNN model and load the pre-trained weights

cnn_model = CNN() cnn_model.load_state_dict(torch.load('cnn_model.pth'))

Create the GCN model

gcn_model = GCN(num_features=8, hidden_size=64, num_labels=8)

Create the dataset and data loader for training

train_dataset = GraphDataset(num_graphs=38, num_nodes=37, image_size=40) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)

Create the dataset and data loader for testing

test_dataset = GraphDataset(num_graphs=4, num_nodes=37, image_size=40) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

Define the loss function

loss_fn = nn.MultiLabelSoftMarginLoss()

Define the optimizer

optimizer = torch.optim.Adam(gcn_model.parameters(), lr=0.01)

Train the model

for epoch in range(100): gcn_model.train() for data in train_loader: optimizer.zero_grad() x = data.x.view(-1, 3, 40, 40) # Reshape node pixel features for CNN input x = cnn_model(x) # Reduce node pixel feature dimension using CNN y = data.y.float() edge_index = data.edge_index output = gcn_model(x, edge_index) loss = loss_fn(output, y) loss.backward() optimizer.step()

Test the model

gcn_model.eval() with torch.no_grad(): total_correct = 0 total_samples = 0 for data in test_loader: x = data.x.view(-1, 3, 40, 40) # Reshape node pixel features for CNN input x = cnn_model(x) # Reduce node pixel feature dimension using CNN y = data.y edge_index = data.edge_index output = gcn_model(x, edge_index) predicted_labels = torch.round(torch.sigmoid(output)) correct = (predicted_labels == y).sum().item() total_correct += correct total_samples += y.size(0) accuracy = total_correct / total_samples print(f"Accuracy: {accuracy}")