def generator(coco, path, batch_size, class_count):
images_directory = os.path.join(os.path.dirname(path), "data")
files = {
int(f.split(".")[0]): os.path.join(images_directory, f)
for f in os.listdir(images_directory)
}
while True:
images = np.zeros((batch_size, 1024, 1024, 3), dtype=np.uint8)
labels = np.zeros((batch_size, 1024 // 8, 1024 // 8, class_count),
dtype=np.uint8)
count = 0
for image_id, segmentation in coco.imgToAnns.items():
image = Image.open(files[image_id])
images[count, :, :, :] = resize_and_crop(image, 1024)
for i, ann in enumerate(segmentation):
arr = coco.annToMask(ann)
image = Image.fromarray(arr)
resized = resize_and_crop(image, 1024 // 8, rgb=False)
labels[count, :, :, ann["category_id"] - 1] += resized
labels[count, labels[count, :, :, :] > 0] = 1
count += 1
if count == batch_size:
yield images, labels
images = np.zeros((batch_size, 1024, 1024, 3),
dtype=np.uint8)
labels = np.zeros(
(batch_size, 1024 // 8, 1024 // 8, class_count),
dtype=np.uint8)
count = 0
def __getitem__(self, idx):
img = readimage(self.dir_img + self.data[idx] + '.nii')
img = resize_and_crop(img)
img = np.expand_dims(img, axis=0)
mask = readimage(self.dir_mask + self.data[idx] + '_mask.nii')
mask = resize_and_crop(mask)
mask = channalize_mask(mask)
mask = np.transpose(mask, axes=[2, 0, 1])
return img.astype(np.float32), mask.astype(np.float32)
def generator(images_dir, labels_file, batch_size, class_count,
picture_size):
while True:
images = np.zeros((batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
labels = np.zeros((batch_size, class_count), dtype=np.uint8)
count = 0
for image_filename in os.listdir(images_dir):
centering = (round(random.random(),
1), round(random.random(), 1))
if image_filename[-3:] not in ("jpg", "png"):
continue
p = os.path.join(images_dir, image_filename)
image = Image.open(p)
images[count, :, :, :] = resize_and_crop(image,
picture_size,
centering=centering)
labels[count, label_dict[image_filename]] = 1
count += 1
if count == batch_size:
yield images, labels
images = np.zeros(
(batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
labels = np.zeros((batch_size, class_count))
count = 0
def prepare_image(image_path):
img_in = scipy.misc.imread(image_path, mode='RGB')
img = img_in.astype(np.float32)
img = utils.resize_and_crop(img, CROP_SIZE)
img = img.transpose(2, 0, 1)
return img
def predict_images(paths, model, output_directory):
images = np.zeros((len(paths), 512, 512, 3))
orig_sizes = []
for i, p in enumerate(paths):
img = Image.open(p)
orig_sizes.append(img.size)
images[i, :, :, :] = resize_and_crop(img, 512)[np.newaxis, :, :, :]
predictions = model.predict(images, batch_size=len(paths))
predictions *= 255
predictions = predictions.astype(np.uint8)
for i in range(predictions.shape[0]):
for j in range(predictions.shape[-1]):
new_w, new_h = (
(int(orig_sizes[i][0] * 512 / orig_sizes[i][1]), 512)
if orig_sizes[i][0] > orig_sizes[i][1]
else (512, int(orig_sizes[i][1] * 512 / orig_sizes[i][0]))
)
img = Image.new("RGB", (new_w, new_h))
ret, th = cv2.threshold(
predictions[i, 0 : images[0].shape[1], 0 : images[0].shape[0], j],
0,
255,
cv2.THRESH_OTSU,
)
img.paste(
Image.fromarray(th).resize((512, 512)),
(int(img.size[0] / 2 - 256), int(img.size[1] / 2 - 256)),
)
img = img.resize(orig_sizes[i])
p = ".".join(os.path.basename(paths[i]).split(".")[0:-1])
img.save(f"{output_directory}/{p}_{j}_segmentation.png")
def prepare_img_data(data):
#if not os.path.exists("/home/dogs/"):
# _download_dogs(DATASET_DIR)
#train_img_names, train_labels, test_img_names,test_labels = utils.read_image_names_and_assign_labels(class_size,test_split_number,DATASET_IMAGES)
train_data = utils.resize_and_crop(new_height, new_width, data)
train_data_arr = substract_mean(train_data)
print("method successful")
return train_data_arr
def predict_img(net,
full_img,
scale_factor=0.5,
out_threshold=0.5,
use_dense_crf=True,
use_gpu=False):
img_height = full_img.size[1]
img_width = full_img.size[0]
img = resize_and_crop(full_img, scale=scale_factor)
img = normalize(img)
left_square, right_square = split_img_into_squares(img)
left_square = hwc_to_chw(left_square)
right_square = hwc_to_chw(right_square)
X_left = torch.from_numpy(left_square).unsqueeze(0)
X_right = torch.from_numpy(right_square).unsqueeze(0)
if use_gpu:
X_left = X_left.cuda()
X_right = X_right.cuda()
with torch.no_grad():
output_left = net(X_left)
output_right = net(X_right)
left_probs = F.sigmoid(output_left).squeeze(0)
right_probs = F.sigmoid(output_right).squeeze(0)
'''
tf = transforms.Compose(
[
transforms.ToPILImage(),
transforms.Resize(img_height),
transforms.ToTensor()
]
)
'''
tf = transforms.Compose([
transforms.ToPILImage(),
transforms.Scale(img_height),
transforms.ToTensor()
])
left_probs = tf(left_probs.cpu())
right_probs = tf(right_probs.cpu())
left_mask_np = left_probs.squeeze().cpu().numpy()
right_mask_np = right_probs.squeeze().cpu().numpy()
full_mask = merge_masks(left_mask_np, right_mask_np, img_width)
if use_dense_crf:
full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
return full_mask > out_threshold
def classification_generator():
while True:
for image_filename in os.listdir(class_images_dir):
centering = (round(random.random(),
1), round(random.random(), 1))
if image_filename[-3:] not in ("jpg", "png"):
continue
p = os.path.join(class_images_dir, image_filename)
image = Image.open(p)
yield (
resize_and_crop(image, 512, centering=centering),
label_dict[image_filename],
)
def segmentation_generator():
while True:
for image_filename in os.listdir(seg_images_dir):
centering = (round(random.random(),
1), round(random.random(), 1))
if image_filename[-3:] not in ("jpg", "png"):
continue
p = os.path.join(os.getcwd(), seg_images_dir, image_filename)
image = Image.open(p)
mask = Image.open(
os.path.join(seg_mask_dir,
f"{image_filename[:-4]}_segmentation.png"))
yield (
resize_and_crop(image, 512, centering=centering),
resize_and_crop(
mask,
512 // (8 if model_size == "large" else 16),
centering=centering,
rgb=False,
),
)
def generator(samples_dir, batch_size, class_count):
files = os.listdir(samples_dir)
while True:
images = np.zeros((batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
labels = np.zeros(
(batch_size, picture_size // 8, picture_size // 8,
class_count),
dtype=np.uint8,
)
count = 0
random.shuffle(files)
for f in files:
centering = (round(random.random(),
1), round(random.random(), 1))
arr, masks = parse_labelme_file(os.path.join(samples_dir, f))
images[count, :, :, :] = resize_and_crop(arr,
picture_size,
centering=centering)
for l, mask in masks:
labels[count, :, :,
l] += resize_and_crop(mask,
picture_size // 8,
centering=centering,
rgb=False)
count += 1
if count == batch_size:
labels[labels > 0] = 1
yield images, labels
images = np.zeros(
(batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
labels = np.zeros(
(batch_size, picture_size // 8, picture_size // 8,
class_count),
dtype=np.uint8,
)
count = 0
def evaluate():
"""Runs inference with pretrained model.
"""
args = parse_arguments()
# Load model
model = MobileNetV3LiteRASPP(shape=(512, 512, 3),
n_class=args.class_count,
task=args.task)
label_dict = {}
with open(args.labels_file, "r") as f:
csvfile = csv.reader(f)
# Skip column description
next(csvfile)
for row in csvfile:
label_dict[row[0]] = row[1:].index("1.0")
if args.model_size == "large":
model = model.build_large()
else:
model = model.build_small()
model.load_weights(args.save_path, by_name=True)
images = np.zeros((len(os.listdir(args.input_dir)), 512, 512, 3))
labels = np.zeros((len(os.listdir(args.input_dir)), args.class_count))
for i, filename in enumerate(os.listdir(args.input_dir)):
img = resize_and_crop(
Image.open(os.path.join(args.input_dir, filename)), 512)
images[i, :, :, :] = np.array(img)
labels[i, label_dict[filename[:-4]]] = 1
preds = model.predict(images)
print(labels.argmax(axis=1))
print(preds.argmax(axis=1))
matrix = metrics.confusion_matrix(labels.argmax(axis=1),
preds.argmax(axis=1))
f1 = metrics.f1_score(labels.argmax(axis=1),
preds.argmax(axis=1),
average=None)
acc = 0
for i in range(matrix.shape[0]):
acc += matrix[i, i] / sum(matrix[i, :])
print(np.around(matrix / np.sum(matrix, axis=1)[:, None], decimals=2))
print(acc / args.class_count)
print(f1)
def predict_img(net,
full_img,
device,
scale_factor=1,
out_threshold=0.5,
use_dense_crf=False):
net.eval()
img_height = full_img.size[1]
img = resize_and_crop(full_img, scale=scale_factor)
img = normalize(img)
img = hwc_to_chw(img)
X = torch.from_numpy(img).unsqueeze(0)
X = X.to(device=device)
with torch.no_grad():
output = net(X)
if net.n_classes > 1:
probs = F.softmax(output, dim=1)
else:
probs = torch.sigmoid(output)
probs = probs.squeeze(0)
tf = transforms.Compose(
[
transforms.ToPILImage(),
transforms.Resize(img_height),
transforms.ToTensor()
]
)
probs = tf(probs.cpu())
full_mask = probs.squeeze().cpu().numpy()
if use_dense_crf:
full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
return full_mask > out_threshold
def predict_img(net,
full_img,
scale_factor=0.5,
out_threshold=0.5,
use_dense_crf=True,
use_gpu=False):
net.eval()
img = resize_and_crop(full_img, scale=scale_factor)
img = normalize(img)
img = hwc_to_chw(img)
X = torch.from_numpy(img).unsqueeze(0)
if use_gpu:
X = X.cuda()
with torch.no_grad():
output = net(X)
probs = output.squeeze(0)
tf = transforms.Compose(
[
transforms.ToPILImage(),
transforms.ToTensor()
]
)
probs = tf(probs.cpu())
print("Probs - ", probs.shape)
mask = probs.squeeze().cpu().numpy()
print("Mask - ", mask.shape)
plt.imshow(mask)
plt.show()
if use_dense_crf:
mask = dense_crf(np.array(full_img).astype(np.uint8), mask)
return mask > out_threshold
def predict_img_batch(net,
full_img,
scale_factor=0.5,
out_threshold=0.5,
use_dense_crf=True,
use_gpu=True):
"""return fullmask with size (C, H, W)"""
net.eval()
img_height = full_img.size[1]
img_width = full_img.size[0]
print('imgheight', img_height)
print('imgwidth', img_width)
img = resize_and_crop(full_img, scale=scale_factor)
img = normalize(img)
if len(img.shape) == 2:
img = img[..., np.newaxis]
print('img.shape', img.shape)
left_square, right_square = split_img_into_squares(img)
left_square = hwc_to_chw(left_square)
right_square = hwc_to_chw(right_square)
X_left = torch.from_numpy(left_square).unsqueeze(0)
X_right = torch.from_numpy(right_square).unsqueeze(0)
if use_gpu:
X_left = X_left.cuda()
X_right = X_right.cuda()
with torch.no_grad():
output_left = net(X_left)
output_right = net(X_right)
print('output_left.shape', output_left.shape)
print('output_right.shape', output_right.shape)
left_probs = output_left.squeeze(0)
right_probs = output_right.squeeze(0)
#
# if not scale_factor==1:
# tf = transforms.Compose(
# [
# transforms.ToPILImage(),
# transforms.Resize(img_height),
# transforms.ToTensor()
# ]
# )
#
# left_probs = tf(left_probs.cpu())
# right_probs = tf(right_probs.cpu())
# print('left_probs', left_probs.shape)
# print('right_probs', right_probs.shape)
left_mask_np = left_probs.squeeze().cpu().numpy()
right_mask_np = right_probs.squeeze().cpu().numpy()
left_mask_np = np.transpose(left_mask_np, axes=[1, 2, 0])
right_mask_np = np.transpose(right_mask_np, axes=[1, 2, 0])
if not scale_factor == 1:
right_mask_np = resize_np(right_mask_np, 1 / scale_factor)
left_mask_np = resize_np(left_mask_np, 1 / scale_factor)
print('left_mask_np', left_mask_np.shape)
print('right_mask_np', right_mask_np.shape)
left_mask_np = np.transpose(left_mask_np, axes=[2, 0, 1])
right_mask_np = np.transpose(right_mask_np, axes=[2, 0, 1])
full_mask = merge_masks(left_mask_np, right_mask_np, img_width)
# if use_dense_crf:
if 0:
full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
return full_mask
def prepare_image(img_in, crop_size):
img = utils.resize_and_crop(img_in, crop_size)
img = img.astype(np.float32)
img = img[None, ...]
return img
def prepare_image(img_in, crop_size):
img = utils.resize_and_crop(img_in, crop_size)
img = img.astype(INPUT_DATA_TYPE)
img = img.transpose(2, 0, 1) # to CHW
return img
def to_cropped_imgs(ids, dir, suffix):
"""From a list of tuples, returns the correct cropped img"""
for id, pos in ids:
im = resize_and_crop(Image.open(dir + id + suffix))
yield get_square(im, pos)
img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2GRAY)
mask = cv2.imread(dir_mask + '/' + id + '.png')
red_mask = np.array(mask[:, :, 2] == 128)
green_mask = np.array(mask[:, :, 1] == 128)
true_mask = np.stack(
[red_mask.astype(np.float32),
green_mask.astype(np.float32)])
img[red_mask] = img[red_mask] * 0.8
img[green_mask] = img[green_mask] * 1.2
img[np.logical_and(mask[:, :, 2] != 128, mask[:, :, 1] != 128)] = 0.2 * \
img[np.logical_and(mask[:, :, 2] != 128, mask[:, :, 1] != 128)]
img = resize_and_crop(img, scale=0.5)
img = normalize(img)[None, None, :, :]
img = torch.from_numpy(img).float()
net = UNet(1, 2)
net.eval()
net.load_state_dict(
torch.load(
'/media/zhuzhu/0C5809B80C5809B8/draft/unet/checkpoint/unet_0.854608765.pth',
map_location='cpu'))
predict = net(img).squeeze(0)
mask_predict = (predict > 0.5).float().numpy()
mask_blue = np.zeros(mask_predict.shape[1:])[np.newaxis, :]
mask_predict = np.concatenate([mask_predict, mask_blue], axis=0)
mask_predict = (mask_predict * 128).astype(np.uint8).transpose([1, 2, 0])
def generator(images_dir, masks_dir, batch_size, class_count):
while True:
images = np.zeros((batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
if model_size == "large":
labels = np.zeros(
(batch_size, picture_size // 8, picture_size // 8,
class_count),
dtype=np.uint8,
)
else:
labels = np.zeros(
(batch_size, picture_size // 16, picture_size // 16,
class_count),
dtype=np.uint8,
)
count = 0
for image_filename in os.listdir(images_dir):
centering = (round(random.random(),
1), round(random.random(), 1))
if image_filename[-3:] not in ("jpg", "png"):
continue
p = os.path.join(os.getcwd(), images_dir, image_filename)
image = Image.open(p)
images[count, :, :, :] = resize_and_crop(image,
picture_size,
centering=centering)
mask = Image.open(
os.path.join(masks_dir,
f"{image_filename[:-4]}_segmentation.png"))
if model_size == "large":
labels[count, :, :,
0] = resize_and_crop(mask,
picture_size // 8,
centering=centering,
rgb=False)
else:
labels[count, :, :,
0] = resize_and_crop(mask,
picture_size // 16,
centering=centering,
rgb=False)
labels[count, labels[count, :, :, :] > 0] = 1
count += 1
if count == batch_size:
yield images, labels
images = np.zeros(
(batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
if model_size == "large":
labels = np.zeros((
batch_size,
picture_size // 8,
picture_size // 8,
class_count,
))
else:
labels = np.zeros((
batch_size,
picture_size // 16,
picture_size // 16,
class_count,
))
count = 0
def predict_img(net,
full_img,
scale_factor=0.5,
out_threshold=0.5,
use_gpu=False):
pst = time.perf_counter()
net.eval()
#print(' 0 running time: %s seconds ' %(( time.clock() -pst)))
img_height = full_img.size[1]
# print(img_height)
img_width = full_img.size[0]
# print(full_img.size) # 2048,1229
img = resize_and_crop(full_img, scale=scale_factor)
#pdb.set_trace()
#print(' 1 running time: %s seconds ' %(( time.clock() -pst)))
img = normalize(img)
#print(' 2 running time: %s seconds ' %(( time.clock() -pst)))
# print(img.shape) # 614,1024,3
left_square, right_square = split_img_into_squares(img)
# print(right_square.shape) # 614,614,3
left_square = hwc_to_chw(left_square)
right_square = hwc_to_chw(right_square)
#print(' 3 running time: %s seconds ' %(( time.clock() -pst)))
X_left = torch.from_numpy(left_square).unsqueeze(0)
X_right = torch.from_numpy(right_square).unsqueeze(0)
#print(' 4 running time: %s seconds ' %(( time.clock() -pst)))
#outstart = time.clock()
if use_gpu:
X_left = X_left.cuda()
X_right = X_right.cuda()
#print(' 5 running time: %s seconds ' %(( time.clock() -pst)))
with torch.no_grad():
torch.cuda.synchronize()
st = time.perf_counter()
# print(X_left.shape) # 1,3,614,614
output_left = net(X_left)
output_right = net(X_right)
torch.cuda.synchronize()
st1 = time.perf_counter()
#outend = time.clock()
print(' Unet++ --------------------> running time: %s seconds ' %
(st1 - st))
left_probs = output_left.squeeze(0)
right_probs = output_right.squeeze(0)
# print(' squeeze running time: %s seconds ' %((time.perf_counter()-st1)))
if (left_probs.shape[1] != img_height):
tf = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(img_height),
transforms.ToTensor()
])
left_probs = tf(left_probs.cpu())
right_probs = tf(right_probs.cpu())
print("Transform done!")
#print(' 8 running time: %s seconds ' %(( time.clock() -pst)))
#lstart = time.clock()
#left_probs.cpu()
#print(' transforms running time: %s seconds ' %(( time.time() -pst)))
st = time.perf_counter()
left_mask_np = left_probs.squeeze().cpu().numpy()
end1 = time.perf_counter() - st
#print(left_probs.shape)
#pdb.set_trace()
# print(' tonumpy1 running time: %s seconds ' %(end1))
st = time.perf_counter()
right_mask_np = right_probs.squeeze().cpu().numpy()
end2 = time.perf_counter() - st
# print(' tonumpy2 running time: %s seconds ' %(end2))
full_mask = merge_masks(left_mask_np, right_mask_np, img_width)
# print(' 9 running time: %s seconds ' %(( time.perf_counter() -pst)))
#pdb.set_trace()
full_mask[full_mask >= out_threshold] = 1
full_mask[full_mask < out_threshold] = 0
#-------------------------------------------------------------------------------
#newmask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
#lend = time.clock()
#print(' running time: %s seconds ' %((lend-pst)))
#pdb.set_trace()
return full_mask > out_threshold
def generator(images_dir, masks_dir, batch_size, class_count):
while True:
random.shuffle(samples)
images = np.zeros((batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
if model_size == "large":
labels = np.zeros(
(batch_size, picture_size // 8, picture_size // 8,
class_count),
dtype=np.uint8,
)
else:
labels = np.zeros(
(batch_size, picture_size // 16, picture_size // 16,
class_count),
dtype=np.uint8,
)
count = 0
for image_path, mask_path, label in samples:
centering = (round(random.random(),
1), round(random.random(), 1))
try:
angle = random.choice([0, 45, 90, 270])
image = Image.open(image_path).rotate(angle)
images[count, :, :, :] = resize_and_crop(
image, picture_size, centering=centering)
#Image.fromarray(images[count, :, :, :]).save(f"out_debug/image_{count}.jpg")
mask = Image.open(mask_path).rotate(angle)
if model_size == "large":
labels[count, :, :,
label] = resize_and_crop(mask,
picture_size // 8,
centering=centering,
rgb=False)
else:
labels[count, :, :,
label] = resize_and_crop(mask,
picture_size // 16,
centering=centering,
rgb=False)
labels[count, labels[count, :, :, :] > 0] = 1
#Image.fromarray(labels[count, :, :, label] * 255).save(f"out_debug/mask_{count}.png")
count += 1
except Exception as ex:
print(ex)
if count == batch_size:
yield images, labels
images = np.zeros(
(batch_size, picture_size, picture_size, 3),
dtype=np.uint8)
if model_size == "large":
labels = np.zeros((
batch_size,
picture_size // 8,
picture_size // 8,
class_count,
))
else:
labels = np.zeros((
batch_size,
picture_size // 16,
picture_size // 16,
class_count,
))
count = 0
def predict_img(net,
full_img,
scale_factor=0.5,
out_threshold=0.5,
use_dense_crf=True,
use_gpu=False):
net.eval()
img_height = full_img.size[1]
img_width = full_img.size[0]
img = resize_and_crop(full_img, scale=scale_factor)
img = normalize(img)
# left_square, right_square = split_img_into_squares(img)
#
# left_square = hwc_to_chw(left_square)
# right_square = hwc_to_chw(right_square)
img = hwc_to_chw(img)
X = torch.from_numpy(img).unsqueeze(0)
# X_left = torch.from_numpy(left_square).unsqueeze(0)
# X_right = torch.from_numpy(right_square).unsqueeze(0)
if use_gpu:
X = X.cuda()
# X_left = X_left.cuda()
# X_right = X_right.cuda()
with torch.no_grad():
# output_left = net(X_left)
# output_right = net(X_right)
output = net(X)
# left_probs = output_left.squeeze(0)
# right_probs = output_right.squeeze(0)
probs = output.squeeze(0)
tf = transforms.Compose([
transforms.ToPILImage(),
# transforms.Resize(img_height),
transforms.ToTensor()
])
# left_probs = tf(left_probs.cpu())
# right_probs = tf(right_probs.cpu())
#
# left_mask_np = left_probs.squeeze().cpu().numpy()
# right_mask_np = right_probs.squeeze().cpu().numpy()
probs = tf(probs.cpu())
mask = probs.squeeze().cpu().numpy()
# full_mask = merge_masks(left_mask_np, right_mask_np, img_width)
if use_dense_crf:
# full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
mask = dense_crf(np.array(full_img).astype(np.uint8), mask)
return mask > out_threshold
