def predict_image(model_path: str, image: np.ndarray, show: bool = False):
pred = Prediction(model_path=model_path,
output_path=None,
patch_size=train_params['patch_size'],
batch_size=train_params['batch_size'],
offset=2 * train_params["n_layers"],
n_classes=len(classes_mask.keys()),
n_filters=train_params["n_filters"],
n_layers=train_params["n_layers"])
# image = np.array(Image.open(image)).astype(np.float32) / 255
predicted = pred.__predict_image__(image)
predicted = np.argmax(predicted, axis=2)
if show:
Image.fromarray(
colorize_mask(np.dstack((predicted, predicted, predicted)),
n_classes=pred.n_classes)).show()
return predicted
np.argmax(predicted, axis=2),
os.path.basename(image_name),
output_path=self.output_path)
# predicted = np.argmax(predicted, axis=-1)
# Image.fromarray(colorize_mask(np.dstack((predicted, predicted, predicted)), n_classes=self.n_classes).astype(np.uint8)).save(os.path.join(self.output_path, os.path.basename(image_name)))
# visualize_prediction_result(image, res, os.path.basename(image_name), output_path=self.output_path)
if __name__ == "__main__":
output_path = os.path.join("output", "prediction")
if not os.path.exists(output_path):
os.makedirs(output_path)
pred = Prediction(
model_path="/home/akochkarev/geology/models/model_46_0.07.hdf5",
output_path=output_path,
patch_size=train_params["patch_size"],
batch_size=train_params["batch_size"],
offset=2 * train_params["n_layers"],
n_classes=len(classes_mask.keys()),
n_filters=train_params["n_filters"],
n_layers=train_params["n_layers"])
# unmarked = get_unmarked_images(os.path.join("input", "UMNIK_2019", "BoxA_DS1", "img"), os.path.join("input", "dataset"))
unmarked = [
f"/home/akochkarev/geology/test_img1/{i+1}.jpg" for i in range(4)
]
print(unmarked)
pred.predict(unmarked)
def __iter__(self):
assert (self.images.shape[:3] == self.masks.shape[:3]), (
"Original images and masks must be equal shape")
classes = {cl: i for i, cl in enumerate(classes_mask.values())}
aa = np.arange(self.masks[0].shape[0] * self.masks[0].shape[1])
a = np.indices(self.masks[0].shape)
_xx = np.ndarray.flatten(a[0])
_yy = np.ndarray.flatten(a[1])
with open(os.path.join("input", "ores.json")) as ores_json:
ores = json.load(ores_json)
while (True):
cur_class = min(self.stat.keys(), key=(lambda k: self.stat[k]))
pp = np.fromiter(ores[cur_class].values(), dtype=np.float64)
pp /= np.sum(pp)
batch_files = np.random.choice(a=self.images.shape[0],
size=self.batch_size,
p=pp)
batch_x = []
batch_y = []
for batch_idx in batch_files:
if self.full_augment:
etha_max = 2
etha = np.random.uniform(1 / etha_max, etha_max)
angle = np.random.random_integers(0, 360)
new_size = int(np.ceil(self.patch_size * sqrt(2) * etha))
p_s = self.heatmaps[batch_idx]
p = p_s[classes[cur_class]]
if self.full_augment:
while (True):
n = np.random.choice(a=aa, p=np.ndarray.flatten(p))
if (_xx[n] + new_size < self.masks[batch_idx].shape[0]
and _yy[n] + new_size <
self.masks[batch_idx].shape[1]):
break
else:
while (True):
n = np.random.choice(a=aa, p=np.ndarray.flatten(p))
if (_xx[n] + self.patch_size <
self.masks[batch_idx].shape[0]
and _yy[n] + self.patch_size <
self.masks[batch_idx].shape[1]):
break
# Choosing patch
if self.balanced:
if self.full_augment:
yy = self.masks[batch_idx, _xx[n]:_xx[n] + new_size,
_yy[n]:_yy[n] + new_size, :]
xx = self.images[batch_idx, _xx[n]:_xx[n] + new_size,
_yy[n]:_yy[n] + new_size, :]
else:
yy = self.masks[batch_idx,
_xx[n]:_xx[n] + self.patch_size,
_yy[n]:_yy[n] + self.patch_size, :]
xx = self.images[batch_idx,
_xx[n]:_xx[n] + self.patch_size,
_yy[n]:_yy[n] + self.patch_size, :]
else:
i = np.random.choice(a=self.images.shape[1] - new_size - 1)
j = np.random.choice(a=self.images.shape[2] - new_size - 1)
if self.full_augment:
xx = self.images[batch_idx, i:i + new_size,
j:j + new_size, :]
yy = self.masks[batch_idx, i:i + new_size,
j:j + new_size, :]
else:
xx = self.images[batch_idx, i:i + self.patch_size,
j:j + self.patch_size, :]
yy = self.masks[batch_idx, i:i + self.patch_size,
j:j + self.patch_size, :]
xx = Image.fromarray((255 * xx).astype(np.uint8))
yy = Image.fromarray((np.argmax(yy, axis=2)).astype(np.uint8))
if self.full_augment:
# Rotate
xx = xx.rotate(angle=angle, resample=Image.BICUBIC)
yy = yy.rotate(angle=angle, resample=Image.NEAREST)
# Rescale
new_size1 = int(np.ceil(self.patch_size * etha))
ii = abs(new_size - new_size1) // 2
xx = xx.crop((ii, ii, ii + new_size1, ii + new_size1))
yy = yy.crop((ii, ii, ii + new_size1, ii + new_size1))
xx = xx.resize(size=(self.patch_size, self.patch_size),
resample=Image.BICUBIC)
yy = yy.resize(size=(self.patch_size, self.patch_size),
resample=Image.NEAREST)
# Flip
if randint(0, 1) == 0:
xx, yy = xx.transpose(Image.FLIP_TOP_BOTTOM), yy.transpose(
Image.FLIP_TOP_BOTTOM)
else:
xx, yy = xx.transpose(Image.FLIP_LEFT_RIGHT), yy.transpose(
Image.FLIP_LEFT_RIGHT)
# Rotate 90
rotate_flg = randint(0, 3)
if rotate_flg == 0:
xx, yy = xx.transpose(Image.ROTATE_90), yy.transpose(
Image.ROTATE_90)
elif rotate_flg == 1:
xx, yy = xx.transpose(Image.ROTATE_180), yy.transpose(
Image.ROTATE_180)
elif rotate_flg == 2:
xx, yy = xx.transpose(Image.ROTATE_270), yy.transpose(
Image.ROTATE_270)
yy = to_categorical(np.array(yy).astype(np.uint8),
num_classes=len(classes_mask.keys()))
batch_x.append(np.array(xx).astype(np.float32) / 255)
batch_y.append(yy)
unique, counts = np.unique(np.argmax(yy, axis=2),
return_counts=True)
aaa = dict(zip(unique, counts))
for class_n in range(len(classes_mask)):
if class_n in aaa:
self.stat[classes_mask[class_n]] += aaa[class_n]
batch_x = np.stack(batch_x)
batch_y = np.stack(batch_y)
yield (batch_x, batch_y)