ximage_to_video

ximage_to_video.c

#include <X11/Xatom.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/extensions/XShm.h>
#include <png.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <sys/ipc.h>
#include <sys/shm.h>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
#include <libswscale/swscale.h>
}

//int flush_encoder(AVCodecContext* codec_ctx, AVFormatContext* fmt_ctx, unsigned int stream_index)
//{
//    int ret;
//    int got_frame;
//    AVPacket enc_pkt;

//    //    if (!(fmt_ctx->streams[stream_index]->codecpar->codec->capabilities & AV_CODEC_CAP_DELAY))
//    //        return 0;

//    while (1) {
//        enc_pkt.data = NULL;
//        enc_pkt.size = 0;
//        av_init_packet(&enc_pkt);
//        //        ret = avcodec_encode_video2(,
//        //            NULL, &got_frame);
//        //        av_frame_free(NULL);

//        avcodec_send_frame(codec_ctx, &enc_pkt);
//        if (ret < 0)
//            break;
//        if (!got_frame) {
//            ret = 0;
//            break;
//        }
//        printf("Flush Encoder: Succeed to encode 1 frame!\tsize:%5d\n", enc_pkt.size);
//        /* mux encoded frame */
//        ret = av_write_frame(fmt_ctx, &enc_pkt);
//        if (ret < 0)
//            break;
//    }
//    return ret;
//}

bool writePngForImage(XImage* image, int width, int height, const char* filename)
{
    FILE* fp;

    png_structp png_ptr;
    png_infop png_info_ptr;
    png_bytep png_row;

    // Open file
    fp = fopen(filename, "wb");
    if (fp == NULL) {
        fprintf(stderr, "Could not open file for writing\n");
        return false;
    }

    // Initialize write structure
    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (png_ptr == nullptr) {
        fprintf(stderr, "Could not allocate write struct\n");
        return false;
    }

    // Initialize info structure
    png_info_ptr = png_create_info_struct(png_ptr);
    if (png_info_ptr == NULL) {
        fprintf(stderr, "Could not allocate info struct\n");
        return false;
    }

    // Setup Exception handling
    if (setjmp(png_jmpbuf(png_ptr))) {
        fprintf(stderr, "Error during png creation\n");
        return false;
    }

    png_init_io(png_ptr, fp);

    // Write header (8 bit colour depth)
    png_set_IHDR(png_ptr, png_info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    /*
    // Set title
    char * title = "Screenshot";

    if (title != NULL)
    {
        png_text title_text;
        title_text.compression = PNG_TEXT_COMPRESSION_NONE;
        title_text.key = "Title";
        title_text.text = title;
        png_set_text(png_ptr, png_info_ptr, &title_text, 1);
    }
    //*/

    png_write_info(png_ptr, png_info_ptr);

    // Allocate memory for one row (3 bytes per pixel - RGB)
    png_row = (png_bytep)malloc(3 * width * sizeof(png_byte));

    unsigned long red_mask = image->red_mask;
    unsigned long green_mask = image->green_mask;
    unsigned long blue_mask = image->blue_mask;

    // Write image data
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            unsigned long pixel = XGetPixel(image, x, y);
            unsigned char blue = pixel & blue_mask;
            unsigned char green = (pixel & green_mask) >> 8;
            unsigned char red = (pixel & red_mask) >> 16;
            png_byte* ptr = &(png_row[x * 3]);
            ptr[0] = red;
            ptr[1] = green;
            ptr[2] = blue;
        }
        png_write_row(png_ptr, png_row);
    }

    // End write
    png_write_end(png_ptr, NULL);

    // Free
    fclose(fp);

    png_free_data(png_ptr, png_info_ptr, PNG_FREE_ALL, -1);

    png_destroy_write_struct(&png_ptr, &png_info_ptr);

    if (png_row != nullptr)
        free(png_row);

    return true;
}

bool gRunning = true;

void signalHandler(int)
{
    gRunning = false;
}

