git.sesse.net Git - nageru/blob - futatabi/export.cpp

   1 #include "clip_list.h"
   2 #include "defs.h"
   3 #include "export.h"
   4 #include "flags.h"
   5 #include "frame_on_disk.h"
   6 #include "shared/ffmpeg_raii.h"
   7 #include "shared/timebase.h"
   8
   9 #include <QMessageBox>
  10 #include <QProgressDialog>
  11
  12 #include <vector>
  13
  14 #include <unistd.h>
  15
  16 extern "C" {
  17 #include <libavformat/avformat.h>
  18 }
  19
  20 using namespace std;
  21
  22 namespace {
  23
  24 // Only used in export_cliplist_clip_multitrack_triggered.
  25 struct BufferedJPEG {
  26         int64_t pts;
  27         unsigned stream_idx;
  28         string jpeg;
  29 };
  30
  31 bool write_buffered_jpegs(AVFormatContext *avctx, const vector<BufferedJPEG> &buffered_jpegs)
  32 {
  33         for (const BufferedJPEG &jpeg : buffered_jpegs) {
  34                 AVPacket pkt;
  35                 av_init_packet(&pkt);
  36                 pkt.stream_index = jpeg.stream_idx;
  37                 pkt.data = (uint8_t *)jpeg.jpeg.data();
  38                 pkt.size = jpeg.jpeg.size();
  39                 pkt.pts = jpeg.pts;
  40                 pkt.dts = jpeg.pts;
  41                 pkt.flags = AV_PKT_FLAG_KEY;
  42
  43                 if (av_write_frame(avctx, &pkt) < 0) {
  44                         return false;
  45                 }
  46         }
  47         return true;
  48 }
  49
  50 }  // namespace
  51
  52 void export_multitrack_clip(const string &filename, const Clip &clip)
  53 {
  54         AVFormatContext *avctx = nullptr;
  55         avformat_alloc_output_context2(&avctx, NULL, NULL, filename.c_str());
  56         if (avctx == nullptr) {
  57                 QMessageBox msgbox;
  58                 msgbox.setText("Could not allocate FFmpeg context");
  59                 msgbox.exec();
  60                 return;
  61         }
  62         AVFormatContextWithCloser closer(avctx);
  63
  64         int ret = avio_open(&avctx->pb, filename.c_str(), AVIO_FLAG_WRITE);
  65         if (ret < 0) {
  66                 QMessageBox msgbox;
  67                 msgbox.setText(QString::fromStdString("Could not open output file '" + filename + "'"));
  68                 msgbox.exec();
  69                 return;
  70         }
  71
  72         // Find the first frame for each stream.
  73         size_t num_frames = 0;
  74         size_t num_streams_with_frames_left = 0;
  75         size_t last_stream_idx = 0;
  76         FrameReader readers[MAX_STREAMS];
  77         bool has_frames[MAX_STREAMS];
  78         size_t first_frame_idx[MAX_STREAMS], last_frame_idx[MAX_STREAMS];  // Inclusive, exclusive.
  79         {
  80                 lock_guard<mutex> lock(frame_mu);
  81                 for (size_t stream_idx = 0; stream_idx < MAX_STREAMS; ++stream_idx) {
  82                         // Find the first frame such that frame.pts <= pts_in.
  83                         auto it = find_first_frame_at_or_after(frames[stream_idx], clip.pts_in);
  84                         first_frame_idx[stream_idx] = distance(frames[stream_idx].begin(), it);
  85                         has_frames[stream_idx] = (it != frames[stream_idx].end());
  86
  87                         // Find the first frame such that frame.pts >= pts_out.
  88                         it = find_first_frame_at_or_after(frames[stream_idx], clip.pts_out);
  89                         last_frame_idx[stream_idx] = distance(frames[stream_idx].begin(), it);
  90                         num_frames += last_frame_idx[stream_idx] - first_frame_idx[stream_idx];
  91
  92                         if (has_frames[stream_idx]) {
  93                                 ++num_streams_with_frames_left;
  94                                 last_stream_idx = stream_idx;
  95                         }
  96                 }
  97         }
  98
  99         // Create the streams. Note that some of them could be without frames
 100         // (we try to maintain the stream indexes in the export).
 101         vector<AVStream *> video_streams;
 102         for (unsigned stream_idx = 0; stream_idx <= last_stream_idx; ++stream_idx) {
 103                 AVStream *avstream_video = avformat_new_stream(avctx, nullptr);
 104                 if (avstream_video == nullptr) {
 105                         fprintf(stderr, "avformat_new_stream() failed\n");
 106                         exit(1);
 107                 }
 108                 avstream_video->time_base = AVRational{1, TIMEBASE};
 109                 avstream_video->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
 110                 avstream_video->codecpar->codec_id = AV_CODEC_ID_MJPEG;
 111                 avstream_video->codecpar->width = global_flags.width;  // Might be wrong, but doesn't matter all that much.
 112                 avstream_video->codecpar->height = global_flags.height;
 113
 114                 // TODO: Deduplicate this against Mux.
 115                 avstream_video->codecpar->color_primaries = AVCOL_PRI_BT709;  // RGB colorspace (inout_format.color_space).
 116                 avstream_video->codecpar->color_trc = AVCOL_TRC_IEC61966_2_1;  // Gamma curve (inout_format.gamma_curve).
 117                 // YUV colorspace (output_ycbcr_format.luma_coefficients).
 118                 avstream_video->codecpar->color_space = AVCOL_SPC_BT709;
 119                 avstream_video->codecpar->color_range = AVCOL_RANGE_MPEG;  // Full vs. limited range (output_ycbcr_format.full_range).
 120                 avstream_video->codecpar->chroma_location = AVCHROMA_LOC_LEFT;  // Chroma sample location. See chroma_offset_0[] in Mixer::subsample_chroma().
 121                 avstream_video->codecpar->field_order = AV_FIELD_PROGRESSIVE;
 122                 video_streams.push_back(avstream_video);
 123         }
 124
 125         if (avformat_write_header(avctx, nullptr) < 0) {
 126                 QMessageBox msgbox;
 127                 msgbox.setText("Writing header failed");
 128                 msgbox.exec();
 129                 unlink(filename.c_str());
 130                 return;
 131         }
 132
 133         QProgressDialog progress(QString::fromStdString("Exporting to " + filename + "..."), "Abort", 0, 1);
 134         progress.setWindowTitle("Futatabi");
 135         progress.setWindowModality(Qt::WindowModal);
 136         progress.setMinimumDuration(1000);
 137         progress.setMaximum(num_frames);
 138         progress.setValue(0);
 139
 140         // We buffer up to 1000 frames at a time, in a hope that we can reduce
 141         // the amount of seeking needed on rotational media.
 142         vector<BufferedJPEG> buffered_jpegs;
 143         size_t frames_written = 0;
 144         while (num_streams_with_frames_left > 0) {
 145                 // Find the stream with the lowest frame. Lower stream indexes win.
 146                 FrameOnDisk first_frame;
 147                 unsigned first_frame_stream_idx = 0;
 148                 {
 149                         lock_guard<mutex> lock(frame_mu);
 150                         for (size_t stream_idx = 0; stream_idx < MAX_STREAMS; ++stream_idx) {
 151                                 if (!has_frames[stream_idx]) {
 152                                         continue;
 153                                 }
 154                                 if (first_frame.pts == -1 || frames[stream_idx][first_frame_idx[stream_idx]].pts < first_frame.pts) {
 155                                         first_frame = frames[stream_idx][first_frame_idx[stream_idx]];
 156                                         first_frame_stream_idx = stream_idx;
 157                                 }
 158                         }
 159                         ++first_frame_idx[first_frame_stream_idx];
 160                         if (first_frame_idx[first_frame_stream_idx] >= last_frame_idx[first_frame_stream_idx]) {
 161                                 has_frames[first_frame_stream_idx] = false;
 162                                 --num_streams_with_frames_left;
 163                         }
 164                 }
 165                 string jpeg = readers[first_frame_stream_idx].read_frame(first_frame);
 166                 int64_t scaled_pts = av_rescale_q(first_frame.pts, AVRational{1, TIMEBASE},
 167                         video_streams[first_frame_stream_idx]->time_base);
 168                 buffered_jpegs.emplace_back(BufferedJPEG{ scaled_pts, first_frame_stream_idx, std::move(jpeg) });
 169                 if (buffered_jpegs.size() >= 1000) {
 170                         if (!write_buffered_jpegs(avctx, buffered_jpegs)) {
 171                                 QMessageBox msgbox;
 172                                 msgbox.setText("Writing frames failed");
 173                                 msgbox.exec();
 174                                 unlink(filename.c_str());
 175                                 return;
 176                         }
 177                         frames_written += buffered_jpegs.size();
 178                         progress.setValue(frames_written);
 179                         buffered_jpegs.clear();
 180                 }
 181                 if (progress.wasCanceled()) {
 182                         unlink(filename.c_str());
 183                         return;
 184                 }
 185         }
 186
 187         if (!write_buffered_jpegs(avctx, buffered_jpegs)) {
 188                 QMessageBox msgbox;
 189                 msgbox.setText("Writing frames failed");
 190                 msgbox.exec();
 191                 unlink(filename.c_str());
 192                 return;
 193         }
 194         frames_written += buffered_jpegs.size();
 195         progress.setValue(frames_written);
 196 }