基于图片颜色特征的图数据分类 - 使用 PyTorch Geometric - 常规

import\u0020os\nimport\u0020pandas\u0020as\u0020pd\nimport\u0020torch\nimport\u0020torch.nn\u0020as\u0020nn\nfrom\u0020torch_geometric.data\u0020import\u0020Data,\u0020DataLoader\nfrom\u0020torch_geometric.nn\u0020import\u0020GCNConv\nimport\u0020torch.nn.functional\u0020as\u0020F\nfrom\u0020sklearn.model_selection\u0020import\u0020train_test_split\nfrom\u0020PIL\u0020import\u0020Image\n\n#\u0020加载数据并创建 PyG 数据集类：\nclass\u0020MyDataset(torch.utils.data.Dataset):\n\u0020\u0020def\u0020__init__(self, root, transform=None, pre_transform=None):\n\u0020\u0020\u0020\u0020self.edges\u0020=\u0020pd.read_csv(os.path.join(root, 'edges_L.csv'))\n\u0020\u0020\u0020\u0020self.transform\u0020=\u0020transform\n\u0020\u0020\u0020\u0020self.pre_transform\u0020=\u0020pre_transform\n\u0020\u0020\u0020\u0020self.num_classes\u0020=\u00208\u0020#\u0020修改成 8 类标签\n\u0020\u0020\u0020\u0020self.features\u0020=\u0020[]\n\u0020\u0020\u0020\u0020self.labels\u0020=\u0020[]\n\u0020\u0020\u0020\u0020for\u0020i\u0020in\u0020range(1, 43):\n\u0020\u0020\u0020\u0020\u0020\u0020for\u0020j\u0020in\u0020range(37):\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020image_path\u0020=\u0020os.path.join(root, 'images_block', f'{i}.png_{j}.png')\u0020#\u0020修改图片路径\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020label_path\u0020=\u0020os.path.join(root, 'labels', f'{i}{j}.txt')\u0020#\u0020修改标签文件名\n\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020#\u0020加载图片并调整大小\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020image\u0020=\u0020Image.open(image_path)\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020image\u0020=\u0020image.resize((40, 40))\n\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020#\u0020提取颜色特征\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020features\u0020=\u0020self.extract_features(image)\n\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020labels\u0020=\u0020pd.read_csv(label_path, header=None, sep=' ', encoding='ansi')\n\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020self.features.append(torch.tensor(features, dtype=torch.float))\n\u0020\u0020\u0020\u0020\u0020\u0020\u0020\u0020self.labels.append(torch.tensor(labels.values.squeeze(), dtype=torch.long))\u0020#\u0020修改标签数据类型为 long 型\n\n\u0020\u0020def\u0020extract_features(self, image):\n\u0020\u0020\u0020\u0020#\u0020提取颜色特征的代码\n\u0020\u0020\u0020\u0020#\u0020...\n\n\u0020\u0020def\u0020__len__(self):\n\u0020\u0020\u0020\u0020return\u0020len(self.features)\n\n\u0020\u0020def\u0020__getitem__(self, idx):\n\u0020\u0020\u0020\u0020edge_index\u0020=\u0020torch.tensor(self.edges.values, dtype=torch.long).t().contiguous()\n\u0020\u0020\u0020\u0020x\u0020=\u0020self.features[idx]\n\u0020\u0020\u0020\u0020y\u0020=\u0020self.labels[idx]\n\n\u0020\u0020\u0020\u0020#\u0020定义图数据的 train_mask 和 val_mask\n\u0020\u0020\u0020\u0020train_mask\u0020=\u0020torch.zeros(y.size(0), dtype=torch.bool)\n\u0020\u0020\u0020\u0020val_mask\u0020=\u0020torch.zeros(y.size(0), dtype=torch.bool)\n\u0020\u0020\u0020\u0020train_mask[:30]\u0020=\u00201\n\u0020\u0020\u0020\u0020val_mask[30:]\u0020=\u00201\n\n\u0020\u0020\u0020\u0020data\u0020=\u0020Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask)\n\n\u0020\u0020\u0020\u0020if\u0020self.transform\u0020is\u0020not\u0020None:\n\u0020\u0020\u0020\u0020\u0020\u0020data\u0020=\u0020self.transform(data)\n\n\u0020\u0020\u0020\u0020return\u0020data\n#\u0020定义 GCN 模型：\nclass\u0020GCN(torch.nn.Module):\n\u0020\u0020def\u0020__init__(self, num_node_features, num_classes):\n\u0020\u0020\u0020\u0020super(GCN, self).init()\n\u0020\u0020\u0020\u0020self.conv1\u0020=\u0020GCNConv(num_node_features, 8)\n\u0020\u0020\u0020\u0020self.conv2\u0020=\u0020GCNConv(8, 16)\n\u0020\u0020\u0020\u0020self.conv3\u0020=\u0020GCNConv(16, num_classes)\n\u0020\u0020def\u0020forward(self, data):\n\u0020\u0020\u0020\u0020x, edge_index\u0020=\u0020data.x, data.edge_index\n\u0020\u0020\u0020\u0020x\u0020=\u0020self.conv1(x, edge_index)\n\u0020\u0020\u0020\u0020x\u0020=\u0020F.relu(x)\n\u0020\u0020\u0020\u0020x\u0020=\u0020self.conv2(x, edge_index)\n\u0020\u0020\u0020\u0020x\u0020=\u0020F.relu(x)\n\u0020\u0020\u0020\u0020x\u0020=\u0020F.dropout(x, training=self.training)\n\u0020\u0020\u0020\u0020x\u0020=\u0020self.conv3(x, edge_index)\n\u0020\u0020\u0020\u0020return\u0020x\n#\u0020创建训练和验证模型：\ndef\u0020train_model(dataset, model, optimizer, device):\n\u0020\u0020model.train()\n\u0020\u0020total_loss\u0020=\u00200.0\n\n\u0020\u0020for\u0020data\u0020in\u0020dataset:\n\u0020\u0020\u0020\u0020data\u0020=\u0020data.to(device)\n\u0020\u0020\u0020\u0020optimizer.zero_grad()\n\u0020\u0020\u0020\u0020output\u0020=\u0020model(data)\n\u0020\u0020\u0020\u0020loss\u0020=\u0020F.cross_entropy(output[data.train_mask], data.y[data.train_mask])\n\u0020\u0020\u0020\u0020loss.backward()\n\u0020\u0020\u0020\u0020optimizer.step()\n\u0020\u0020\u0020\u0020total_loss\u0020+=\u0020loss.item()\n\u0020\u0020return\u0020total_loss\u0020/\u0020len(dataset)\n\ndef\u0020validate_model(dataset, model, device):\n\u0020\u0020model.eval()\n\u0020\u0020correct\u0020=\u00200\n\u0020\u0020total\u0020=\u00200\n\n\u0020\u0020for\u0020data\u0020in\u0020dataset:\n\u0020\u0020\u0020\u0020data\u0020=\u0020data.to(device)\n\u0020\u0020\u0020\u0020output\u0020=\u0020model(data)\n\u0020\u0020\u0020\u0020, predicted\u0020=\u0020torch.max(output[data.val_mask], 1)\n\u0020\u0020\u0020\u0020total\u0020+=\u0020data.val_mask.sum().item()\n\u0020\u0020\u0020\u0020correct\u0020+=\u0020(predicted\u0020==\u0020data.y[data.val_mask]).sum().item()\n\u0020\u0020return\u0020correct\u0020/\u0020total\n#\u0020加载数据集、创建模型、定义优化器和训练循环，以及验证模型：\nif\u0020__name__\u0020==\u0020'main':\n\u0020\u0020dataset\u0020=\u0020MyDataset(root="C:\Users\jh\Desktop\data\input")\n\u0020\u0020device\u0020=\u0020torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n\u0020\u0020model\u0020=\u0020GCN(num_node_features=1600, num_classes=8).to(device)\u0020#\u0020修改 num_classes 为 8\n\u0020\u0020optimizer\u0020=\u0020torch.optim.Adam(model.parameters(), lr=0.01)\n\n\u0020\u0020train_dataset, val_dataset\u0020=\u0020train_test_split(dataset, test_size=0.1)\n\u0020\u0020train_loader\u0020=\u0020DataLoader(train_dataset, batch_size=1, shuffle=False)\n\u0020\u0020val_loader\u0020=\u0020DataLoader(val_dataset, batch_size=1, shuffle=False)\n\n\u0020\u0020epochs\u0020=\u00202\n\u0020\u0020for\u0020epoch\u0020in\u0020range(epochs):\n\u0020\u0020\u0020\u0020train_loss\u0020=\u0020train_model(train_loader, model, optimizer, device)\n\u0020\u0020\u0020\u0020print(f'Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}')\n\n\u0020\u0020val_accuracy\u0020=\u0020validate_model(val_loader, model, device)\n\u0020\u0020print(f'Val_Acc: {val_accuracy:.4f}')