HLS segmenter written in C, as a python extension

gistfile1.c

/*

Video file segmenter, implemented as a python C extension

*/

#include <Python.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <libavformat/avformat.h>
#include <libavutil/mathematics.h>  // For av_rescale_q

// Exceptions
static PyObject *SegmentError;


// Prototypes
char *create_basename(const char *dir, const char *format);
int process_video(const char *input_filename, const char *dir, const char *filename, const int s_length, const char *playlist, int start_segment, int end_segment);
static AVStream *add_output_stream(AVFormatContext *output_format_context, AVStream *input_stream);
static int write_frame_fixed(AVPacket *packet, AVFormatContext *input_context, AVFormatContext *output_context, AVBitStreamFilterContext *bsf);
double get_playlist_timestamp(char *filename, int start_segment);
int file_exists(const char *filename);

/*
    Segment an input videofile and produce either a playlist, video segments
    and/or only a specific interval of segments [segment X-Y out of Z]

    Arguments:
    - (String) Input filename (absolute path)
    - (String) Destination directory (absolute path)
    - (String) Format (human readable description, such as 'video_medium')
    - (Int) Segment approximate length (in seconds)
    - (String) Playlist filename. Will be created if it does not exist!

    Optional arguments:
    - (Int) Start segment # (1-indexed, -1 means do not write segments)
    - (Int) End segment # (1-indexed, -1 to include everything)

    If a specific interval of segments is request, no playlist exists and
    "create playlist" is False, it will throw an error.

    Produces:
    - /folder/format.m3u8 playlist
    - /folder/format_%05u.ts chunks
*/
static PyObject *
segmenter_segment(PyObject *self, PyObject *args)
{
    // Required arguments
    const char *filename;
    const char *dir;
    const char *format;
    const int segment_length;
    const char *playlist_name;

    // Optional arguments with default values
    int start_segment = 1;
    int end_segment = 10;

    // See http://docs.python.org/2.7/c-api/arg.html
    if (!PyArg_ParseTuple(args, "sssis|ii", &filename,
                                            &dir,
                                            &format,
                                            &segment_length,
                                            &playlist_name,
                                            &start_segment,
                                            &end_segment))
        return NULL;

    if (end_segment != -1 && end_segment < start_segment)
    {
        PyErr_SetString(SegmentError, "End segment must be after the starting segment");
        return NULL;
    }

    // Run it!
    int result;
    result = process_video(filename, dir, format, segment_length, playlist_name, start_segment, end_segment);
    if (result != 0) // Exception occurred!
        return NULL;

    const char *fake = "SUCCESS!";
    return Py_BuildValue("s", fake);
}

static PyMethodDef SegmenterMethods[] = {
    {"segment", segmenter_segment, METH_VARARGS,
     "Process a video creating a playlist or video segments."},
    {NULL, NULL, 0, NULL} /* Sentinel */
};

PyMODINIT_FUNC
initsegmenter(void)
{
    PyObject *m;
    m = Py_InitModule("segmenter", SegmenterMethods);
    if (m == NULL)
        return;

    SegmentError = PyErr_NewException("segmenter.SegmentError", NULL, NULL);
    Py_INCREF(SegmentError);
    PyModule_AddObject(m, "SegmentError", SegmentError);

    // Get ready for libav
    avcodec_register_all();
    av_register_all();
}

// C Code
char *create_basename(const char *dir, const char *filename)
{
    size_t dlen = strlen(dir);
    size_t flen = strlen(filename);
    char *basename = (char *) malloc(sizeof(char) * (dlen +flen + 2));
    if (dir[dlen-1] != '/')
        snprintf(basename, dlen + flen + 2, "%s/%s", dir, filename);
    else
        snprintf(basename, dlen + flen + 1, "%s%s", dir, filename);

    return basename;
}

/*

*/
double get_playlist_timestamp(char *filename, int start_segment)
{
    FILE *file = fopen(filename, "r");
    if (!file)
    {
        PyErr_Format(SegmentError, "Playlist '%s' not found.", filename);
        return -1;
    }
    char *durtext = "#EXTINF:";
    char buffer[64];
    char duration[16];
    int cur_segment = 1;
    double timestamp = -1;
    double total = 0;

    while (fgets(buffer, sizeof buffer, file) != NULL)
    {
        if (strncmp(buffer, durtext, strlen(durtext)) == 0)
        {
            cur_segment++;
            strncpy(duration, buffer+strlen(durtext), 16);
            timestamp = atof(duration);
            if (timestamp == 0) // atof error
            {
                total = -1;
                break;
            }
            total += timestamp;

            if (cur_segment == start_segment)
            {
                break;
            }
        }
    }

    fclose(file);
    return total;
}

