]>
Commit | Line | Data |
---|---|---|
7a340e0d NA |
1 | import copy |
2 | import heapq | |
3 | import os | |
4 | ||
5 | from .common import PostProcessor | |
6 | from .ffmpeg import ( | |
7 | FFmpegPostProcessor, | |
8 | FFmpegSubtitlesConvertorPP | |
9 | ) | |
10 | from .sponsorblock import SponsorBlockPP | |
11 | from ..utils import ( | |
7a340e0d NA |
12 | orderedSet, |
13 | PostProcessingError, | |
14 | prepend_extension, | |
7a340e0d NA |
15 | ) |
16 | ||
17 | ||
c6af2dd8 | 18 | _TINY_CHAPTER_DURATION = 1 |
7a340e0d NA |
19 | DEFAULT_SPONSORBLOCK_CHAPTER_TITLE = '[SponsorBlock]: %(category_names)l' |
20 | ||
21 | ||
22 | class ModifyChaptersPP(FFmpegPostProcessor): | |
2d9ec704 | 23 | def __init__(self, downloader, remove_chapters_patterns=None, remove_sponsor_segments=None, remove_ranges=None, |
24 | *, sponsorblock_chapter_title=DEFAULT_SPONSORBLOCK_CHAPTER_TITLE, force_keyframes=False): | |
7a340e0d NA |
25 | FFmpegPostProcessor.__init__(self, downloader) |
26 | self._remove_chapters_patterns = set(remove_chapters_patterns or []) | |
8157a09d | 27 | self._remove_sponsor_segments = set(remove_sponsor_segments or []) - set(SponsorBlockPP.POI_CATEGORIES.keys()) |
2d9ec704 | 28 | self._ranges_to_remove = set(remove_ranges or []) |
7a340e0d NA |
29 | self._sponsorblock_chapter_title = sponsorblock_chapter_title |
30 | self._force_keyframes = force_keyframes | |
31 | ||
32 | @PostProcessor._restrict_to(images=False) | |
33 | def run(self, info): | |
e619d8a7 | 34 | # Chapters must be preserved intact when downloading multiple formats of the same video. |
7a340e0d | 35 | chapters, sponsor_chapters = self._mark_chapters_to_remove( |
e619d8a7 NA |
36 | copy.deepcopy(info.get('chapters')) or [], |
37 | copy.deepcopy(info.get('sponsorblock_chapters')) or []) | |
7a340e0d NA |
38 | if not chapters and not sponsor_chapters: |
39 | return [], info | |
40 | ||
5ce1d13e | 41 | real_duration = self._get_real_video_duration(info['filepath']) |
7a340e0d NA |
42 | if not chapters: |
43 | chapters = [{'start_time': 0, 'end_time': real_duration, 'title': info['title']}] | |
44 | ||
45 | info['chapters'], cuts = self._remove_marked_arrange_sponsors(chapters + sponsor_chapters) | |
46 | if not cuts: | |
47 | return [], info | |
48 | ||
165efb82 | 49 | if self._duration_mismatch(real_duration, info.get('duration')): |
50 | if not self._duration_mismatch(real_duration, info['chapters'][-1]['end_time']): | |
7a340e0d NA |
51 | self.to_screen(f'Skipping {self.pp_key()} since the video appears to be already cut') |
52 | return [], info | |
53 | if not info.get('__real_download'): | |
54 | raise PostProcessingError('Cannot cut video since the real and expected durations mismatch. ' | |
55 | 'Different chapters may have already been removed') | |
7a340e0d NA |
56 | else: |
57 | self.write_debug('Expected and actual durations mismatch') | |
58 | ||
59 | concat_opts = self._make_concat_opts(cuts, real_duration) | |
60 | ||
61 | def remove_chapters(file, is_sub): | |
62 | return file, self.remove_chapters(file, cuts, concat_opts, self._force_keyframes and not is_sub) | |
63 | ||
64 | in_out_files = [remove_chapters(info['filepath'], False)] | |
65 | in_out_files.extend(remove_chapters(in_file, True) for in_file in self._get_supported_subs(info)) | |
66 | ||
67 | # Renaming should only happen after all files are processed | |
68 | files_to_remove = [] | |
69 | for in_file, out_file in in_out_files: | |
70 | uncut_file = prepend_extension(in_file, 'uncut') | |
71 | os.replace(in_file, uncut_file) | |
72 | os.replace(out_file, in_file) | |
73 | files_to_remove.append(uncut_file) | |
74 | ||
75 | return files_to_remove, info | |
76 | ||
77 | def _mark_chapters_to_remove(self, chapters, sponsor_chapters): | |
78 | if self._remove_chapters_patterns: | |
79 | warn_no_chapter_to_remove = True | |
80 | if not chapters: | |
81 | self.to_screen('Chapter information is unavailable') | |
82 | warn_no_chapter_to_remove = False | |
83 | for c in chapters: | |
84 | if any(regex.search(c['title']) for regex in self._remove_chapters_patterns): | |
85 | c['remove'] = True | |
86 | warn_no_chapter_to_remove = False | |
87 | if warn_no_chapter_to_remove: | |
88 | self.to_screen('There are no chapters matching the regex') | |
89 | ||
90 | if self._remove_sponsor_segments: | |
91 | warn_no_chapter_to_remove = True | |
92 | if not sponsor_chapters: | |
93 | self.to_screen('SponsorBlock information is unavailable') | |
94 | warn_no_chapter_to_remove = False | |
95 | for c in sponsor_chapters: | |
96 | if c['category'] in self._remove_sponsor_segments: | |
97 | c['remove'] = True | |
98 | warn_no_chapter_to_remove = False | |
99 | if warn_no_chapter_to_remove: | |
100 | self.to_screen('There are no matching SponsorBlock chapters') | |
101 | ||
2d9ec704 | 102 | sponsor_chapters.extend({ |
103 | 'start_time': start, | |
104 | 'end_time': end, | |
105 | 'category': 'manually_removed', | |
106 | '_categories': [('manually_removed', start, end)], | |
107 | 'remove': True, | |
108 | } for start, end in self._ranges_to_remove) | |
109 | ||
7a340e0d NA |
110 | return chapters, sponsor_chapters |
111 | ||
7a340e0d NA |
112 | def _get_supported_subs(self, info): |
113 | for sub in (info.get('requested_subtitles') or {}).values(): | |
114 | sub_file = sub.get('filepath') | |
115 | # The file might have been removed by --embed-subs | |
116 | if not sub_file or not os.path.exists(sub_file): | |
117 | continue | |
118 | ext = sub['ext'] | |
119 | if ext not in FFmpegSubtitlesConvertorPP.SUPPORTED_EXTS: | |
120 | self.report_warning(f'Cannot remove chapters from external {ext} subtitles; "{sub_file}" is now out of sync') | |
121 | continue | |
122 | # TODO: create __real_download for subs? | |
123 | yield sub_file | |
124 | ||
125 | def _remove_marked_arrange_sponsors(self, chapters): | |
126 | # Store cuts separately, since adjacent and overlapping cuts must be merged. | |
127 | cuts = [] | |
128 | ||
129 | def append_cut(c): | |
e619d8a7 | 130 | assert 'remove' in c, 'Not a cut is appended to cuts' |
7a340e0d NA |
131 | last_to_cut = cuts[-1] if cuts else None |
132 | if last_to_cut and last_to_cut['end_time'] >= c['start_time']: | |
133 | last_to_cut['end_time'] = max(last_to_cut['end_time'], c['end_time']) | |
134 | else: | |
135 | cuts.append(c) | |
136 | return len(cuts) - 1 | |
137 | ||
138 | def excess_duration(c): | |
139 | # Cuts that are completely within the chapter reduce chapters' duration. | |
140 | # Since cuts can overlap, excess duration may be less that the sum of cuts' durations. | |
141 | # To avoid that, chapter stores the index to the fist cut within the chapter, | |
142 | # instead of storing excess duration. append_cut ensures that subsequent cuts (if any) | |
143 | # will be merged with previous ones (if necessary). | |
144 | cut_idx, excess = c.pop('cut_idx', len(cuts)), 0 | |
145 | while cut_idx < len(cuts): | |
146 | cut = cuts[cut_idx] | |
147 | if cut['start_time'] >= c['end_time']: | |
148 | break | |
149 | if cut['end_time'] > c['start_time']: | |
150 | excess += min(cut['end_time'], c['end_time']) | |
151 | excess -= max(cut['start_time'], c['start_time']) | |
152 | cut_idx += 1 | |
153 | return excess | |
154 | ||
155 | new_chapters = [] | |
156 | ||
7a340e0d | 157 | def append_chapter(c): |
e619d8a7 | 158 | assert 'remove' not in c, 'Cut is appended to chapters' |
c6af2dd8 | 159 | length = c['end_time'] - c['start_time'] - excess_duration(c) |
7a340e0d NA |
160 | # Chapter is completely covered by cuts or sponsors. |
161 | if length <= 0: | |
162 | return | |
163 | start = new_chapters[-1]['end_time'] if new_chapters else 0 | |
164 | c.update(start_time=start, end_time=start + length) | |
c6af2dd8 | 165 | new_chapters.append(c) |
7a340e0d NA |
166 | |
167 | # Turn into a priority queue, index is a tie breaker. | |
168 | # Plain stack sorted by start_time is not enough: after splitting the chapter, | |
169 | # the part returned to the stack is not guaranteed to have start_time | |
170 | # less than or equal to the that of the stack's head. | |
171 | chapters = [(c['start_time'], i, c) for i, c in enumerate(chapters)] | |
172 | heapq.heapify(chapters) | |
173 | ||
174 | _, cur_i, cur_chapter = heapq.heappop(chapters) | |
175 | while chapters: | |
176 | _, i, c = heapq.heappop(chapters) | |
177 | # Non-overlapping chapters or cuts can be appended directly. However, | |
178 | # adjacent non-overlapping cuts must be merged, which is handled by append_cut. | |
179 | if cur_chapter['end_time'] <= c['start_time']: | |
180 | (append_chapter if 'remove' not in cur_chapter else append_cut)(cur_chapter) | |
181 | cur_i, cur_chapter = i, c | |
182 | continue | |
183 | ||
184 | # Eight possibilities for overlapping chapters: (cut, cut), (cut, sponsor), | |
185 | # (cut, normal), (sponsor, cut), (normal, cut), (sponsor, sponsor), | |
186 | # (sponsor, normal), and (normal, sponsor). There is no (normal, normal): | |
187 | # normal chapters are assumed not to overlap. | |
188 | if 'remove' in cur_chapter: | |
189 | # (cut, cut): adjust end_time. | |
190 | if 'remove' in c: | |
191 | cur_chapter['end_time'] = max(cur_chapter['end_time'], c['end_time']) | |
192 | # (cut, sponsor/normal): chop the beginning of the later chapter | |
193 | # (if it's not completely hidden by the cut). Push to the priority queue | |
194 | # to restore sorting by start_time: with beginning chopped, c may actually | |
195 | # start later than the remaining chapters from the queue. | |
196 | elif cur_chapter['end_time'] < c['end_time']: | |
197 | c['start_time'] = cur_chapter['end_time'] | |
198 | c['_was_cut'] = True | |
199 | heapq.heappush(chapters, (c['start_time'], i, c)) | |
200 | # (sponsor/normal, cut). | |
201 | elif 'remove' in c: | |
202 | cur_chapter['_was_cut'] = True | |
203 | # Chop the end of the current chapter if the cut is not contained within it. | |
204 | # Chopping the end doesn't break start_time sorting, no PQ push is necessary. | |
205 | if cur_chapter['end_time'] <= c['end_time']: | |
206 | cur_chapter['end_time'] = c['start_time'] | |
207 | append_chapter(cur_chapter) | |
208 | cur_i, cur_chapter = i, c | |
209 | continue | |
210 | # Current chapter contains the cut within it. If the current chapter is | |
211 | # a sponsor chapter, check whether the categories before and after the cut differ. | |
212 | if '_categories' in cur_chapter: | |
213 | after_c = dict(cur_chapter, start_time=c['end_time'], _categories=[]) | |
214 | cur_cats = [] | |
215 | for cat_start_end in cur_chapter['_categories']: | |
216 | if cat_start_end[1] < c['start_time']: | |
217 | cur_cats.append(cat_start_end) | |
218 | if cat_start_end[2] > c['end_time']: | |
219 | after_c['_categories'].append(cat_start_end) | |
220 | cur_chapter['_categories'] = cur_cats | |
221 | if cur_chapter['_categories'] != after_c['_categories']: | |
222 | # Categories before and after the cut differ: push the after part to PQ. | |
223 | heapq.heappush(chapters, (after_c['start_time'], cur_i, after_c)) | |
224 | cur_chapter['end_time'] = c['start_time'] | |
225 | append_chapter(cur_chapter) | |
226 | cur_i, cur_chapter = i, c | |
227 | continue | |
228 | # Either sponsor categories before and after the cut are the same or | |
229 | # we're dealing with a normal chapter. Just register an outstanding cut: | |
230 | # subsequent append_chapter will reduce the duration. | |
231 | cur_chapter.setdefault('cut_idx', append_cut(c)) | |
232 | # (sponsor, normal): if a normal chapter is not completely overlapped, | |
233 | # chop the beginning of it and push it to PQ. | |
234 | elif '_categories' in cur_chapter and '_categories' not in c: | |
235 | if cur_chapter['end_time'] < c['end_time']: | |
236 | c['start_time'] = cur_chapter['end_time'] | |
237 | c['_was_cut'] = True | |
238 | heapq.heappush(chapters, (c['start_time'], i, c)) | |
239 | # (normal, sponsor) and (sponsor, sponsor) | |
240 | else: | |
e619d8a7 | 241 | assert '_categories' in c, 'Normal chapters overlap' |
7a340e0d NA |
242 | cur_chapter['_was_cut'] = True |
243 | c['_was_cut'] = True | |
244 | # Push the part after the sponsor to PQ. | |
245 | if cur_chapter['end_time'] > c['end_time']: | |
246 | # deepcopy to make categories in after_c and cur_chapter/c refer to different lists. | |
247 | after_c = dict(copy.deepcopy(cur_chapter), start_time=c['end_time']) | |
248 | heapq.heappush(chapters, (after_c['start_time'], cur_i, after_c)) | |
249 | # Push the part after the overlap to PQ. | |
250 | elif c['end_time'] > cur_chapter['end_time']: | |
251 | after_cur = dict(copy.deepcopy(c), start_time=cur_chapter['end_time']) | |
252 | heapq.heappush(chapters, (after_cur['start_time'], cur_i, after_cur)) | |
253 | c['end_time'] = cur_chapter['end_time'] | |
254 | # (sponsor, sponsor): merge categories in the overlap. | |
255 | if '_categories' in cur_chapter: | |
256 | c['_categories'] = cur_chapter['_categories'] + c['_categories'] | |
257 | # Inherit the cuts that the current chapter has accumulated within it. | |
258 | if 'cut_idx' in cur_chapter: | |
259 | c['cut_idx'] = cur_chapter['cut_idx'] | |
260 | cur_chapter['end_time'] = c['start_time'] | |
261 | append_chapter(cur_chapter) | |
262 | cur_i, cur_chapter = i, c | |
263 | (append_chapter if 'remove' not in cur_chapter else append_cut)(cur_chapter) | |
c6af2dd8 NA |
264 | return self._remove_tiny_rename_sponsors(new_chapters), cuts |
265 | ||
266 | def _remove_tiny_rename_sponsors(self, chapters): | |
267 | new_chapters = [] | |
268 | for i, c in enumerate(chapters): | |
269 | # Merge with the previous/next if the chapter is tiny. | |
270 | # Only tiny chapters resulting from a cut can be skipped. | |
271 | # Chapters that were already tiny in the original list will be preserved. | |
272 | if (('_was_cut' in c or '_categories' in c) | |
273 | and c['end_time'] - c['start_time'] < _TINY_CHAPTER_DURATION): | |
274 | if not new_chapters: | |
275 | # Prepend tiny chapter to the next one if possible. | |
276 | if i < len(chapters) - 1: | |
277 | chapters[i + 1]['start_time'] = c['start_time'] | |
278 | continue | |
279 | else: | |
280 | old_c = new_chapters[-1] | |
281 | if i < len(chapters) - 1: | |
282 | next_c = chapters[i + 1] | |
283 | # Not a typo: key names in old_c and next_c are really different. | |
284 | prev_is_sponsor = 'categories' in old_c | |
285 | next_is_sponsor = '_categories' in next_c | |
286 | # Preferentially prepend tiny normals to normals and sponsors to sponsors. | |
287 | if (('_categories' not in c and prev_is_sponsor and not next_is_sponsor) | |
288 | or ('_categories' in c and not prev_is_sponsor and next_is_sponsor)): | |
289 | next_c['start_time'] = c['start_time'] | |
290 | continue | |
291 | old_c['end_time'] = c['end_time'] | |
292 | continue | |
7a340e0d | 293 | |
7a340e0d NA |
294 | c.pop('_was_cut', None) |
295 | cats = c.pop('_categories', None) | |
296 | if cats: | |
297 | category = min(cats, key=lambda c: c[2] - c[1])[0] | |
298 | cats = orderedSet(x[0] for x in cats) | |
299 | c.update({ | |
300 | 'category': category, | |
301 | 'categories': cats, | |
302 | 'name': SponsorBlockPP.CATEGORIES[category], | |
303 | 'category_names': [SponsorBlockPP.CATEGORIES[c] for c in cats] | |
304 | }) | |
8157a09d | 305 | c['title'] = self._downloader.evaluate_outtmpl(self._sponsorblock_chapter_title, c.copy()) |
c6af2dd8 NA |
306 | # Merge identically named sponsors. |
307 | if (new_chapters and 'categories' in new_chapters[-1] | |
308 | and new_chapters[-1]['title'] == c['title']): | |
309 | new_chapters[-1]['end_time'] = c['end_time'] | |
310 | continue | |
311 | new_chapters.append(c) | |
312 | return new_chapters | |
7a340e0d NA |
313 | |
314 | def remove_chapters(self, filename, ranges_to_cut, concat_opts, force_keyframes=False): | |
315 | in_file = filename | |
316 | out_file = prepend_extension(in_file, 'temp') | |
317 | if force_keyframes: | |
165efb82 | 318 | in_file = self.force_keyframes(in_file, (t for c in ranges_to_cut for t in (c['start_time'], c['end_time']))) |
7a340e0d NA |
319 | self.to_screen(f'Removing chapters from {filename}') |
320 | self.concat_files([in_file] * len(concat_opts), out_file, concat_opts) | |
321 | if in_file != filename: | |
322 | os.remove(in_file) | |
323 | return out_file | |
324 | ||
325 | @staticmethod | |
326 | def _make_concat_opts(chapters_to_remove, duration): | |
327 | opts = [{}] | |
328 | for s in chapters_to_remove: | |
329 | # Do not create 0 duration chunk at the beginning. | |
330 | if s['start_time'] == 0: | |
331 | opts[-1]['inpoint'] = f'{s["end_time"]:.6f}' | |
332 | continue | |
333 | opts[-1]['outpoint'] = f'{s["start_time"]:.6f}' | |
334 | # Do not create 0 duration chunk at the end. | |
335 | if s['end_time'] != duration: | |
336 | opts.append({'inpoint': f'{s["end_time"]:.6f}'}) | |
337 | return opts |