cuframes-composer/examples/grid_record.c

/* grid_record — Phase 2 smoke test.
 *
 * Подписывается на 4 cuframes-источника (4 камеры), композирует их в 2×2 grid
 * через cfc_composer, кодирует через NVENC, пишет H.264 в файл.
 *
 * Layout 2×2 1080p:
 *   Output: 3840×2160 (4K)
 *   Cells:  4 шт. 1920×1080 в углах
 *
 * Использование:
 *   grid_record --out 4k.h264 \
 *               --cell cam-parking,0,0,1920,1080 \
 *               --cell cam-back_yard,1920,0,1920,1080 \
 *               --cell cam-front_yard,0,1080,1920,1080 \
 *               --cell cam-gate_lpr,1920,1080,1920,1080 \
 *               --seconds 15
 *
 * Лицензия: LGPL-2.1+
 */

#include "../include/cuframes_composer/composer.h"
#include "../include/cuframes_composer/nvenc.h"

#include <cuda.h>

#include <errno.h>
#include <getopt.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define MAX_CELLS 16

static volatile sig_atomic_t g_stop = 0;

static void on_sig(int s) { (void)s; g_stop = 1; }

typedef struct write_ctx {
    FILE *fp;
    uint64_t bytes_written;
    uint64_t frames_encoded;
    uint64_t idr_count;
} write_ctx_t;

static void on_bitstream(const uint8_t *bs, size_t size, int64_t pts_ns,
                          int is_idr, void *user)
{
    (void)pts_ns;
    write_ctx_t *ctx = (write_ctx_t *)user;
    if (fwrite(bs, 1, size, ctx->fp) == size) {
        ctx->bytes_written += size;
        ctx->frames_encoded++;
        if (is_idr) ctx->idr_count++;
    }
}

