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