int file_exists(const char *filename)
{
    struct stat buffer;
    return (stat(filename, &buffer) == 0);
}

/*

*/
int create_playlist(char *filename, const char *base, double *timestamps)
{
    FILE *fp;
    fp = fopen(filename, "w");
    if (!fp)
    {
        PyErr_Format(SegmentError, "Could not open file '%s' for writing.", filename);
        return -1;
    }

    // Determine real max segment length
    int max = 0;
    unsigned int i = 0;
    while(1)
    {
        double current = timestamps[i];
        if (current == 0)
            break;
        if (current > max)
            max = (int)(current + 1); // Round up

        i++;
    }

    // Write file headers
    fprintf(fp, "#EXTM3U\n");
    fprintf(fp, "#EXT-X-VERSION:3\n");
    fprintf(fp, "#EXT-X-TARGETDURATION:%d\n", max);
    fprintf(fp, "#EXT-X-MEDIA-SEQUENCE:0\n");
    fprintf(fp, "#EXT-X-ALLOW-CACHE:YES\n");

    // Write segments
    i = 0;
    while(1)
    {
        double current = timestamps[i];
        if (current == 0)
            break;

        fprintf(fp, "#EXTINF:%f,\n", current);
        fprintf(fp, "%s_%05u.ts\n", base, i+1);
        i++;
    }

    fprintf(fp, "#EXT-X-ENDLIST\n");

    fclose(fp);

    return 0;
}

/*
    Creates a new audio or video stream with the properties copied from the input stream.
*/
static AVStream *add_output_stream(AVFormatContext *output_format_context, AVStream *input_stream)
{
    AVCodecContext *input_codec_context;
    AVCodecContext *output_codec_context;
    AVStream *output_stream;

    output_stream = avformat_new_stream(output_format_context, NULL);
    if (!output_stream) {
        PyErr_SetString(SegmentError, "Could not create new output stream.");
        return NULL;
    }

    input_codec_context = input_stream->codec;
    output_codec_context = output_stream->codec;

    output_codec_context->codec_id = input_codec_context->codec_id;
    output_codec_context->codec_type = input_codec_context->codec_type;
    output_codec_context->codec_tag = input_codec_context->codec_tag;
    output_codec_context->bit_rate = input_codec_context->bit_rate;

    // Copy this, as otherwise releasing either the input or output context
    // will fail (as they try to free() the same pointer).
    output_codec_context->extradata = av_malloc(input_codec_context->extradata_size + FF_INPUT_BUFFER_PADDING_SIZE);
    memcpy(output_codec_context->extradata, input_codec_context->extradata, input_codec_context->extradata_size);
    output_codec_context->extradata_size = input_codec_context->extradata_size;

    // Some codecs put the wrong ticks_per_frame in (in the thousands instead of tens)
    if(av_q2d(input_codec_context->time_base) * input_codec_context->ticks_per_frame > av_q2d(input_stream->time_base) && av_q2d(input_stream->time_base) < 1.0/1000) {
        output_codec_context->time_base = input_codec_context->time_base;
        output_codec_context->time_base.num *= input_codec_context->ticks_per_frame;
    }
    else {
        output_codec_context->time_base = input_stream->time_base;
    }

    switch (input_codec_context->codec_type) {
        case AVMEDIA_TYPE_AUDIO:
            output_codec_context->channel_layout = input_codec_context->channel_layout;
            output_codec_context->sample_rate = input_codec_context->sample_rate;
            output_codec_context->channels = input_codec_context->channels;
            output_codec_context->frame_size = input_codec_context->frame_size;
            if ((input_codec_context->block_align == 1 && input_codec_context->codec_id == AV_CODEC_ID_MP3) || input_codec_context->codec_id == AV_CODEC_ID_AC3) {
                output_codec_context->block_align = 0;
            }
            else {
                output_codec_context->block_align = input_codec_context->block_align;
            }
            break;
        case AVMEDIA_TYPE_VIDEO:
            output_codec_context->pix_fmt = input_codec_context->pix_fmt;
            output_codec_context->width = input_codec_context->width;
            output_codec_context->height = input_codec_context->height;
            output_codec_context->has_b_frames = input_codec_context->has_b_frames;

            if (output_format_context->oformat->flags & AVFMT_GLOBALHEADER) {
                output_codec_context->flags |= CODEC_FLAG_GLOBAL_HEADER;
            }
            break;
        default:
            break;
    }

    return output_stream;
}

