<commit_before>import io
import logging
import math
from itertools import tee
from struct import pack, unpack_from

import av
from av import AVError
from av.codec.context import CodecContext
from av.packet import Packet
from av.video.frame import VideoFrame

from ..contrib.media import frame_from_bgr, frame_to_bgr

logger = logging.getLogger('codec.h264')

PACKET_MAX = 1300

NAL_TYPE_FU_A = 28
NAL_TYPE_STAP_A = 24

NAL_HEADER_SIZE = 1
FU_A_HEADER_SIZE = 2
LENGTH_FIELD_SIZE = 2
STAP_A_HEADER_SIZE = NAL_HEADER_SIZE + LENGTH_FIELD_SIZE


def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


class H264PayloadDescriptor:
    def __init__(self, first_fragment):
        self.first_fragment = first_fragment

    def __repr__(self):
        return 'H264PayloadDescriptor(FF={})'.format(self.first_fragment)

    @classmethod
    def parse(cls, data):
        output = bytes()

        nal_type = data[0] & 0x1f
        if nal_type == NAL_TYPE_FU_A:
            assert len(data) >= FU_A_HEADER_SIZE, 'FU-A NAL units truncated.'

            f_nri = data[0] & (0x80 | 0x60)
            original_nal_type = data[1] & 0x1f
            first_fragment = bool(data[1] & 0x80)

            if first_fragment:
                original_nal_header = bytes([f_nri | original_nal_type])
                output += bytes([0, 0, 0, 1])
                output += original_nal_header

            output += data[2:]

            obj = cls(first_fragment=first_fragment)
        else:
            offsets = []
            if nal_type == NAL_TYPE_STAP_A:
                assert len(data) > STAP_A_HEADER_SIZE,  'StapA header truncated.'

                offset = 1
                while offset < len(data):
                    (nulu_size,) = unpack_from('!H', data, offset)
                    offset += LENGTH_FIELD_SIZE
                    assert offset < len(data), 'StapA length field truncated.'
                    offsets.append(offset)
                    offset += nulu_size
                    assert offset <= len(data), 'StapA packet with incorrect NALU packet lengths.'

                nal_type = data[STAP_A_HEADER_SIZE] & 0x1f
            else:
                offsets.append(0)

            offsets.append(len(data) + LENGTH_FIELD_SIZE)
            for start, end in pairwise(offsets):
                end -= LENGTH_FIELD_SIZE
                output += bytes([0, 0, 0, 1])
                output += data[start:end]

            obj = cls(first_fragment=True)

        return obj, output


class H264Decoder:
    def __init__(self):
        self.codec = CodecContext.create('h264', 'r')

    def decode(self, data):
        try:
            packet = Packet(data)
            frames = self.codec.decode(packet)
        except AVError as e:
            logger.warning('failed to decode, skipping package: ' + str(e))
            return []

        video_frames = []
        for frame in frames:
            # TODO: avoid convert twice
            bgr_frame = frame.to_nd_array(format='bgr24')
            video_frame = frame_from_bgr(bgr_frame)
            video_frames.append(video_frame)

        return video_frames

    def parse(self, packet):
        descriptor, data = H264PayloadDescriptor.parse(packet.payload)
        packet._data = data
        packet._first_in_frame = descriptor.first_fragment


class H264Encoder:
    timestamp_increment = 6000

    def __init__(self):
        # TODO: i think the correct way to go is using CodecContext direct
        # but by now I could not figure out how

        buffer = io.BytesIO()
        self.container = av.open(buffer, format='h264', mode='w')
        self.stream = self.container.add_stream('libx264', rate=25)
        self.stream.pix_fmt = 'yuv420p'
        self.stream.codec_context.options = {
            'profile': 'baseline',
            'level': '31',
            'tune': 'zerolatency'
        }

    @staticmethod
    def _packetize_fu_a(data):<commit_msg>[h264] make timestamp / framerate consistent<commit_after>import io
import logging
import math
from itertools import tee
from struct import pack, unpack_from

