def visualise_optical_flow(first_frame_index,
second_frame_index,
video_path: str,
workspace_path: str = "workspace/"):
"""Visualise the optical flow between two video frames from a generated optical flow dataset."""
video_name = Path(video_path).stem
base_path = os.path.join(workspace_path, video_name)
frame_pair = (first_frame_index, second_frame_index)
optical_flow_path = os.path.join(base_path, "dataset", "optical_flow",
"{:04d}-{:04d}.pkl".format(*frame_pair))
with open(optical_flow_path, 'rb') as f:
optical_flow, valid_mask = pickle.load(f)
frame_path_fmt = os.path.join(base_path, "frames", "{:04d}.png")
frame = cv2.cvtColor(cv2.imread(frame_path_fmt.format(frame_pair[0])),
cv2.COLOR_BGR2RGB)
frame2 = cv2.cvtColor(cv2.imread(frame_path_fmt.format(frame_pair[1])),
cv2.COLOR_BGR2RGB)
optical_flow_viz = flow2img(optical_flow.transpose((1, 2, 0)))
optical_flow_viz *= valid_mask.transpose((1, 2, 0))
optical_flow_viz[optical_flow_viz == 0] = 255
gs = gridspec.GridSpec(4, 4)
ax1 = plt.subplot(gs[:2, :2])
ax1.imshow(frame)
ax1.set_title("First Frame")
ax2 = plt.subplot(gs[:2, 2:])
ax2.imshow(frame2)
ax2.set_title("Second Frame")
ax3 = plt.subplot(gs[2:, 1:3])
ax3.imshow(optical_flow_viz)
X, Y = np.meshgrid(np.arange(0, frame.shape[1], 1),
np.arange(0, frame.shape[0], 1))
U, V = optical_flow * valid_mask
every_n = 32
ax3.quiver(X[::every_n, ::every_n], Y[::every_n, ::every_n],
U[::every_n, ::every_n], V[::every_n, ::every_n])
ax3.set_title("Optical Flow")
plt.tight_layout()
plt.show()
def forward(self, input1, input2):
if self.cropper is None or self.image_size != input1.shape[:2] or \
self.render_size != [((input1.shape[0]) // 64) * 64, ((input1.shape[1]) // 64) * 64]:
self.image_size = input1.shape[:2]
self.render_size = [((input1.shape[0]) // 64) * 64,
((input1.shape[1]) // 64) * 64]
self.cropper = StaticCenterCrop(self.image_size, self.render_size)
images = [input1, input2]
images = list(map(self.cropper, images))
images = torch.stack(images).transpose(0,
1).transpose(0,
2).transpose(0, 3)
images = images.unsqueeze(0)
result = self.net(images).squeeze()
flow = result.transpose(1, 2).transpose(0, 2)
result = torch.tensor(flow2img(flow.cpu().numpy()),
dtype=torch.float32).cuda()
return result
net.load_state_dict(dict["state_dict"])
# end = time.time()
# print('NNloading time : ' + str(end - start))
images = [pim1, pim2]
images = np.array(images).transpose(3, 0, 1, 2)
im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()
# start = time.time()
flow = net(im)
# end = time.time()
flow = flow.squeeze()
# print('NNcomputing time : ' + str(end - start))
flow = flow.data.cpu().numpy().transpose(1, 2, 0)
# estimated_pim2N = warp_flow(pim1,data) # 这个warp是把1warp到2
estimated_pim2P = image_warp(pim2, flow,
mode='nearest') # 这个warp是把2warp回1
# cv2.imwrite(estimated_pim2pathN, estimated_pim2N)
cv2.imwrite(estimated_pim2pathP, estimated_pim2P)
img = flow2img(flow) # data是flow array 这里直接得到bgr图片
img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) # 下面四行 JINJING 改为用HV eoncde 光流
img[:, :, 2] = img[:, :, 1]
img[:, :, 1] = 255 # 此时的img为HSV的
img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
# 写入之前最好设置一个清空文件夹操作
cv2.imwrite(RGBflowpath, img) #如果想保存光流图片，以备后面调用flames2video得到光流video.
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--dataset_name', type=str, default='-1')
parser.add_argument('--thresh_hold', type=int, default=10)
args = parser.parse_args()
dataset_name = args.dataset_name
thresh_hold = args.thresh_hold
optflow_path = os.path.join('./store',
'{}_optflow.hdf5'.format(dataset_name))
optvideo_dir = './store/optvideos'
if not os.path.exists(optvideo_dir):
os.makedirs(optvideo_dir)
print('=' * 20, 'construct optvideos begin', '=' * 20)
with h5py.File(optflow_path, 'r') as f:
vid_list = list(f.keys())[:thresh_hold]
for vid in tqdm.tqdm(vid_list):
optflows = f[vid][()]
optflows_len = optflows.shape[0]
optflows_save_path = os.path.join(optvideo_dir,
'{}.avi'.format(vid))
writer_handle = cv2.VideoWriter(optflows_save_path,
cv2.VideoWriter_fourcc(*'MJPG'), 5,
(256, 256))
for i in range(optflows_len):
single_flow = optflows[i, ...].transpose(1, 2, 0)
im = flow2img(single_flow)
writer_handle.write(im)
writer_handle.release()
print('=' * 20, 'construct optvideos end', '=' * 20)
images = [pim1, pim2]
images = np.array(images).transpose(3, 0, 1, 2)
im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()
# im = torch.cat([im, im], dim=0)
# process the image pair to obtian the flow
print(net._get_name())
# print(im.shape)
print('im.shape:{}'.format(im.shape))
result = net(im).squeeze()
print('result.shape:{}'.format(result.shape))
# save flow, I reference the code in scripts/run-flownet.py in flownet2-caffe project
def writeFlow(name, flow):
f = open(name, 'wb')
f.write('PIEH'.encode('utf-8'))
np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)
flow = flow.astype(np.float32)
flow.tofile(f)
f.flush()
f.close()
data = result.data.cpu().numpy().transpose(1, 2, 0)
print(data.shape)
temp = flow2img(data)
temp = Image.fromarray(temp)
temp.show()
# writeFlow("/home/hjj/flownet2-master/data/FlyingChairs_examples/0000007-img.flo", data)
ret, frame = cap.read()
frame = cv.flip(frame, 1)
if not ret:
print("Can't receive frame (stream end?). Exiting ...")
break
curr_image = padding(
transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((120, 160)),
transforms.ToTensor()
])(frame)).unsqueeze(1).unsqueeze(0)
if prev_image == None:
prev_image = curr_image
inputs = torch.cat([prev_image, curr_image], dim=2).cuda()
start = time.time()
with torch.no_grad():
output = model(inputs).cpu().numpy()
end = time.time()
result = flow_utils.flow2img(output[0].transpose(1, 2, 0))
prev_image = curr_image
# print (1 / (end - start))
# write the flipped frame
cv.imshow('origin', frame)
cv.imshow('result', result)
if cv.waitKey(1) == ord('q'):
break
# Release everything if job is finished
cap.release()
cv.destroyAllWindows()
def FN(videopath, sigDIR, String0, range, gridnumX, gridnumY):
crop_size = (324, 274)
size = (854, 480)
PATHMAG = sigDIR + "FN" + "\size" + String0 + "\\" + "MAG" + "\\" + "SIGS_" + "FN" + "_" + String0 + "_" + str(
range[0]) + "_" + str(range[1]) + "mag.npy"
PATHANG = sigDIR + "FN" + "\size" + String0 + "\\" + "ANG" + "\\" + "SIGS_" + "FN" + "_" + String0 + "_" + str(
range[0]) + "_" + str(range[1]) + "ang.npy"
cap = cv2.VideoCapture(videopath)
_, frame1 = cap.read()
ret, frame1 = cap.read()
prvs = frame1
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
lower = fps * range[0]
upper = fps * range[1]
framenum = upper - lower
gridWidth = int(
width / gridnumX
) # int 0.6 = 0; 320 required to be zhengchu, but if not, then directly negelact the boundary
gridHeight = int(height / gridnumY) # 240
SIGS_ang = np.zeros((framenum, gridnumY, gridnumX))
SIGS_mag = np.zeros((framenum, gridnumY, gridnumX))
start = time.time()
parser = argparse.ArgumentParser()
parser.add_argument(
'--fp16',
action='store_true',
help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).')
parser.add_argument("--rgb_max", type=float, default=255.)
args = parser.parse_args()
# initial a Net
# net = FlowNet2C(args).cuda()
# net = FlowNet2S(args).cuda()
# net = FlowNet2CS(args).cuda()
net = FlowNet2(args).cuda()
# net = FlowNet2(args).cuda()
# load the state_dict
dict = torch.load("checkpoints/FlowNet2_checkpoint.pth.tar") # flownet2
net.load_state_dict(dict["state_dict"])
end = time.time()
print(end - start)
print('///////')
timepoint = 0
i = 1
# Till you scan the video
start = time.time()
while (i):
ret, frame2 = cap.read()
#
if ret != True:
break
next = frame2
# next = cv2.cvtColor(next, cv2.COLOR_BGR2GRAY)
# next = cv2.GaussianBlur(next, (21, 21), 0)
# start = time.time()
pim1 = cv2.resize(prvs, crop_size, interpolation=cv2.INTER_LINEAR)
pim2 = cv2.resize(next, crop_size, interpolation=cv2.INTER_LINEAR)
images = [pim1, pim2]
images = np.array(images).transpose(3, 0, 1, 2)
im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()
# print (cap.get(3))
start = time.time()
result = net(im).squeeze()
# result = net(im)
end = time.time()
print(end - start)
data = result.data.cpu().numpy().transpose(1, 2, 0)
img = flow2img(data)
img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) # JINJING 改为用HV eoncde 光流
img[:, :, 2] = img[:, :, 1]
img[:, :, 1] = 255
next_ang = img[..., 0].reshape(gridnumY, gridHeight, gridnumX,
gridWidth)
next_mag = img[..., 2].reshape(gridnumY, gridHeight, gridnumX,
gridWidth)
SIGS_ang[timepoint] = next_ang.mean(axis=(1, 3))
SIGS_mag[timepoint] = next_mag.mean(axis=(1, 3))
img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
double_representation = cv2.addWeighted(img, 0.5, pim2, 0.5, 0)
cv2.imshow("window", double_representation)
cv2.imshow("window", img)
kk = cv2.waitKey(20) & 0xff # 实时可视化光流图片，（人自己写好了flow2img函数）
# Press 'e' to exit the video
if kk == ord('e'):
break
# plt.imshow(img)
# plt.show()
i = i + 1
timepoint = timepoint + 1
prvs = next
# pic2video(save_path,crop_size)
print(np.array(SIGS_ang).size)
print(np.array(SIGS_mag).size)
np.save(PATHANG, SIGS_ang)
np.save(PATHMAG, SIGS_mag)
pim1 = cv2.resize(prvs, crop_size, interpolation=cv2.INTER_AREA)
pim2 = cv2.resize(next, crop_size, interpolation=cv2.INTER_AREA)
images = [pim1, pim2]
# images = [prvs,next]
images = np.array(images).transpose(3, 0, 1, 2)
im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()
start = time.time()
result = net(im).squeeze()
end = time.time()
print(end - start)
data = result.data.cpu().numpy().transpose(1, 2, 0)
img = flow2img(data)
# 写入之前最好设置一个清空文件夹操作
cv2.imwrite(save_path + str(i) + '.png',
img) #如果想保存光流图片，以备后面调用flames2video得到光流video.
cv2.imshow("window", img)
kk = cv2.waitKey(20) & 0xff #实时可视化光流图片，（人自己写好了flow2img函数）
# Press 'e' to exit the video
if kk == ord('e'):
break
# plt.imshow(img)
# plt.show()
i = i + 1
prvs = next
# pic2video(path,crop_size)
# pic2video(path,crop_size)