def main():
input_size = (320, 480, 3)
classes = 21
weight_path = './checkpoint/final_stage.h5'
filename = '../seg_train_images/seg_train_images/train_1118.jpg'
color_index = [[0, 0, 255], [193, 214, 0], [180, 0, 129], [255, 121, 166],
[255, 0, 0], [65, 166, 1], [208, 149, 1], [255, 255, 0],
[255, 134, 0], [0, 152, 225], [0, 203, 151], [85, 255, 50],
[92, 136, 125], [69, 47, 142], [136, 45, 66], [0, 255, 255],
[215, 0, 255], [180, 131, 135], [81, 99, 0], [86, 62, 67]]
net = Unet(input_size, classes)
net.load_weights(weight_path)
img = Image.open(filename).resize((input_size[1], input_size[0]))
img = np.asarray(img, dtype='float32')
img = img / 255.0
img = np.expand_dims(img, axis=0)
pred = net.predict(img)
b, w, h, c = pred.shape
res = []
for i in range(w):
tmp = []
for j in range(h):
tmp.append(color_index[np.argmax(pred[0, i, j, :])])
res.append(tmp)
img = Image.fromarray(np.asarray(res, dtype='uint8'))
img.save('result.png')
return x, y
x_all_all = np.load('img.npy')
y_all_all = np.load('label.npy')
print(y_all_all.shape, x_all_all.shape)
train_db = tf.data.Dataset.from_tensor_slices((x_all_all, y_all_all))
train_db = train_db.shuffle(200).batch(1)
test_db = tf.data.Dataset.from_tensor_slices(
(x_all_all[0:50], y_all_all[0:50]))
test_db = test_db.batch(1)
next(iter(train_db))
next(iter(test_db))
print(y_all_all.shape, x_all_all.shape)
model = Unet(num_class=2)
model.load_weights('C:/stone_segmataion/callbacks/slate_models.h5')
def focal(y_ture, y_pre):
loss = -tf.reduce_mean(
y_ture * tf.multiply(1 - y_pre, 1 - y_pre) *
tf.math.log(tf.clip_by_value(y_pre, 1e-10, 1.0))) * 0.25
return loss
def meanIOU(y_ture, y_pre):
y_pre = y_pre[0]
y_ture = y_ture[0]
y_ture = tf.argmax(y_ture, axis=2)
y_pre = tf.argmax(y_pre, axis=2)
parser.add_argument(
'-g',
'--generate_images',
help=
'saves representative image of mask in addition to the mask prediction',
action="store_true")
args = parser.parse_args()
assert os.path.isfile(args.weights), "Could not find weights file!"
assert os.path.isdir(args.images), "Could not find images directory!"
unet = Unet(DESIRED_SIZE,
DESIRED_SIZE,
nclasses=NUM_CLASSES,
filters=UNET_FILTERS)
unet.load_weights(args.weights)
unet.compile(optimizer='adam',
loss=SparseCategoricalFocalLoss(gamma=2),
metrics=['accuracy'])
image_paths = glob.glob(args.images + '/*.jpg')
images = [
x[0] for x in [
get_resized_image_and_mask_label(cv2.imread(p), NEGATIVE)
for p in image_paths
]
]
mpreds = unet.predict(np.array(images, dtype=np.uint8))
preds = maskpred(mpreds)
for (img, mask) in train_generator:
img = img / 255.
mask = mask / 255.
mask[mask > 0.5] = 1
mask[mask <= 0.5] = 0
yield (img, mask)
model = Unet(1, image_size=512)
trainset = DataGenerator("dataset/train", batch_size=2)
model.fit_generator(trainset, steps_per_epoch=100, epochs=1)
model.save_weights("model.h5")
testSet = DataGenerator("dataset/test", batch_size=1)
alpha = 0.3
model.load_weights("model.h5")
if not os.path.exists("./results"): os.mkdir("./results")
for idx, (img, mask) in enumerate(testSet):
oring_img = img[0]
# 开始用模型进行预测
pred_mask = model.predict(img)[0]
pred_mask[pred_mask > 0.5] = 1
pred_mask[pred_mask <= 0.5] = 0
# 如果这里展示的预测结果一片黑，请调整lr，同时注意图片的深度是否为8
# cv2.imshow('pred_mask', pred_mask)
# cv2.waitKey()
img = cv2.cvtColor(img[0], cv2.COLOR_GRAY2RGB)
H, W, C = img.shape
