NVIDIA / DeepLearningExamples

State-of-the-Art Deep Learning scripts organized by models - easy to train and deploy with reproducible accuracy and performance on enterprise-grade infrastructure.
13.61k stars 3.24k forks source link

How to write the bitstream obtained from `PyNvVideoCodec` encoding into a video file (e.g., MP4). #1434

Open xingguang12 opened 2 weeks ago

xingguang12 commented 2 weeks ago

`import cv2 import numpy as np

cap = cv2.VideoCapture(0)

while True: ret, frame = cap.read() if not ret: break

height, width, channels = frame.shape
print(f"Width: {width}, Height: {height}")

# Convert to YUV format
yun_image = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)

# Separate Y, U, V channels
y = yun_image[:, :, 0]
u = yun_image[:, :, 1]
v = yun_image[:, :, 2]

# Downsample U and V (subsample by 2 in both width and height)
u = u[::2, ::2]
v = v[::2, ::2]

# Flatten Y, U, V channels
y_flat = y.flatten()
uv_flat = np.dstack((u.flatten(), v.flatten())).flatten()

# Combine Y and UV for NV12 format
nv12_data = np.concatenate((y_flat, uv_flat))
import PyNvVideoCodec as nvc
import numpy as np

encoder = nvc.CreateEncoder(
    640,
    480,
    "NV12",
    True)
frame_size = 640 * 480 * 1.5
bitstream = encoder.Encode(nv12_data)  # encode frame one by one
if cv2.waitKey(1) & 0xFF == ord('q'):
    break

cap.release() cv2.destroyAllWindows()`