int main(int argc, const char* argv[])
{
    // ###############################################################
    // Setup SIGINT Handler

    struct sigaction sigIntHandler;

    sigIntHandler.sa_handler = signalHandler;
    sigemptyset(&sigIntHandler.sa_mask);
    sigIntHandler.sa_flags = 0;

    sigaction(SIGINT, &sigIntHandler, NULL);

    // ###############################################################
    // Setup libav stuff

    AVPacket pkt;

    AVCodecID chosenEncoder = AV_CODEC_ID_H264;
    AVPixelFormat destinationPixelFormat = AV_PIX_FMT_YUV444P;

    int y_size;

    //FILE *in_file = fopen("src01_480x272.yuv", "rb");	//Input raw YUV data
    // FILE *in_file = fopen("../ds_480x272.yuv", "rb");   //Input raw YUV data
    int in_w = 800, in_h = 600; //Input data's width and height
    int framenum = 100; //Frames to encode
    //const char* out_file = "src01.h264";              //Output Filepath
    //const char* out_file = "src01.ts";
    //const char* out_file = "src01.hevc";

    const char* out_file = "salida.mp4";

    // ###############################################################
    // Register all muxers, etc
    printf("Initializing libavformat... \n");
    av_register_all();

    // #################################################################################################################
    // AVFormatContext: Allocate output format context. It holds information about the container (mp4 file)

    printf("Allocating format context... \n");
    AVFormatContext* pFormatCtx;

    if (avformat_alloc_output_context2(&pFormatCtx, nullptr, nullptr, out_file) != 0) {
        printf("Error building format context \n");
        return -1;
    }

    printf("Guessed format: %s \n", pFormatCtx->oformat->long_name);

    // #################################################################################################################
    // AVStream: create video stream

    printf("Adding new stream... \n");
    AVStream* pVideoStream = avformat_new_stream(pFormatCtx, nullptr);

    if (pVideoStream == nullptr) {
        printf("Failed adding new stream \n");
        return -1;
    }

    // #################################################################################################################
    // AVCodec: find chosen codec

    printf("Finding encoder... \n");
    AVCodec* pCodec = avcodec_find_encoder(chosenEncoder);
    if (!pCodec) {
        printf("Cannot find encoder! \n");
        return -1;
    }

    // #################################################################################################################
    // Build the AVCodecContext based in the created codec

    printf("Allocating codec context... \n");
    AVCodecContext* pCodecCtx = avcodec_alloc_context3(pCodec);

    if (!pCodecCtx) {
        printf("Cannot alloc codec context! \n");
        return -1;
    }

    // #################################################################################################################
    // Setup basic attributes for the coding procedure

    //    //    pCodecCtx->codec_id = fmt->video_codec;
    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->pix_fmt = destinationPixelFormat;
    pCodecCtx->width = in_w;
    pCodecCtx->height = in_h;
    pCodecCtx->bit_rate = 800000;
    pCodecCtx->gop_size = 250;

    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = 25;

    pVideoStream->time_base = pCodecCtx->time_base;

    //H264
    //pCodecCtx->me_range = 16;
    //pCodecCtx->max_qdiff = 4;
    //pCodecCtx->qcompress = 0.6;
    pCodecCtx->qmin = 40;
    pCodecCtx->qmax = 51;

    //Optional Param
    pCodecCtx->max_b_frames = 3;

    // #################################################################################################################
    // Open the codec

    printf("Opening encoder... \n");
    if (avcodec_open2(pCodecCtx, pCodec, nullptr) < 0) {
        printf("Failed to open encoder! \n");
        return -1;
    }

    // #################################################################################################################
    // Copy parameters from context to the stream

    if(avcodec_parameters_from_context(pVideoStream->codecpar, pCodecCtx) < 0)
    {
        printf("Error building parameters from context \n");
        return -1;
    }

//    AVCodecParameters * pCodecParameters = pVideoStream->codecpar;
//    pCodecParameters->codec_id = chosenEncoder;
//    pCodecParameters->codec_type = AVMEDIA_TYPE_VIDEO;
//    pCodecParameters->format = AV_PIX_FMT_YUV420P;
//    pCodecParameters->width = in_w;
//    pCodecParameters->height = in_h;
//    pCodecParameters->bit_rate = 800000;
//    pCodecParameters->time_base.num = 1;
//    pCodecParameters->time_base.den = 25;



    /*
//    pCodecParameters->gop_size = 250;
//
//    pCodecParameters->time_base.num = 1;
//    pCodecParameters->time_base.den = 25;
//
//    //H264
//    //pCodecCtx->me_range = 16;
//    //pCodecCtx->max_qdiff = 4;
//    //pCodecCtx->qcompress = 0.6;
//    pCodecCtx->qmin = 40;
//    pCodecCtx->qmax = 51;
//
//    //Optional Param
//    pCodecCtx->max_b_frames = 3;
////    avcodec_parameters_copy(pVideoStream->codecpar, )
     */
//
//    pVideoStream->time_base.num = 1;
//    pVideoStream->time_base.den = 25;




//    if(avcodec_parameters_to_context(pCodecCtx, pCodecParameters) < 0)
//    {
//        printf("Error calling avcodec_parameters_to_context \n");
//        return -1;
//    }

    // Set Option
    AVDictionary* param = 0;
    //H.264
    if (pCodecCtx->codec_id == AV_CODEC_ID_H264) {
        av_dict_set(&param, "preset", "fast", 0);
        //        av_dict_set(&param, "tune", "zerolatency", 0);
        av_dict_set(&param, "profile", "main", 0);
    }

    // #################################################################################################################
    // Open output file

    printf("Opening output file... \n");
    if (avio_open(&pFormatCtx->pb, out_file, AVIO_FLAG_READ_WRITE) < 0) {
        printf("Failed to open output file! \n");
        return -1;
    }

//    printf("Setting codec parameters... \n");
    //    AVCodecContext* pCodecCtx = pVideoStream->codec;

    //    //H.265
    //    if (pCodecCtx->codecpar_id == AV_CODEC_ID_H265) {
    //        av_dict_set(&param, "preset", "ultrafast", 0);
    //        av_dict_set(&param, "tune", "zero-latency", 0);
    //    }

    //Show some Information
    av_dump_format(pFormatCtx, 0, out_file, 1);

    printf("Writing file header... \n");
    if (avformat_write_header(pFormatCtx, nullptr) < 0) {
        printf("Error writing stream header \n");
        return -1;
    }

    // ------------------------------------------------------------------------
    // ------------------------------------------------------------------------

    printf("Allocating frame... \n");
    AVFrame* pFrame = av_frame_alloc();
    pFrame->format = destinationPixelFormat;
    pFrame->width = in_w;
    pFrame->height = in_h;

    //    int picture_size = avpicture_get_size(pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);


    printf("Computing image size... \n");
    int image_size = av_image_get_buffer_size(destinationPixelFormat, in_w, in_h, 1);

    printf("Image size: %framecnt bytes \n", image_size);

/*
    uint8_t* image_buf = (uint8_t*)av_malloc(image_size);

    printf("Filling frame fields... \n");
    //    avpicture_fill((AVPicture*)pFrame, picture_buf, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
    av_image_fill_arrays(pFrame->data, pFrame->linesize, image_buf, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height, 1);

    //*/

    printf("Allocating image buffer... \n");
    if (av_image_alloc(pFrame->data, pFrame->linesize, pCodecCtx->width, pCodecCtx->height, pCodecCtx->pix_fmt, 1) < 0)
    {
        printf("Could not allocate picture buffer \n");
        return -1;
    }

//*/




    printf("Allocating packet... \n");
    av_new_packet(&pkt, image_size);

    y_size = pCodecCtx->width * pCodecCtx->height;

    // Initialize sws scaler

    printf("Initializing sws context...\n");
    SwsContext* sws_context = sws_getContext(
            in_w, in_h, AV_PIX_FMT_RGB24,
            in_w, in_h, destinationPixelFormat,
            SWS_LANCZOS | SWS_ACCURATE_RND, nullptr, nullptr, nullptr);

    if (!sws_context) {
        printf("Failed getting sws context! \n");
        return -1;
    }

    // ------------------------------------------------------------------------------
    // Setup X11 stuff

    int screen;
    Window root;
    Display* display;
    XImage* img;

    int shm = 0;
    XShmSegmentInfo shminfo;

    /* My Desktop Screen Resolution */

    int width = 800;
    int height = 600;

    display = XOpenDisplay(NULL);

    shm = XShmQueryExtension(display);

    if (!shm) {
        printf("Error, no SHM extension \n");
        return 1;
    }

    int scr = XDefaultScreen(display);

    img = XShmCreateImage(display, DefaultVisual(display, scr),
        DefaultDepth(display, scr), ZPixmap, nullptr, &shminfo,
        width, height);
    // printf("\n Bytes Per Line %d ", img->bytes_per_line);

    shminfo.shmid = shmget(IPC_PRIVATE, img->bytes_per_line * img->height,
        IPC_CREAT | 0777);

    if (shminfo.shmid == -1) {
        printf("Can not get the shared Memory ... \n");
        return 1;
    }

    shminfo.shmaddr = img->data = (char*)shmat(shminfo.shmid, 0, 0);
    shminfo.readOnly = False;

    if (!XShmAttach(display, &shminfo)) {
        printf("Unable to attach now... \n");
        return 1;
    }

    char imageName[1024];

    // -------------------------------------------------------------------------
    // MAIN LOOP

    auto pixel_rgb_data = new uint8_t[width * height * 4];
    int framecnt = 0;

    while (gRunning) {
        printf("Processing frame framecnt=%i \n", framecnt);

        if (!XShmGetImage(display, RootWindow(display, DefaultScreen(display)), img, 0, 0, AllPlanes)) {
            printf("\n Ooops.. Something is wrong.");
            break;
        }

        snprintf(imageName, sizeof(imageName), "salida_%i.png", framecnt);
        writePngForImage(img, width, height, imageName);

        unsigned long red_mask = img->red_mask;
        unsigned long green_mask = img->green_mask;
        unsigned long blue_mask = img->blue_mask;

        // Write image data
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                unsigned long pixel = XGetPixel(img, x, y);

                unsigned char blue = pixel & blue_mask;
                unsigned char green = (pixel & green_mask) >> 8;
                unsigned char red = (pixel & red_mask) >> 16;

                pixel_rgb_data[y * width + x * 3] = red;
                pixel_rgb_data[y * width + x * 3 + 1] = green;
                pixel_rgb_data[y * width + x * 3 + 2] = blue;
            }
        }

        uint8_t* inData[1] = { pixel_rgb_data };
        int inLinesize[1] = { in_w };

        printf("Scaling frame... \n");
        int sliceHeight = sws_scale(sws_context, inData, inLinesize, 0, height, pFrame->data, pFrame->linesize);

        printf("Obtained slice height: %i \n", sliceHeight);
        pFrame->pts = framecnt * (pVideoStream->time_base.den) / ((pVideoStream->time_base.num) * 25);

        printf("Frame pts: %li \n", pFrame->pts);
        int got_picture = 0;

        printf("Encoding frame... \n");
        int ret = avcodec_encode_video2(pCodecCtx, &pkt, pFrame, &got_picture);

