modelnet40 example pass!

Author: yifanfeng97

HyperG/utils/data/pathology/patch_funs.py

@@ -0,0 +1,2 @@
+def sample_patch():
+    pass

examples/clssification/modelnet40/data_helper.py

@@ -0,0 +1,25 @@
+import numpy as np
+import scipy.io as scio
+import torch
+
+
+def load_ft(data_dir, feature_name='GVCNN'):
+    data = scio.loadmat(data_dir)
+    lbls = data['Y'].astype(np.long)
+    if lbls.min() == 1:
+        lbls = lbls - 1
+    idx = data['indices'].item()
+
+    if feature_name == 'MVCNN':
+        fts = data['X'][0].item().astype(np.float32)
+    elif feature_name == 'GVCNN':
+        fts = data['X'][1].item().astype(np.float32)
+    else:
+        print(f'wrong feature name{feature_name}!')
+        raise IOError
+
+    idx_train = (idx == 1)
+    idx_test = (idx == 0)
+    return torch.tensor(fts), torch.tensor(lbls).squeeze(), \
+           torch.tensor(idx_train).squeeze().bool(), \
+           torch.tensor(idx_test).squeeze().bool()

examples/clssification/modelnet40/train.py

@@ -0,0 +1,67 @@
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+from data_helper import load_ft
+
+from HyperG.hyedge import neighbor_distance
+from HyperG.hygraph import hyedge_concat
+from HyperG.models import HGNN
+from HyperG.utils.meter import trans_class_acc
+
+# initialize parameters
+data_root = '/repository/HyperG_example/example_data/modelnet40/processed'
+result_root = '/repository/HyperG_example/tmp/modelnet40'
+
+k_nearest = 10
+feature_dir = osp.join(data_root, 'ModelNet40_mvcnn_gvcnn.mat')
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# load data
+gvcnn_ft, target, mask_train, mask_val = load_ft(feature_dir, feature_name='GVCNN')
+mvcnn_ft, _, _, _ = load_ft(feature_dir, feature_name='MVCNN')
+
+# init H and X
+gvcnn_H = neighbor_distance(gvcnn_ft, k_nearest)
+mvcnn_H = neighbor_distance(mvcnn_ft, k_nearest)
+
+ft = torch.cat([mvcnn_ft, gvcnn_ft], dim=1)
+H = hyedge_concat([mvcnn_H, gvcnn_H])
+
+x_ch = ft.size(1)
+n_class = target.max().item() + 1
+model = HGNN(x_ch, n_class, hiddens=[128])
+
+model, ft, H, target, mask_train, mask_val = model.to(device), ft.to(device), H.to(device), \
+                                             target.to(device), mask_train.to(device), mask_val.to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+
+def train():
+    model.train()
+    optimizer.zero_grad()
+    pred = model(ft, H)
+    F.nll_loss(pred[mask_train], target[mask_train]).backward()
+    optimizer.step()
+
+
+def val():
+    model.eval()
+    pred = model(ft, H)
+
+    _train_acc = trans_class_acc(pred, target, mask_train)
+    _val_acc = trans_class_acc(pred, target, mask_val)
+
+    return _train_acc, _val_acc
+
+
+if __name__ == '__main__':
+    best_acc, best_iou = 0.0, 0.0
+    for epoch in range(1, 101):
+        train()
+        train_acc, val_acc = val()
+        if val_acc > best_acc:
+            best_acc = val_acc
+        print(f'Epoch: {epoch}, Train:{train_acc:.4f}, Val:{val_acc:.4f}, '
+              f'Best Val acc:{best_acc:.4f}')

examples/segment/heart_mri/train.py

@@ -1,14 +1,13 @@
-import os
 import os.path as osp
 
 import torch
 import torch.nn.functional as F
 from data_helper import preprocess, split_train_val
 
 from HyperG.models import HGNN
+from HyperG.utils import check_dir
 from HyperG.utils.meter import trans_class_acc, trans_iou_socre
 from HyperG.utils.visualization import trans_vis_pred_target
-from HyperG.utils import check_dir
 
 # initialize parameters
 data_root = '/repository/HyperG_example/example_data/heart_mri/processed'

test/utils/data/mri/test_read.py

@@ -5,11 +5,14 @@
 
 @pytest.mark.skip(reason='unpleasure')
 def test_read_mri():
-    img_dir = '/repository/HyperG_example/example_data/heart_mri/processed/0001.mha'
-    img = read_mri(img_dir)
-
-    # print(img.shape)
-    # import matplotlib.pyplot as plt
-    # img = img.squeeze().numpy()
-    # plt.imshow(img, cmap=plt.cm.bone)
-    # plt.show()
+    for _img_name in ['0001', '0002', '0003', '0004', '0005']:
+        img_dir = f'/repository/HyperG_example/example_data/heart_mri/processed/{_img_name}.mha'
+        save_dir = f'/repository/HyperG_example/tmp/heart_mri/{_img_name}.jpg'
+        img = read_mri(img_dir)
+
+        print(img.shape)
+        import matplotlib.pyplot as plt
+        img = img.squeeze().numpy()
+        plt.imshow(img, cmap=plt.cm.bone)
+        plt.imsave(save_dir, img, cmap=plt.cm.bone)
+        # plt.show()