FlowCalculator.py 文件源码

def writeOpticalFlowImage(self, index, optical_flow):
        filename = "flow_" + str(index) + ".png"
        output_path = os.path.join(self.optical_flow_output_directory, filename)

        # create hsv image
        shape_optical_flow = optical_flow.shape[:-1]
        shape_hsv = [shape_optical_flow[0], shape_optical_flow[1], 3]
        hsv = np.zeros(shape_hsv, np.float32)

        # set saturation to 255
        hsv[:,:,1] = 255

        # create colorful illustration of optical flow
        mag, ang = cv2.cartToPolar(optical_flow[:,:,0], optical_flow[:,:,1])
        hsv[:,:,0] = ang*180/np.pi/2
        hsv[:,:,2] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX)
        bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)

        cv2.imwrite(output_path, bgr)