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')
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)
cmap = copy.copy(plt.cm.get_cmap("gist_rainbow", 20))
cmap.set_under(color="black")
print(f"Generating predictions for {len(images)} images...")
if args.generate_images:
for img_path, pred in zip(image_paths, preds):
np.savez(f"{img_path.split('.jpg')[0]}_PRED.npz", pred)
pred[pred == 18] = -1
plt.imshow(pred, cmap=cmap, interpolation='none', vmin=0, vmax=19)
plt.colorbar(ticks=range(0, 20))
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
for i in range(H):
for j in range(W):
if pred_mask[i][j][0] <= 0.5:
img[i][j] = (1 - alpha) * img[i][j] * 255 + alpha * np.array(
[0, 0, 255])
else:
