create_dataset.py 文件源码

def create_dataset(path, dest, dist, transpose, rgb, resize_factor):
    cap = cv2.VideoCapture(path)
    length = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
    width  = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))

    fac = 3 if rgb else 1

    rwidth = int(width * resize_factor * fac)
    rheight = int(height * resize_factor * fac)

    frames = 0
    I = np.zeros([rwidth*rheight * fac, length])
    Phi = np.zeros([3, length])

    while True:
        (grabbed, frame) = cap.read()

        if not grabbed: break       

        p, r = detect_ball(frame)

        if p != None:
            Phi[:,frames] = dist(p, r)[0:3]
        else:
            Phi[:,frames] = np.nan

        img = frame

        if not rgb:
            img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        img = imutils.resize(img, width=rwidth, height=rheight, inter=cv2.INTER_AREA)

        I[:,frames] = img.T.ravel() if transpose else img.ravel()

        frames += 1

    print("Got {0} frames".format(frames))

    cap.release()
    cv2.destroyAllWindows()

    data = {'I' : I, 'h' : rheight, 'w' : rwidth, 'Phi' : Phi}
    scipy.io.savemat(dest, data , do_compression=True)