Untitled

import PyNvCodec as nvc
import numpy as np
import av  # PyAV library
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import glob

def process_video(input_video_path, output_video_path, gpu_id=0, bitrate='1M'):
    try:
        if not os.path.isfile(input_video_path):
            print(f"Input file does not exist: {input_video_path}")
            return

        print(f"Processing {input_video_path} on GPU {gpu_id}...")

        # Initialize the decoder
        nvDec = nvc.PyNvDecoder(input_video_path, gpu_id)

        # Get video properties
        width = nvDec.Width()
        height = nvDec.Height()
        fps_in = nvDec.Framerate()
        fps_out = fps_in
        codec = 'h264'

        # Optionally reduce resolution
        new_width = (width // 4)   # Ensure width is even
        new_height = (height // 4)   # Ensure height is even
        print(f"Original resolution: {width}x{height}, Output resolution: {new_width}x{new_height}")

        # Use PyAV to get input time base
        input_container = av.open(input_video_path)
        input_stream = input_container.streams.video[0]
        input_time_base = input_stream.time_base
        input_container.close()

        # Initialize the encoder with desired settings
        enc_params = {
            'preset': 'P1',        # Fastest preset
            'codec': codec,
            's': f'{new_width}x{new_height}',  # Force encoder scaling
            'bitrate': bitrate,
            'fps': str(int(fps_out)),
            'profile': 'baseline',
            'rc': 'cbr',
            'gop': '500',
        }
        nvEnc = nvc.PyNvEncoder(enc_params, gpu_id)

        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_video_path), exist_ok=True)

        # Initialize PyAV container and stream
        output_container = av.open(output_video_path, mode='w')
        av_codec = 'h264_nvenc' if codec == 'h264' else 'hevc_nvenc'
        stream = output_container.add_stream(av_codec, rate=int(fps_out))
        stream.width = new_width
        stream.height = new_height
        stream.pix_fmt = 'yuv420p'
        stream.time_base = input_time_base

        # Initialize packet buffer and packet data
        packet = np.empty(0, dtype=np.uint8)
        pkt_data = nvc.PacketData()

        # Create PySurfaceResizer to resize frames on GPU
        resizer = nvc.PySurfaceResizer(new_width, new_height, nvc.PixelFormat.NV12, gpu_id)

        # Processing loop
        last_output_pts = None
        while True:
            # Decode a single surface with packet data
            surface = nvDec.DecodeSingleSurface(pkt_data)
            if surface.Empty():
                break

            # Resize surface on GPU
            resized_surface = resizer.Execute(surface)
            if resized_surface.Empty():
                print("Failed to resize surface.")
                continue

            # Confirm resized dimensions
            #print(f"Resized surface dimensions: {resized_surface.Width()}x{resized_surface.Height()}")

            # Encode the resized surface
            success = nvEnc.EncodeSingleSurface(resized_surface, packet)
            if not success:
                continue

            if packet.size > 0:
                output_pts = pkt_data.pts
                if last_output_pts is not None and output_pts <= last_output_pts:
                    output_pts = last_output_pts + 1
                last_output_pts = output_pts

                av_packet = av.packet.Packet(packet.tobytes())
                av_packet.stream = stream
                av_packet.pts = output_pts
                av_packet.dts = output_pts
                output_container.mux(av_packet)

        # Flush the encoder
        while True:
            success = nvEnc.FlushSinglePacket(packet)
            if not success or packet.size == 0:
                break

            output_pts = (last_output_pts + 1) if last_output_pts is not None else 0
            last_output_pts = output_pts

            av_packet = av.packet.Packet(packet.tobytes())
            av_packet.stream = stream
            av_packet.pts = output_pts
            av_packet.dts = output_pts
            output_container.mux(av_packet)

        # Close the output container
        output_container.close()
        print(f"Processed: {input_video_path} -> {output_video_path}")

    except Exception as e:
        print(f"An error occurred while processing {input_video_path}: {e}")

def main():
    input_dir = 'in_movies/stream'
    output_dir = 'out_movies/compressed_stream'
    os.makedirs(output_dir, exist_ok=True)

    # List of video files to process
    video_files = glob.glob(os.path.join(input_dir, '*'))
    print(f"Found {len(video_files)} video files to process:")
    for vf in video_files:
        print(f" - {vf}")

    start_time = time.time()

    # Specify available GPU IDs for cycling
    gpu_ids = [0, 1]  # Change based on your system's available GPUs
    max_workers = len(gpu_ids)  # One worker per GPU

    # Parallel processing with ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        for index, input_file in enumerate(video_files):
            filename = os.path.basename(input_file)
            output_file = os.path.join(output_dir, filename)
            gpu_id = gpu_ids[index % len(gpu_ids)]  # Cycle through GPUs
            futures.append(executor.submit(process_video, input_file, output_file, gpu_id))

        # Track completion
        for future in as_completed(futures):
            try:
                future.result()
                print("Processing completed for a video.")
            except Exception as e:
                print(f"An error occurred: {e}")

    total_time = time.time() - start_time
    print(f"All files processed in {total_time:.2f} seconds.")

if __name__ == '__main__':
    main()