/*
    Function for writing a frame to the correct output context.

    Makes sure that:
    - The mp4_h264toannexb bitfilter for converting mp4 to MPEG-TS
    - Correct rescaling of time unit (pts / dts) to the destination format
*/
static int write_frame_fixed(AVPacket *packet, AVFormatContext *input_context, AVFormatContext *output_context, AVBitStreamFilterContext *bsf)
{
    int si = packet->stream_index;
    if (output_context->streams[si]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
    {
        if (input_context->streams[si]->codec->codec_id == AV_CODEC_ID_H264)
        {
            int a = av_bitstream_filter_filter(bsf, output_context->streams[si]->codec, NULL,
                                               &packet->data, &packet->size,
                                               packet->data, packet->size,
                                               packet->flags & AV_PKT_FLAG_KEY);

            if (a < 0)
            {
                PyErr_SetString(SegmentError, "Bitstream filter failed (h264_mp4toannexb).");
                return -1;
            }
        }
    }

    packet->pts = av_rescale_q(packet->pts, input_context->streams[si]->time_base, output_context->streams[si]->time_base);
    packet->dts = av_rescale_q(packet->dts, input_context->streams[si]->time_base, output_context->streams[si]->time_base);

    return av_interleaved_write_frame(output_context, packet);
}

/*
    The actual processing function
*/
int process_video(const char *input_filename, const char *dir, const char *filename, const int s_length, const char *playlist_name, int start_segment, int end_segment)
{
    // Definitions
    char *output_filename = NULL;
    char *playlist_filename = NULL;
    int playlist = 0;
    int error = 0;
    double *timestamps = NULL;
    AVFormatContext *input_context = NULL;
    AVFormatContext *output_context = NULL;
    AVBitStreamFilterContext *bsf = NULL;
    AVOutputFormat *output_format = NULL;

    char *basename = create_basename(dir, filename);
    // Segment name: basename + "_XXXXX.ts"
    output_filename = (char*) malloc(sizeof(char) * (strlen(basename) + 11));
    if (!output_filename)
    {
        PyErr_SetString(SegmentError, "Could not allocate memory for output filename");
        error = 1;
        goto cleanup;
    }

    // Playlist name: basename + ".m3u8"
    playlist_filename = create_basename(dir, playlist_name);
    if (!playlist_filename)
    {
        PyErr_SetString(SegmentError, "Could not allocate memory for playlist filename");
        error = 1;
        goto cleanup;
    }
    playlist = file_exists(playlist_filename);

    // Assign input and output context - guess input format based on filename
    int ret; // Return status
    input_context = NULL;
    ret = avformat_open_input(&input_context, input_filename, NULL, NULL);
    if (ret != 0)
    {
        PyErr_Format(SegmentError, "Could not create input context from %s", input_filename);
        error = 1;
        goto cleanup;
    }

    // Read "sample" packets to discover additional info about the file/format
    ret = avformat_find_stream_info(input_context, NULL);
    if (ret < 0)
    {
        PyErr_SetString(SegmentError, "Unable to read stream info from input context");
        error = 1;
        goto cleanup;
    }

    // Create the format for the output context - av_guess_format = find format
    output_format = av_guess_format("mpegts", NULL, NULL);
    if (!output_format)
    {
        PyErr_SetString(SegmentError, "Unable to open the MPEG-TS output format (muxer)");
        error = 1;
        goto cleanup;
    }

    // Allocate the output context and set the output format
    output_context = avformat_alloc_context();
    if (!output_context)
    {
        PyErr_SetString(SegmentError, "Unable to allocate output context");
        error = 1;
        goto cleanup;
    }
    output_context->oformat = output_format;
    av_dict_set(&output_context->metadata, "service_name", "23Video Segmenter", 0);

    // Copy one video and one audio stream - additional streams will be ignored
    AVStream *video_stream = NULL;
    AVStream *audio_stream = NULL;

    int i;
    int video_index = -1;
    int audio_index = -1;
    int main_stream = -1;
    int off_stream = -1;
    for (i = 0; i < input_context->nb_streams && (video_stream == NULL || audio_stream == NULL); i++)
    {
        switch (input_context->streams[i]->codec->codec_type)
        {
            case AVMEDIA_TYPE_VIDEO:
                video_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                video_stream = add_output_stream(output_context, input_context->streams[i]);
                av_dict_set(&video_stream->metadata, "service_name", "23Video Segmenter", 0);
                if (video_stream == NULL)
                {
                    error = 1;
                    goto cleanup;
                }
                break;

            case AVMEDIA_TYPE_AUDIO:
                audio_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                audio_stream = add_output_stream(output_context, input_context->streams[i]);
                av_dict_set(&audio_stream->metadata, "service_name", "23Video Segmenter", 0);
                if (audio_stream == NULL)
                {
                    error = 1;
                    goto cleanup;
                }
                break;

            default:
                input_context->streams[i]->discard = AVDISCARD_ALL;
                break;
        }
    }

    if (video_stream == NULL && audio_stream == NULL)
    {
        PyErr_SetString(SegmentError, "File has no video and audio streams!");
        error = 1;
        goto cleanup;
    }
    else
    {
        if (video_index != -1)
            main_stream = video_index;
        else
            main_stream = audio_index;

        if (audio_index != -1 && main_stream != audio_index)
            off_stream = audio_index;
    }

    // First file to be written, if -1 it will default to 1
    unsigned int output_index = 1;

    // Write frames (or just read and skip?)
    int write_frames = 0;

    // Timestamps for playlist - 5000 * 10 seconds = ~13.8 hours
    timestamps = (double *)calloc(5000, sizeof(double));

    // Tracking variables
    int decode_done;
    double prev_segment_time = 0;
    double final_time = 0;
    double segment_time;
    int64_t start_pts = -1;
    double time_base = av_q2d(input_context->streams[main_stream]->time_base);
    double offbase = 0;
    if (off_stream != -1)
        offbase = av_q2d(input_context->streams[off_stream]->time_base);

    // Start reading frames
    AVPacket packet;
    av_init_packet(&packet);
    decode_done = av_read_frame(input_context, &packet);

    // Movies don't always start at 0 :(
    start_pts = packet.pts;
    prev_segment_time = start_pts * time_base;

    // Do we need to seek and find the keyframe of [start_segment]
    if (playlist == 1 && start_segment > 1)
    {
        double seek_to = get_playlist_timestamp(playlist_filename, start_segment);
        if (seek_to == -1)
        {
            PyErr_SetString(SegmentError, "Error parsing playlist (nonexistant or malformed)");
            error = 1;
            goto cleanup;
        }


        // Seek to before the requested time to avoid overshooting it
        int64_t target = (seek_to / time_base) + start_pts;
        // 10000 pts is around ~0.8 seconds for most of our videos
        av_seek_frame(input_context, main_stream, target - 10000, 0);
        decode_done = av_read_frame(input_context, &packet);

        while (1)
        {
            decode_done = av_read_frame(input_context, &packet);
            if (packet.stream_index == main_stream)
            {
                if (packet.pts >= target)
                    break;
            }

            if (decode_done)
            {
                PyErr_SetString(SegmentError, "Reached EOF when seeking!");
                error = 1;
                goto cleanup;
            }
        }

        output_index = start_segment;
        prev_segment_time = packet.pts * time_base;
    }

    // Only actually copy anything if we are writing
    if (start_segment != -1)
    {
        // Assign and open the first file if we're starting from scratch OR we have seek'd to the right location
        if (playlist != 0 || start_segment == 1)
        {
            snprintf(output_filename, strlen(basename) + 10, "%s_%05u.ts", basename, output_index);
            if (avio_open(&output_context->pb, output_filename, AVIO_FLAG_WRITE) < 0)
            {
                PyErr_SetString(SegmentError, "Could not open segment file for writing");
                error = 1;
                goto cleanup;
            }

            if (avformat_write_header(output_context, NULL))
            {
                PyErr_SetString(SegmentError, "Could not write header to segment file");
                error = 1;
                goto cleanup;
            }

            write_frames = 1;
        }

        bsf = av_bitstream_filter_init("h264_mp4toannexb");
        if (!bsf)
        {
            PyErr_SetString(SegmentError, "Could not initialize the h264_mp4toannexb filter!");
            error = 1;
            goto cleanup;
        }
    }

    double totaldur = (input_context->duration) / ((double)AV_TIME_BASE);

    do
    {
        if (packet.stream_index != main_stream && packet.stream_index != off_stream)
        {
            av_free_packet(&packet);
            decode_done = av_read_frame(input_context, &packet);
            continue;
        }

        if (packet.stream_index == main_stream)
            final_time = (packet.pts + packet.duration) * time_base;
        else
            final_time = (packet.pts + packet.duration) * offbase;

        if (packet.stream_index == main_stream && (packet.flags & AV_PKT_FLAG_KEY))
        {
            segment_time = packet.pts * time_base;
        }
        else
        {
            segment_time = prev_segment_time;
        }

        if (segment_time - prev_segment_time >= s_length)
        {
            // Include the last bit if there's less than a second left
            // We do this to avoid "micro" segments
            if ((segment_time + 1) < totaldur)
            {
                if (playlist == 0)
                {
                    double length = segment_time - prev_segment_time;
                    timestamps[output_index - 1] = length;
                }

                // Is the current file within our interval
                if (write_frames)
                {
                    // Close it
                    av_write_trailer(output_context);
                    avio_flush(output_context->pb);
                    avio_close(output_context->pb);
                }

                output_index++;

                // If we are writing at all?
                // And have we reached a segment within our interval
                if (start_segment != -1 &&
                    start_segment <= output_index &&
                    (end_segment >= output_index || end_segment == -1))
                    write_frames = 1;
                else
                    write_frames = 0;

                // Write the new segment to disk?
                if (write_frames)
                {
                    // Next segment filename
                    snprintf(output_filename, strlen(basename) + 10, "%s_%05u.ts", basename, output_index);
                    if (avio_open(&output_context->pb, output_filename, AVIO_FLAG_WRITE) < 0)
                    {
                        PyErr_Format(SegmentError, "Could not open segment file '%s' for writing", output_filename);
                        error = 1;
                        goto cleanup;
                    }

                    avformat_write_header(output_context, NULL);
                }
                else if (playlist == 1)
                {
                    // Playlist has already been made, drop out
                    av_free_packet(&packet);
                    break; 
                }
                prev_segment_time = segment_time;
            }
        }

        if (write_frames)
        {
            ret = write_frame_fixed(&packet, input_context, output_context, bsf);
            avio_flush(output_context->pb);
            if (ret < 0)
            {
                char *errormsg = malloc(sizeof(char) * 100);
                int success = av_strerror(ret, errormsg, 100);
                if (success == 0) {
                    PyErr_Format(SegmentError, "Unable to write frame. Error: %s", errormsg);
                }
                else {
                    PyErr_Format(SegmentError, "Unable to write frame. Error No: %x", ret);
                }
                free(errormsg);
                error = 1;
                goto cleanup;
            }
        }
        av_free_packet(&packet);
        decode_done = av_read_frame(input_context, &packet);
    } while (!decode_done);

    if (write_frames)
    {
        av_write_trailer(output_context);
        avio_flush(output_context->pb);
        avio_close(output_context->pb);
    }

    if (playlist == 0)
    {
        double length = final_time - prev_segment_time;
        timestamps[output_index - 1] = length;
        ret = create_playlist(playlist_filename, filename, timestamps);
        if (ret != 0)
            error = 1;
    }

cleanup:
    // Libav stuff
    while (bsf) {
        AVBitStreamFilterContext *next = bsf->next;
        av_bitstream_filter_close(bsf);
        bsf = next;
    }
    if (output_context)
        avformat_free_context(output_context);
    if (input_context)
        avformat_close_input(&input_context);

    // Strings
    if (basename)
        free(basename);
    if (output_filename)
        free(output_filename);
    if (playlist_filename)
        free(playlist_filename);

    if (timestamps)
        free(timestamps);

    if (error)
        return -1;
    return 0;
}