import av
from av import AVError
from av.codec.context import CodecContext
from av.packet import Packet
from av.video.frame import VideoFrame

from ..contrib.media import frame_from_bgr, frame_to_bgr

logger = logging.getLogger('codec.h264')

PACKET_MAX = 1300

NAL_TYPE_FU_A = 28
NAL_TYPE_STAP_A = 24

NAL_HEADER_SIZE = 1
FU_A_HEADER_SIZE = 2
LENGTH_FIELD_SIZE = 2
STAP_A_HEADER_SIZE = NAL_HEADER_SIZE + LENGTH_FIELD_SIZE


def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


class H264PayloadDescriptor:
    def __init__(self, first_fragment):
        self.first_fragment = first_fragment

    def __repr__(self):
        return 'H264PayloadDescriptor(FF={})'.format(self.first_fragment)

    @classmethod
    def parse(cls, data):
        output = bytes()

        nal_type = data[0] & 0x1f
        if nal_type == NAL_TYPE_FU_A:
            assert len(data) >= FU_A_HEADER_SIZE, 'FU-A NAL units truncated.'

            f_nri = data[0] & (0x80 | 0x60)
            original_nal_type = data[1] & 0x1f
            first_fragment = bool(data[1] & 0x80)

            if first_fragment:
                original_nal_header = bytes([f_nri | original_nal_type])
                output += bytes([0, 0, 0, 1])
                output += original_nal_header

            output += data[2:]

            obj = cls(first_fragment=first_fragment)
        else:
            offsets = []
            if nal_type == NAL_TYPE_STAP_A:
                assert len(data) > STAP_A_HEADER_SIZE,  'StapA header truncated.'

                offset = 1
                while offset < len(data):
                    (nulu_size,) = unpack_from('!H', data, offset)
                    offset += LENGTH_FIELD_SIZE
                    assert offset < len(data), 'StapA length field truncated.'
                    offsets.append(offset)
                    offset += nulu_size
                    assert offset <= len(data), 'StapA packet with incorrect NALU packet lengths.'

                nal_type = data[STAP_A_HEADER_SIZE] & 0x1f
            else:
                offsets.append(0)

            offsets.append(len(data) + LENGTH_FIELD_SIZE)
            for start, end in pairwise(offsets):
                end -= LENGTH_FIELD_SIZE
                output += bytes([0, 0, 0, 1])
                output += data[start:end]

            obj = cls(first_fragment=True)

        return obj, output


class H264Decoder:
    def __init__(self):
        self.codec = CodecContext.create('h264', 'r')

    def decode(self, data):
        try:
            packet = Packet(data)
            frames = self.codec.decode(packet)
        except AVError as e:
            logger.warning('failed to decode, skipping package: ' + str(e))
            return []

        video_frames = []
        for frame in frames:
            # TODO: avoid convert twice
            bgr_frame = frame.to_nd_array(format='bgr24')
            video_frame = frame_from_bgr(bgr_frame)
            video_frames.append(video_frame)

        return video_frames

    def parse(self, packet):
        descriptor, data = H264PayloadDescriptor.parse(packet.payload)
        packet._data = data
        packet._first_in_frame = descriptor.first_fragment


class H264Encoder:
    timestamp_increment = 3000

    def __init__(self):
        # TODO: i think the correct way to go is using CodecContext direct
        # but by now I could not figure out how

        buffer = io.BytesIO()
        self.container = av.open(buffer, format='h264', mode='w')
        self.stream = self.container.add_stream('libx264', rate=30)
        self.stream.pix_fmt = 'yuv420p'
        self.stream.codec_context.options = {
            'profile': 'baseline',
            'level': '31',
            'tune': 'zerolatency'
        }

    @staticmethod
    def _packetize_fu_a(data):
{'fraction_non_alphanumeric': 0.09152713080699987, 'fraction_numerical': 0.02077300767971799, 'mean_word_length': 2.1189634864546525, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1317606', 'n_tokens_mistral': 2585, 'n_tokens_neox': 2351, 'n_words': 581}