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