Hep2-Segmentation/train_utils.py at main · prasathlab/Hep2-Segmentation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
from Rand_Augment.rand_augment import *
from image_utils import *


def img_convert(img, target_type_min, target_type_max, target_type):
    img_min = img.min()
    img_max = img.max()
    print(img_max)
    print(img_min)
    if img_max - img_min <=0:
        #pdb.set_trace()
        print("Here it is> Divide by Zero")

    a = (target_type_max - target_type_min) / (img_max - img_min)
    b = target_type_max - a * img_max
    new_img = (a * img + b).astype(target_type)
    return img_min, img_max, new_img

def get_batch(img_files, mask_dir, batch_size):
    sample_idxs = random.sample(range(len(img_files)), k=batch_size)
    batch_imgs, batch_masks = [], []
    for idx in sample_idxs:
        img = np.load(img_files[idx])
        batch_imgs.append(img)
        img_basename = os.path.basename(img_files[idx])
        img_name = os.path.splitext(img_basename)[0]
        mask_name = img_name + "_mask.npy"
        mask_path = os.path.join(mask_dir, mask_name)
        mask = np.load(mask_path)
        batch_masks.append(mask)
    batch_imgs = np.array(batch_imgs)
    batch_masks = np.array(batch_masks)
    check = batch_imgs.shape == batch_masks.shape
    return batch_imgs, batch_masks


def make_generator_w_aug(img_files, mask_dir, batch_size, aug_img_files, aug_mask_dir):
    unaug_batch_size = 1
    #unaug_batch_size = int(batch_size//2)
    aug_batch_size = batch_size - unaug_batch_size
    while 1:
        unaug_batch_imgs, unaug_batch_masks = get_batch(img_files, mask_dir, unaug_batch_size)
        aug_batch_imgs, aug_batch_masks = get_batch(aug_img_files, aug_mask_dir, aug_batch_size)
        batch_imgs = np.concatenate([unaug_batch_imgs, aug_batch_imgs], axis=0)
        batch_masks = np.concatenate([unaug_batch_masks, aug_batch_masks], axis=0)
        yield (batch_imgs, batch_masks)


def make_generator(img_files, mask_dir, batch_size ):
    while 1:
        batch_imgs, batch_masks = get_batch(img_files, mask_dir, batch_size)
        yield (batch_imgs, batch_masks)

def split_train_val_set(img_files, valid_ratio):
    n_val_samples = int(valid_ratio*len(img_files))
    val_img_files = random.sample(img_files, k=n_val_samples)
    train_img_files = [img_file for img_file in img_files if img_file not in val_img_files]
    check = [x for x in train_img_files if x in val_img_files]
    return train_img_files, val_img_files

def get_validation_data(val_img_files, mask_dir):
    val_imgs, val_masks  = [], []
    for val_img_file in val_img_files:
        val_img = np.load(val_img_file)
        val_imgs.append(val_img)
        img_basename = os.path.basename(val_img_file)
        img_name = os.path.splitext(img_basename)[0]
        mask_name = img_name + "_mask.npy"
        mask_path = os.path.join(mask_dir, mask_name)
        val_mask = np.load(mask_path)
        val_masks.append(val_mask)
    val_imgs = np.array(val_imgs)
    val_masks = np.array(val_masks)
    return val_imgs, val_masks