support ssl

add semi- and weakly-supervised learning

Author: taigw

pymic/io/nifty_dataset.py

@@ -94,8 +94,6 @@ def __init__(self, root_dir, csv_file, modal_num = 1, class_num = 2,
         super(ClassificationDataset, self).__init__(root_dir, 
             csv_file, modal_num, with_label, transform)
         self.class_num = class_num
-        print("class number for ClassificationDataset", self.class_num)
-        print("self.transform", self.transform)
 
     def __getlabel__(self, idx):
         csv_keys = list(self.csv_items.keys())
@@ -127,7 +125,6 @@ def __getitem__(self, idx):
             sample['label'] = self.__getlabel__(idx) 
         if (self.image_weight_idx is not None):
             sample['image_weight'] = self.__getweight__(idx) 
-        print("***transform", self.transform)
         if self.transform:
             sample = self.transform(sample)
         return sample

pymic/net_run/agent_abstract.py

@@ -42,6 +42,7 @@ def __init__(self, config, stage = 'train'):
         self.scheduler = None 
         self.loss_dict = None 
         self.transform_dict  = None
+        self.inferer   = None
         self.tensor_type   = config['dataset']['tensor_type']
         self.task_type     = config['dataset']['task_type'] #cls, cls_mtbc, seg
         self.deterministic = config['training'].get('deterministic', True)
@@ -69,6 +70,9 @@ def set_optimizer(self, optimizer):
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
+    
+    def set_inferer(self, inferer):
+        self.inferer = inferer
 
     def get_checkpoint_name(self):
         ckpt_mode = self.config['testing']['ckpt_mode']

pymic/net_run/agent_cls.py

@@ -7,28 +7,20 @@
 import sys
 import time
 import random
-import scipy
 import torch
-import torchvision
-from torchvision import datasets, models, transforms
+from torchvision import transforms
 import numpy as np
 import torch.nn as nn
 import torch.optim as optim
 import torch.nn.functional as F
-from torch.optim import lr_scheduler
 import matplotlib.pyplot as plt
 from PIL import Image 
-from scipy import special
 from datetime import datetime
 from tensorboardX import SummaryWriter
-from pymic.io.image_read_write import save_nd_array_as_image
 from pymic.io.nifty_dataset import ClassificationDataset
 from pymic.loss.loss_dict_cls import PyMICClsLossDict
 from pymic.net.net_dict_cls import TorchClsNetDict
-from pymic.net_run.get_optimizer import get_optimiser
 from pymic.transform.trans_dict import TransformDict
-from pymic.util.image_process import convert_label
-from pymic.util.parse_config import parse_config
 from pymic.net_run.agent_abstract import NetRunAgent
 import warnings
 warnings.filterwarnings('ignore', '.*output shape of zoom.*')
@@ -207,6 +199,8 @@ def train_valid(self):
         self.net.to(self.device)
 
         ckpt_dir    = self.config['training']['ckpt_save_dir']
+        if(ckpt_dir[-1] == "/"):
+            ckpt_dir = ckpt_dir[:-1]
         ckpt_prefx  = ckpt_dir.split('/')[-1]
         iter_start  = self.config['training']['iter_start']
         iter_max    = self.config['training']['iter_max']
@@ -292,10 +286,7 @@ def infer(self):
         with torch.no_grad():
             for data in self.test_loader:
                 names  = data['names']
-                if(names[0] != "data3_process/20190711_1005487059.png"):
-                    continue
                 inputs = self.convert_tensor_type(data['image'])
-                print("intensity mean and std", inputs.mean(), inputs.std())
                 inputs = inputs.to(device) 
 
                 start_time = time.time()

pymic/net_run/agent_seg.py

@@ -221,21 +221,22 @@ def training(self):
 
     def validation(self):
         class_num = self.config['network']['class_num']
-        infer_cfg = self.config['testing']
-        infer_cfg['class_num'] = class_num
+        if(self.inferer is None):
+            infer_cfg = self.config['testing']
+            infer_cfg['class_num'] = class_num
+            self.inferer = Inferer(infer_cfg)
 
         valid_loss_list = []
         valid_dice_list = []
         validIter  = iter(self.valid_loader)
         with torch.no_grad():
             self.net.eval()
-            infer_obj = Inferer(self.net, infer_cfg)
             for data in validIter:
                 inputs      = self.convert_tensor_type(data['image'])
                 labels_prob = self.convert_tensor_type(data['label_prob'])
                 inputs, labels_prob  = inputs.to(self.device), labels_prob.to(self.device)
                 batch_n = inputs.shape[0]
-                outputs = infer_obj.run(inputs)
+                outputs = self.inferer.run(self.net, inputs)
 
                 # The tensors are on CPU when calculating loss for validation data
                 loss = self.get_loss_value(data, outputs, labels_prob)
@@ -286,6 +287,8 @@ def train_valid(self):
             self.device = torch.device("cuda:{0:}".format(device_ids[0]))
         self.net.to(self.device)
         ckpt_dir    = self.config['training']['ckpt_save_dir']
+        if(ckpt_dir[-1] == "/"):
+            ckpt_dir = ckpt_dir[:-1]
         ckpt_prefx  = ckpt_dir.split('/')[-1]
         iter_start  = self.config['training']['iter_start']
         iter_max    = self.config['training']['iter_max']
@@ -392,9 +395,10 @@ def test_time_dropout(m):
         checkpoint = torch.load(ckpt_name, map_location = device)
         self.net.load_state_dict(checkpoint['model_state_dict'])
 
-        infer_cfg = self.config['testing']
-        infer_cfg['class_num'] = self.config['network']['class_num']
-        infer_obj = Inferer(self.net, infer_cfg)
+        if(self.inferer is None):
+            infer_cfg = self.config['testing']
+            infer_cfg['class_num'] = self.config['network']['class_num']
+            self.inferer = Inferer(infer_cfg)
         infer_time_list = []
         with torch.no_grad():
             for data in self.test_loader:
@@ -412,7 +416,7 @@ def test_time_dropout(m):
                 # continue
                 start_time = time.time()
 
-                pred = infer_obj.run(images)
+                pred = self.inferer.run(self.net, images)
                 # convert tensor to numpy
                 if(isinstance(pred, (tuple, list))):
                     pred = [item.cpu().numpy() for item in pred]
@@ -438,10 +442,11 @@ def infer_with_multiple_checkpoints(self):
         device_ids = self.config['testing']['gpus']
         device = torch.device("cuda:{0:}".format(device_ids[0]))
 
+        if(self.inferer is None):
+            infer_cfg  = self.config['testing']
+            infer_cfg['class_num'] = self.config['network']['class_num']
+            self.inferer = Inferer(infer_cfg)
         ckpt_names = self.config['testing']['ckpt_name']
-        infer_cfg  = self.config['testing']
-        infer_cfg['class_num'] = self.config['network']['class_num']
-        infer_obj = Inferer(self.net, infer_cfg)
         infer_time_list = []
         with torch.no_grad():
             for data in self.test_loader:
@@ -463,7 +468,7 @@ def infer_with_multiple_checkpoints(self):
                     checkpoint = torch.load(ckpt_name, map_location = device)
                     self.net.load_state_dict(checkpoint['model_state_dict'])
 
-                    pred = infer_obj.run(images)
+                    pred = self.inferer.run(self.net, images)
                     # convert tensor to numpy
                     if(isinstance(pred, (tuple, list))):
                         pred = [item.cpu().numpy() for item in pred]

pymic/net_run/infer_func.py

@@ -9,9 +9,8 @@
 from torch.nn.functional import interpolate
 
 class Inferer(object):
-    def __init__(self, model, config):
-        self.model   = model 
-        self.config  = config
+    def __init__(self, config):
+        self.config = config
 
     def __infer(self, image):
         use_sw  = self.config.get('sliding_window_enable', False)
@@ -131,8 +130,9 @@ def __infer_with_sliding_window(self, image):
                 output_list[i] = output_list[i] / counter_i
             return output_list
 
-    def run(self, image):
-        tta_mode  = self.config.get('tta_mode', 0)
+    def run(self, model, image):
+        self.model = model
+        tta_mode   = self.config.get('tta_mode', 0)
         if(tta_mode == 0):
             outputs = self.__infer(image)
         elif(tta_mode == 1): # test time augmentation with flip in 2D

pymic/net_run_noise/cl.py

@@ -77,10 +77,11 @@ def test_time_dropout(m):
         checkpoint = torch.load(ckpt_name, map_location = device)
         self.net.load_state_dict(checkpoint['model_state_dict'])
 
-        infer_cfg = self.config['testing']
-        class_num = self.config['network']['class_num']
-        infer_cfg['class_num'] = class_num 
-        infer_obj = Inferer(self.net, infer_cfg)
+        if(self.inferer is None):
+            infer_cfg = self.config['testing']
+            class_num = self.config['network']['class_num']
+            infer_cfg['class_num'] = class_num 
+            self.inferer = Inferer(infer_cfg)
         pred_list  = []
         gt_list    = []
         filename_list = []
@@ -102,7 +103,7 @@ def test_time_dropout(m):
                 #     save_nd_array_as_image(label_i, label_name, reference_name = None)
                 # continue
 
-                pred = infer_obj.run(images)
+                pred = self.inferer.run(self.net, images)
                 # convert tensor to numpy
                 if(isinstance(pred, (tuple, list))):
                     pred = [item.cpu().numpy() for item in pred]