static int parse_cell(const char *arg, cfc_composer_cell_t *out,
                       char *key_storage)
{
    /* Формат: key,x,y,w,h */
    char buf[128];
    strncpy(buf, arg, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = '\0';
    char *tok = strtok(buf, ",");
    if (!tok) return -1;
    strncpy(key_storage, tok, 63);
    key_storage[63] = '\0';
    out->source_key = key_storage;
    tok = strtok(NULL, ","); if (!tok) return -1; out->x = atoi(tok);
    tok = strtok(NULL, ","); if (!tok) return -1; out->y = atoi(tok);
    tok = strtok(NULL, ","); if (!tok) return -1; out->w = atoi(tok);
    tok = strtok(NULL, ","); if (!tok) return -1; out->h = atoi(tok);
    return 0;
}

static const char *cu_err(CUresult r)
{
    const char *s = NULL;
    cuGetErrorString(r, &s);
    return s ? s : "?";
}

int main(int argc, char **argv)
{
    const char *out_path = NULL;
    int fps = 25, bitrate = 10000, max_seconds = 0;
    int out_w = 3840, out_h = 2160;
    cfc_composer_cell_t cells[MAX_CELLS] = { 0 };
    static char cell_keys[MAX_CELLS][64];
    int num_cells = 0;

    static struct option opts[] = {
        {"out",     required_argument, 0, 'o'},
        {"cell",    required_argument, 0, 'c'},
        {"fps",     required_argument, 0, 'f'},
        {"bitrate", required_argument, 0, 'b'},
        {"width",   required_argument, 0, 'W'},
        {"height",  required_argument, 0, 'H'},
        {"seconds", required_argument, 0, 's'},
        {0, 0, 0, 0},
    };
    int c;
    while ((c = getopt_long(argc, argv, "o:c:f:b:W:H:s:", opts, NULL)) != -1) {
        switch (c) {
        case 'o': out_path = optarg; break;
        case 'c':
            if (num_cells >= MAX_CELLS) {
                fprintf(stderr, "max %d cells\n", MAX_CELLS);
                return 1;
            }
            if (parse_cell(optarg, &cells[num_cells], cell_keys[num_cells]) != 0) {
                fprintf(stderr, "invalid --cell '%s' (key,x,y,w,h)\n", optarg);
                return 1;
            }
            num_cells++;
            break;
        case 'f': fps = atoi(optarg); break;
        case 'b': bitrate = atoi(optarg); break;
        case 'W': out_w = atoi(optarg); break;
        case 'H': out_h = atoi(optarg); break;
        case 's': max_seconds = atoi(optarg); break;
        default: return 1;
        }
    }
    if (!out_path || num_cells == 0) {
        fprintf(stderr,
            "Использование: %s --out <file.h264> --cell key,x,y,w,h [--cell ...]\n"
            "                [--width 3840] [--height 2160] [--fps 25]\n"
            "                [--bitrate 10000] [--seconds N]\n",
            argv[0]);
        return 1;
    }

    signal(SIGINT, on_sig);
    signal(SIGTERM, on_sig);

    /* CUDA primary context. */
    CUresult cr = cuInit(0);
    if (cr != CUDA_SUCCESS) { fprintf(stderr, "cuInit: %s\n", cu_err(cr)); return 1; }
    CUdevice dev;
    cuDeviceGet(&dev, 0);
    CUcontext ctx;
    cuDevicePrimaryCtxRetain(&ctx, dev);
    cuCtxPushCurrent(ctx);

    /* Composer. */
    cfc_composer_config_t ccfg = {
        .width = out_w,
        .height = out_h,
        .cells = cells,
        .num_cells = num_cells,
        .cuda_device = 0,
    };
    cfc_composer_t *comp = NULL;
    if (cfc_composer_create(&ccfg, &comp) != 0) {
        fprintf(stderr, "cfc_composer_create failed\n");
        return 1;
    }
    fprintf(stderr, "[grid_record] composer %dx%d, %d ячеек\n",
            out_w, out_h, num_cells);

    /* Encoder. */
    cfc_encoder_config_t ecfg = {
        .cuda_ctx = ctx,
        .width = out_w,
        .height = out_h,
        .fps_num = fps,
        .fps_den = 1,
        .bitrate_kbps = bitrate,
        .gop_size = fps,
        .num_b_frames = 0,
        .preset = "ll",
    };
    cfc_encoder_t *enc = NULL;
    if (cfc_encoder_create(&ecfg, &enc) != 0) {
        fprintf(stderr, "cfc_encoder_create failed\n");
        cfc_composer_destroy(comp);
        return 1;
    }

    /* Output file. */
    write_ctx_t wctx = { 0 };
    wctx.fp = fopen(out_path, "wb");
    if (!wctx.fp) {
        fprintf(stderr, "fopen(%s): %s\n", out_path, strerror(errno));
        cfc_encoder_destroy(enc);
        cfc_composer_destroy(comp);
        return 1;
    }
    fprintf(stderr, "[grid_record] начало записи в %s (Ctrl+C для остановки)\n",
            out_path);

    /* Main loop — frame cadence по wall clock'у. */
    struct timespec ts_start;
    clock_gettime(CLOCK_MONOTONIC, &ts_start);
    int64_t start_us = (int64_t)ts_start.tv_sec * 1000000 + ts_start.tv_nsec / 1000;
    int64_t frame_us = 1000000 / fps;
    int64_t next_us = start_us;

    while (!g_stop) {
        struct timespec now;
        clock_gettime(CLOCK_MONOTONIC, &now);
        int64_t now_us = (int64_t)now.tv_sec * 1000000 + now.tv_nsec / 1000;
        if (now_us < next_us) {
            int64_t sleep_us = next_us - now_us;
            if (sleep_us > 1000000) sleep_us = 1000000;
            struct timespec ts = {
                .tv_sec = sleep_us / 1000000,
                .tv_nsec = (sleep_us % 1000000) * 1000,
            };
            nanosleep(&ts, NULL);
            continue;
        }
        next_us += frame_us;

        CUdeviceptr out_y = 0;
        int out_pitch = 0, oW = 0, oH = 0;
        if (cfc_composer_compose(comp, &out_y, &out_pitch, &oW, &oH) != 0) {
            fprintf(stderr, "[grid_record] compose failed\n");
            break;
        }

        int64_t pts_ns = (now_us - start_us) * 1000;
        if (cfc_encoder_encode_frame(enc, out_y, out_pitch, pts_ns,
                                      on_bitstream, &wctx) != 0) {
            fprintf(stderr, "[grid_record] encode failed\n");
            break;
        }

        if (wctx.frames_encoded > 0 && wctx.frames_encoded % 50 == 0) {
            double elapsed = (now_us - start_us) / 1e6;
            cfc_composer_health_t h;
            cfc_composer_get_health(comp, &h);
            fprintf(stderr,
                "[grid_record] %llu кадров, %llu IDR, %.1f МБ за %.1fс (%.1f fps) | "
                "src active=%d stale=%d dead=%d\n",
                (unsigned long long)wctx.frames_encoded,
                (unsigned long long)wctx.idr_count,
                wctx.bytes_written / 1048576.0,
                elapsed,
                wctx.frames_encoded / elapsed,
                h.active, h.stale, h.dead);
        }

        if (max_seconds > 0 && (now_us - start_us) / 1000000 >= max_seconds) {
            fprintf(stderr, "[grid_record] лимит %dс\n", max_seconds);
            break;
        }
    }

    fprintf(stderr, "[grid_record] flush encoder\n");
    cfc_encoder_flush(enc, on_bitstream, &wctx);

    fprintf(stderr,
        "[grid_record] итого: %llu кадров, %llu IDR, %.2f МБ\n",
        (unsigned long long)wctx.frames_encoded,
        (unsigned long long)wctx.idr_count,
        wctx.bytes_written / 1048576.0);

    fclose(wctx.fp);
    cfc_encoder_destroy(enc);
    cfc_composer_destroy(comp);
    cuCtxPopCurrent(NULL);
    cuDevicePrimaryCtxRelease(dev);
    return 0;
}