一、使用GCN根据颗粒图像预测对应性能【含数据生成+源码实现】

之前做一个小实验写的代码，本想创建个git repo，想了想好像没必要，直接用篇博文记录一下吧。
对应资源 ： https://download.youkuaiyun.com/download/rayso9898/87865298

0. 大纲

0.1 代码说明

1. dataGeneration.py ->
   RSA生成n张图像，可以指定颗粒个数为m,半径为r。

2. exGraph.py ->
   将一张sim仿真图像，转换成一个node_feature matrix, n * 4 : n个节点，3个特征-质心x y,等效直径，面积
   有r12 - r22共 6 种不同粒径的图像，组成字典，每种类别有100张图的node_feature

3. create_dataset.py ->
   主要是 create_graph函数

4. gcn.py ->
   gcn回归模型

5. task_script ->
   提供模型训练 预测保存 loss查看等功能

0.2 数据说明

1. sim_data.zip -> dataGeneration.py 生成
   r12-r22共6个文件夹，每个文件夹100张图像。

2. sim_gdata.pt -> exGraph.py 生成
   sim_gdata structure:
   all - r12 r14 r16 r18 r20 r22
   r12 - img1 img2 … img100
   img1 - particle1 … paritle n
   particle1 - [x,y,dalimeter,area]

3. dataset_img_property.pt -> create_dataset.py 生成
   x = [] # 图所对应的节点 num4 （x,y,r,area）
   pos = []# 图所对应节点的位置 num2 (x,y)
   y = stress[i] # 图所对应的力学性能数据
   x = torch.tensor(x,dtype=torch.float32)
   # 构造一张img的一个图
   y = torch.tensor([y],dtype=torch.float32)
   pos = torch.tensor(pos,dtype=torch.float32)
   edge_index = knn_graph(pos, k=5)
   g = Data(x=x, y=y, pos=pos,edge_index=edge_index)
   all.append(g)

4. gcn.pt ->
   模型 后续加载即可使用

5. loss_gcn.pt
   训练过程中的训练集和测试集loss变化

1. dataGeneration.py ->随机序列吸附法（RSA）生成颗粒图像

可以指定生成颗粒的个数和颗粒半径，以及生成的图像个数。

# -*- coding: utf-8 -*-
# @Time    : 2021/5/1 15:12
# @Author  : Ray_song
# @File    : dataGeneration.py
# @Software: PyCharm

import cv2
import math
import numpy as np

def calcDis(a,b,c,d):
    return math.sqrt(math.pow(a-c,2)+math.pow(b-d,2))

def generate(n,r,filename,save_path):
    '''
    基于RSA算法，生成一张图像
    :param n: 该图像中有n个颗粒
    :param r: 颗粒的半径大小为r
    :param filename: 图像的名字  如 1 2 3 ....
    :param save_path: 图像的保存路径 - 文件夹
    :return: None
    '''
    # 1.创建白色背景图片
    d = 512
    img = np.ones((d, d, 3), np.uint8) * 0
    #testing

    list = []
    center_x = np.random.randint(0, high=d)
    center_y = np.random.randint(0, high=d)
    list.append([center_x,center_y])

    # 随机半径与颜色
    radius = r
    color = (0, 255, 0)
    cv2.circle(img, (center_x, center_y), radius, color, -1)

    # 2.循环随机绘制实心圆
    for i in range(1, n):
        flag = True
        # 随机中心点
        while flag:
            center_x_new = np.random.randint(radius, high=d-radius)
            center_y_new = np.random.randint(radius, high=d-radius)
            panduan = True
            for per in list:
                Dis = calcDis(center_x_new, center_y_new, per[0], per[1])
                if Dis<2*r:
                    panduan = False
                    break
                else:
                    continue
            if panduan:
                list.append([center_x_new,center_y_new])
                cv2.circle(img, (center_x_new, center_y_new), radius, color, -1)
                break
    # 3.显示结果
    # cv2.imshow("img", img)
    # cv2.waitKey()
    # cv2.destroyAllWindows()

    # 4.保存结果
    root = f'{
     
     save_path}/{
     
     filename}.jpg'
    cv2.imwrite(root,img)

def main():
    # example1 ： 随机生成 100张 颗粒个数为80 半径为20 的 图像
    save_path = 'sim_data/r20'
    for i in range(100):
        generate(80,20,i+1,save_path)

if __name__ == '__main__':
    main()

2. exGraph.py

将一张sim仿真图像，转换成一个node_feature matrix, n * 4 : n个节点，3个特征-质心x y,等效直径，面积
有r12 - r22共 6 种不同粒径的图像，组成字典，每种类别有100张图的node_feature

# -*- coding: utf-8 -*-
# @Time    : 2021/11/3 22:18
# @Author  : Ray_song
# @File    : exGraph.py
# @Software: PyCharm

import os
import torch

# 计算面积占比
def countArea(img):
    # 返回面积占比
    area = 0
    size = img.shape
    height,width = size[0