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