//                int ret = avcodec_send_frame(pCodecCtx, pFrame);

        if (ret != 0) {
            printf("Failed to encode! Error: %i\n", ret);
            return -1;
        }

        printf("Succeed to encode frame: %5d - size: %5d\n", framecnt, pkt.size);

        framecnt++;

        pkt.stream_index = pVideoStream->index;
        ret = av_write_frame(pFormatCtx, &pkt);

        if (ret != 0) {
            printf("Error writing frame! Error: %framecnt \n", ret);
            return -1;
        }

        av_packet_unref(&pkt);
    }

    // X11 CLEANUP
    XShmDetach(display, &shminfo);
    img->data = NULL;
    XDestroyImage(img);

    XCloseDisplay(display);

    // libav CLEANUP

    sws_freeContext(sws_context);

    //Flush Encoder
    //    int ret = flush_encoder(pCodecCtx, pFormatCtx, 0);
    //    if (ret < 0) {
    //        printf("Flushing encoder failed\n");
    //        return -1;
    //    }

    //Write file trailer
    av_write_trailer(pFormatCtx);

    //Clean
    if (pVideoStream) {
        //avcodec_close(pVideoStream->codec);
        av_free(pFrame);
//        av_free(image_buf);
    }
    avio_close(pFormatCtx->pb);
    avformat_free_context(pFormatCtx);
}