def test(image_size=(512, 512, 3),num_class=3):
PSPNet = pspnet(image_size = image_size,num_class=num_class)
PSPNet.load()
for flag in range(500):
print(str(flag).zfill(5))
image = Image.open("../marine_data/11/images/"+str(flag+1).zfill(5)+".jpg")
image,label = preprocess(image)
plt.subplot(1, 2, 1)
plt.imshow(np.array(image))
prediction = PSPNet.predict(image/255)
result = np.argmax(prediction[0,:,:,:],-1)
plt.subplot(1, 2, 2)
plt.imshow(result)
plt.pause(0.01)
plt.clf()
def _eval(dataset,image_size=(512, 512, 3),num_class=3):
images, truths = dataset.eval_data(batch_size=8,image_size=(512, 512, 3),labels=num_class)
print(truths.shape)
PSPNet = pspnet(image_size = image_size,num_class=num_class)
PSPNet.load()
mean_acc=0
mean_mIoU = 0
for i in range(images.shape[0]):
prediction = PSPNet.predict(images[i,:,:,:])
metric = SegmentationMetric(num_class)
prediction = np.argmax(prediction[0,:,:,:],-1)
plt.subplot(1,2,1)
plt.imshow(images[i,:,:,:])
plt.subplot(1,2,2)
plt.imshow(prediction)
# plt.show()
truth = np.argmax(truths[i,:,:,:],-1)
metric.addBatch(prediction, truth)
acc = metric.pixelAccuracy()
mIoU = metric.meanIntersectionOverUnion()
mean_acc+=acc
mean_mIoU+=mIoU
print(mean_acc/images.shape[0], mean_mIoU/images.shape[0])
def test(args):
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args.model=='td4-psp18':
path_num = 4
vid_seq = cityscapesLoader(img_path=args.img_path,in_size=(769,1537))
vid_seq.load_frames()
model = td4_psp18.td4_psp18(nclass=19,path_num=path_num,model_path=args._td4_psp18_path)
elif args.model=='td2-psp50':
path_num = 2
vid_seq = cityscapesLoader(img_path=args.img_path,in_size=(769,1537))
vid_seq.load_frames()
model = td2_psp50.td2_psp50(nclass=19,path_num=path_num,model_path=args._td2_psp50_path)
elif args.model=='psp101':
path_num = 1
vid_seq = cityscapesLoader(img_path=args.img_path,in_size=(769,1537))
vid_seq.load_frames()
model = pspnet.pspnet(nclass=19,model_path=args._psp101_path)
model.eval()
model.to(device)
timer = 0.0
with torch.no_grad():
for i, (image, img_name, folder, ori_size) in enumerate(vid_seq.data):
image = image.to(device)
torch.cuda.synchronize()
start_time = timeit.default_timer()
output = model(image, pos_id=i%path_num)
torch.cuda.synchronize()
elapsed_time = timeit.default_timer() - start_time
if i > 5:
timer += elapsed_time
pred = np.squeeze(output.data.max(1)[1].cpu().numpy(), axis=0)
pred = pred.astype(np.int8)
pred = cv2.resize(pred, (ori_size[0]//4,ori_size[1]//4), interpolation=cv2.INTER_NEAREST)
decoded = vid_seq.decode_segmap(pred)
save_dir = os.path.join(args.output_path,folder)
res_path = os.path.join(save_dir,img_name)
if not os.path.exists(save_dir):
os.mkdir(save_dir)
imageio.imwrite(res_path, decoded.astype(np.uint8))
cv2.namedWindow("Image")
cv2.imshow("Image", decoded.astype(np.uint8))
cv2.waitKey(1)
print(" Frame {0:2d} RunningTime/Latency={1:3.5f} s".format(i + 1, elapsed_time))
print("---------------------")
print(" Model: {0:s}".format(args.model))
print(" Average RunningTime/Latency={0:3.5f} s".format(timer/(i-5)))
print("---------------------")
def train(md,image_size=(512, 512, 3),num_class=12):
PSPNet = pspnet(image_size = image_size,num_class=num_class)
PSPNet.batch_generator = md.BatchGenerator(batch_size=2, image_size=(512, 512, 3), labels=num_class)
PSPNet.train(epochs=10, steps_per_epoch=500)
PSPNet.save()