import sys
from time import sleep
import cv2 as cv

# Using nvv4l2h264enc
out_gst_pipe = 'appsrc ! queue ! videoconvert ! video/x-raw,format=(string)BGRx ! nvvidconv ! video/x-raw(memory:NVMM),format=(string)NV12 ! nvv4l2h264enc ! h264parse ! qtmux ! filesink location=testvid1.mp4'

# Using nvv4l2h264enc and matroskamux
# out_gst_pipe = 'appsrc ! queue ! videoconvert ! video/x-raw,format=(string)BGRx ! nvvidconv ! video/x-raw(memory:NVMM),format=(string)NV12 ! nvv4l2h264enc ! h264parse ! matroskamux ! filesink location=testvid2.mkv'

# Using omxh264enc
# out_gst_pipe = "appsrc ! videoconvert ! omxh264enc profile=2 preset-level=1 control-rate=2 bitrate=10000000 ! video/x-h264, streamformat=(string)byte-stream ! h264parse ! qtmux ! filesink location=testvid3.mp4"

cc = cv.VideoWriter_fourcc(*'mp4v')
out = cv.VideoWriter(out_gst_pipe, cv.CAP_GSTREAMER, cc, 21, (1280, 960))

if not out.isOpened():
    out.release()
    print('shucks')
    sys.exit(1)


pipe = 'nvarguscamerasrc sensor-id=0 maxperf=1 ! video/x-raw(memory:NVMM),format=(string)NV12,width=(int)3280,height=(int)2464,framerate=(fraction)21/1 ! nvvidconv flip-method=2 ! video/x-raw,width=(int)1280,height=(int)960,format=(string)BGRx ! videoconvert ! video/x-raw,format=(string)BGR ! appsink'
cap = cv.VideoCapture(pipe, cv.CAP_GSTREAMER)
if not cap.isOpened():
    cap.release()
    out.release()
    print('oh, come on...')
    sys.exit(2)


for i in range(21 * 20):
    ret, frame = cap.read()
    if not ret:
        print('no read')
        break
    cv.imshow('frame', frame)
    cv.waitKey(1)
    out.write(frame)

cv.destroyAllWindows()
out.release()
cap.release()