def inference(model, args):
image_folder = "/media/pandongwei/ExtremeSSD/work_relative/extract_img/2020.10.16_1/"
video_save_path = "/home/pandongwei/work_repository/erfnet_pytorch/eval/"
# parameters about saving video
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(video_save_path + 'output_new.avi', fourcc, 10.0, (640, 480))
cuda = True
model.eval()
paths = []
for root, dirs, files in os.walk(image_folder, topdown=True):
for file in files:
image_path = os.path.join(image_folder, file)
paths.append(image_path)
paths.sort()
font = cv2.FONT_HERSHEY_SIMPLEX
angle_pre = 0
for i, path in enumerate(paths):
start_time = time.time()
image = cv2.imread(path)
image = (image / 255.).astype(np.float32)
image = ToTensor()(image).unsqueeze(0)
if (cuda):
image = image.cuda()
input = Variable(image)
with torch.no_grad():
output = model(input)
label = output[0].max(0)[1].byte().cpu().data
# label_cityscapes = cityscapes_trainIds2labelIds(label.unsqueeze(0))
label_color = Colorize()(label.unsqueeze(0))
label_save = label_color.numpy()
label_save = label_save.transpose(1, 2, 0)
# 加上路径规划
label_save, angle = perception_to_angle(label_save, angle_pre)
# 加一个滤波以防止角度突然跳变 TODO
if abs(angle - angle_pre) > 10:
angle = angle_pre
angle_pre = angle
# label_save.save(filenameSave)
image = image.cpu().numpy().squeeze(axis=0).transpose(1, 2, 0)
image = (image * 255).astype(np.uint8)
output = cv2.addWeighted(image, 0.5, label_save, 0.5, 0)
cv2.putText(output,str(round(angle,3)),(50,50),cv2.FONT_HERSHEY_SIMPLEX,2,(0,0,0),2)
#output = np.hstack([label_save, image])
out.write(output)
print(i, " time: %.2f s" % (time.time() - start_time))
out.release()
def sr(im, scale):
im = im.convert('YCbCr')
im, cb, cr = im.split()
h, w = im.size
im = ToTensor()(im)
im = Variable(im).view(1, -1, w, h)
im = im.cuda()
with torch.no_grad():
im = espcn(im)
im = torch.clamp(im, 0., 1.)
im = im.cpu()
im = im.data[0]
im = ToPILImage()(im)
cb = cb.resize(im.size, Image.BICUBIC)
cr = cr.resize(im.size, Image.BICUBIC)
im = Image.merge('YCbCr', [im, cb, cr])
im = im.convert('RGB')
return im
def save_video_and_path_planning(model_geoMat, model_freiburgForest, cfg):
image_folder = cfg["image_folder_raw"]
video_save_path = "/home/pandongwei/work_repository/erfnet_pytorch/eval/"
# parameters about saving video
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(video_save_path + 'output.avi', fourcc, 10.0,
(640, 480))
cuda = cfg['cuda']
model_geoMat.eval()
model_freiburgForest.eval()
coef = np.ones([512, 1024, 5]).astype(np.float32)
traversability = (0, 1.0, 0.6, 0.8, -1)
for i in range(5):
coef[:, :, i] = coef[:, :, i] * traversability[i]
paths = []
for root, dirs, files in os.walk(image_folder, topdown=True):
for file in files:
image_path = os.path.join(image_folder, file)
paths.append(image_path)
paths.sort()
font = cv2.FONT_HERSHEY_SIMPLEX
for i, path in enumerate(paths):
start_time = time.time()
image = cv2.imread(path)
image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR)
image = image / 255.
image = ToTensor()(image).unsqueeze(0)
if (cuda):
image = image.cuda()
input = Variable(image)
with torch.no_grad():
outputs_1 = model_freiburgForest(input)
outputs_2 = model_geoMat(input)
outputs_2 = outputs_2[0] # TODO
outputs_1 = outputs_1[0].max(0)[1].byte().cpu().data.unsqueeze(0)
# print(outputs_1.shape)
outputs_1 = outputs_1.numpy().transpose(1, 2, 0).astype(np.int64)
# print(outputs_1.shape)
outputs_1 = np.take(coef, outputs_1)
outputs_2 = outputs_2[0, 0, :, :].cpu().data.unsqueeze(0)
# print(outputs_2.shape)
outputs_2 = outputs_2.numpy().transpose(1, 2, 0)
# 加上momenton来稳定算法 TODO
if i == 0:
momenton = 0.9
output_pre = outputs_2.copy()
else:
outputs_2 = (momenton * output_pre + (1 - momenton) * outputs_2)
output_pre = outputs_2.copy()
outputs_combine = (outputs_1 * outputs_2 * 255).astype(np.uint8)
outputs_combine_color = Colorize()(outputs_combine).transpose(
1, 2, 0).astype(np.uint8)
outputs_combine_color = cv2.cvtColor(outputs_combine_color,
cv2.COLOR_RGB2BGR)
image = image.cpu().numpy()
image = (image[0].transpose(1, 2, 0) * 255).astype(np.uint8)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# print(outputs_combine.shape)
# output = np.hstack([outputs_combine_color, image])
# print(" post process time: %.2f s" % (time.time() - start_time))
# start_time = time.time()
outputs_combine = outputs_combine[:, :, 0]
output = path_planning(outputs_combine, outputs_combine_color)
# print(" path_planning time: %.2f s" % (time.time() - start_time))
# start_time = time.time()
output = np.hstack([output, image])
cv2.putText(output,
str(round(1 / (time.time() - start_time), 1)) + " Hz",
(50, 50), font, 1, (0, 0, 255), 2)
out.write(image)
print(i, " time: %.2f s" % (time.time() - start_time))
out.release()
parser.add_argument("--test", type=str, help="path to load test images")
opt = parser.parse_args()
print(opt)
net = Net(opt.rb)
net.load_state_dict(torch.load(opt.checkpoint)['state_dict'])
net.eval()
net = nn.DataParallel(net, device_ids=[0, 1, 2, 3]).cuda()
print(net)
images = utils.load_all_image(opt.test)
for im_path in tqdm(images):
filename = im_path.split('/')[-1]
print(filename)
im = Image.open(im_path)
h, w = im.size
print(h, w)
im = ToTensor()(im)
im = Variable(im).view(1, -1, w, h)
im = im.cuda()
with torch.no_grad():
im = net(im)
im = torch.clamp(im, 0., 1.)
im = im.cpu()
im = im.data[0]
im = ToPILImage()(im)
im.save('output/%s' % filename)
def compute_psnr(ref_im, res_im):
ref_img = scipy.misc.fromimage(ref_im).astype(float) / 255.0
res_img = scipy.misc.fromimage(res_im).astype(float) / 255.0
squared_error = np.square(ref_img - res_img)
mse = np.mean(squared_error)
psnr = 10 * np.log10(1.0 / mse)
return psnr
# print(images)
for im_path in tqdm(images):
filename = im_path.split('/')[-1]
im2_filename = filename.split('_')[0] + '.png'
im = Image.open(im_path)
im2_path = str_label + im2_filename
im2 = Image.open(im2_path)
h, w = im.size
im = ToTensor()(im)
im = im.unsqueeze(0).cuda()
im = Variable(im, volatile=True)
im = net(im)
im = torch.clamp(im, 0., 1.)
im = ToPILImage()(im.cpu().data[0])
ssims.append(compare_ssim(np.array(im), np.array(im2), multichannel=True))
psnrs2.append(compare_psnr(np.array(im), np.array(im2)))
print(np.mean(ssims), np.mean(psnrs2))
def visualize_probabilities(y_softmax=None,
y=None,
global_patch=None,
grid_size_in_global=None,
probability_mask=None,
return_mode="buf",
axs=None):
""""""
color = torch.ones(3, int(2 * grid_size_in_global + 1),
int(2 * grid_size_in_global + 1)) * torch.tensor(
[1., 0., 0.]).view(-1, 1, 1)
recon_img = np.clip(y[0][0].cpu().numpy(), 0, 1)
recon_img = Image.fromarray(recon_img)
recon_img = ToTensor()(np.array(
recon_img.resize((int(2 * grid_size_in_global + 1),
int(2 * grid_size_in_global + 1)))))
recon_img = torch.cat((color, recon_img.cpu()), dim=0).permute(1, 2, 0)
softmax_img = np.clip(y_softmax[0][0].cpu().numpy(), 0, 1)
softmax_img /= np.max(softmax_img)
softmax_img = Image.fromarray(softmax_img)
softmax_img = ToTensor()(np.array(
softmax_img.resize((int(2 * grid_size_in_global + 1),
int(2 * grid_size_in_global + 1)))))
softmax_img = torch.cat((color, softmax_img.cpu()), dim=0).permute(1, 2, 0)
gt_img = np.clip(probability_mask, 0, 1)
gt_img = Image.fromarray(gt_img)
gt_img = ToTensor()(np.array(
gt_img.resize((int(2 * grid_size_in_global + 1),
int(2 * grid_size_in_global + 1)))))
gt_img = torch.cat((color, gt_img.cpu()), dim=0).permute(1, 2, 0)
real_img = np.clip(global_patch, 0, 1)
if axs is None:
fig, axs = plt.subplots(1, 3, figsize=(9, 3), sharey=True)
axs[0].imshow(np.transpose(real_img, (1, 2, 0)),
interpolation='nearest') #, origin='upper')
axs[0].imshow(softmax_img, interpolation='nearest')
axs[0].set_title("Probability Map")
axs[0].axis('off')
axs[1].imshow(np.transpose(real_img, (1, 2, 0)),
interpolation='nearest') #, origin='upper')
axs[1].imshow(
recon_img,
interpolation='nearest',
)
axs[1].set_title("Goal Realisation")
axs[1].axis('off')
axs[2].imshow(np.transpose(real_img, (1, 2, 0)),
interpolation='nearest') #, origin='upper')
axs[2].imshow(gt_img, interpolation='nearest')
axs[2].set_title("GT Goal")
axs[2].axis('off')
if return_mode == "buf":
buf = io.BytesIO()
plt.savefig(buf, format='jpeg')
plt.close()
buf.seek(0)
image = PIL.Image.open(buf)
image = ToTensor()(image)
return image
elif return_mode == "ax":
return axs
else:
return fig
