-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathutils.py
122 lines (107 loc) · 4.76 KB
/
utils.py
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import os
import cv2
import shutil
import numpy as np
from os import listdir
from PIL import Image
import torch
from torch.utils.data import DataLoader
from dataset import MammoDataset
Image_Size = [512, 512]
#generate patches
def preprocess_img(img_path):
im = cv2.imread(img_path)
imgray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
figure_size = 5 # the dimension of the x and y axis of the kernal.
new_image = cv2.medianBlur(imgray,figure_size)
return new_image
def get_slices(n): # help to determine where to start splitting patches
slices = []
for i in range(0, n, 480):
if i + Image_Size[0] >= n:
slices.append(n-Image_Size[0])
break
slices.append(i)
return slices
def crop_img(img, name, size, savepath): # crop mammogram into patches and save them to a temp place
x, y = img.shape
Xs = get_slices(x)
Ys = get_slices(y)
save_ext = ".png"
cnt = 0
for i in Xs:
for j in Ys:
cropped_image = img[i:i+size,j:j+size]
if np.mean(cropped_image) <= 20 or np.std(cropped_image) <= 10:
continue
cnt+=1
save_imgname = savepath+'/'+name[:-4]+"_"+str(i)+"_"+str(j)+save_ext
cv2.imwrite(save_imgname, np.array(cropped_image))
def generate_patches(datapath, temppath, imgname):
if not os.path.exists(os.path.join(temppath, imgname[:-4])):
os.makedirs(os.path.join(temppath, imgname[:-4]))
patchpath = os.path.join(temppath, imgname[:-4])
img = preprocess_img(datapath+imgname)
wholeimg_rootdir = patchpath+"whole"
if os.path.exists(wholeimg_rootdir):
return patchpath, img
crop_img(img,imgname,Image_Size[0],patchpath)
return patchpath, img
def evaluate(img, imgname, wholeimg_rootdir, prob_thre):
# area, probability, threshold, intensity>100
breastarea = np.sum((img>0).astype(int))
pre_mask = np.load(wholeimg_rootdir+'/'+imgname[:-4]+"_premask.npy")
prob = np.sum(pre_mask)
pre_mask_thre = (pre_mask>prob_thre).astype(int)
area_thre = np.sum(pre_mask_thre)
pre_image = img[np.where(pre_mask>prob_thre)]
pre_image2 = pre_image[np.where(pre_image >= 100)]
sum_intensity = sum(pre_image)
sum_intensity_100 = sum(pre_image2)
sum_pixels_100 = sum(pre_image >= 100)
imgsize = img.shape
pre_mask_thre = np.array(pre_mask_thre*255, dtype=np.uint8)
pre_mask_img = Image.fromarray(pre_mask_thre)
pre_mask_img.save(wholeimg_rootdir+'/'+imgname[:-4]+"_premask_"+str(prob_thre)+".png")
print("breastarea:", breastarea, " prob:", prob, " area_0.65:", area_thre, " sum_intensity:", sum_intensity, " sum_intensity_100:", sum_intensity_100, " area_100:", sum_pixels_100, " image.size:", imgsize)
return breastarea, prob, area_thre, sum_intensity, sum_intensity_100, sum_pixels_100, imgsize
def predict(net, datapath, temppath, imgname, prob_thre):
testdir, image = generate_patches(datapath, temppath, imgname)
wholeimg_rootdir = testdir+"whole"
print("wholeimg_rootdir: ",wholeimg_rootdir)
if os.path.exists(wholeimg_rootdir):
return evaluate(image, imgname, wholeimg_rootdir, prob_thre)
testset = MammoDataset(rootdir=testdir)
testloader = DataLoader(testset, batch_size=8, shuffle=False, pin_memory=torch.cuda.is_available(),num_workers=8)
patch_pre_mask_dir = testdir+"patch"
print("patch_pre_mask_dir: ",patch_pre_mask_dir)
with torch.no_grad():
for batch_idx, (img, imgnames) in enumerate(testloader):
img = img.type(torch.DoubleTensor)
pre_mask = net(img)
pmask = pre_mask.data.cpu().numpy()
for i in range(pmask.shape[0]):
name = imgnames[i]
if not os.path.exists(patch_pre_mask_dir):
os.makedirs(patch_pre_mask_dir)
np.save(patch_pre_mask_dir+'/'+name[:-4]+".npy", pmask[i][0])
print("Collect.....")
premask_names = os.listdir(patch_pre_mask_dir)
pre_mask = np.zeros(image.shape)
patchsize = 512
for prename in premask_names:
_, x, y = prename[:-4].split("_")
x = int(x)
y = int(y)
patch_mask = np.load(patch_pre_mask_dir+'/'+prename)
for i in range(x, x+patchsize):
for j in range(y, y+patchsize):
pre_mask[i][j] = max(pre_mask[i][j], patch_mask[i-x][j-y])
if not os.path.exists(wholeimg_rootdir):
os.makedirs(wholeimg_rootdir)
np.save(wholeimg_rootdir+'/'+imgname[:-4]+"_premask.npy",pre_mask)
if os.path.exists(testdir):
shutil.rmtree(testdir)
if os.path.exists(patch_pre_mask_dir):
shutil.rmtree(patch_pre_mask_dir)
return evaluate(image, imgname, wholeimg_rootdir, prob_thre)