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