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

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