yt_dlp/postprocessor/modify_chapters.py

   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 (
  12     orderedSet,
  13     PostProcessingError,
  14     prepend_extension,
  15 )
  16
  17
  18 _TINY_CHAPTER_DURATION = 1
  19 DEFAULT_SPONSORBLOCK_CHAPTER_TITLE = '[SponsorBlock]: %(category_names)l'
  20
  21
  22 class ModifyChaptersPP(FFmpegPostProcessor):
  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):
  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 []) - set(SponsorBlockPP.POI_CATEGORIES.keys())
  28         self._ranges_to_remove = set(remove_ranges or [])
  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):
  34         # Chapters must be preserved intact when downloading multiple formats of the same video.
  35         chapters, sponsor_chapters = self._mark_chapters_to_remove(
  36             copy.deepcopy(info.get('chapters')) or [],
  37             copy.deepcopy(info.get('sponsorblock_chapters')) or [])
  38         if not chapters and not sponsor_chapters:
  39             return [], info
  40
  41         real_duration = self._get_real_video_duration(info['filepath'])
  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
  49         if self._duration_mismatch(real_duration, info.get('duration')):
  50             if not self._duration_mismatch(real_duration, info['chapters'][-1]['end_time']):
  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')
  56             else:
  57                 self.write_debug('Expected and actual durations mismatch')
  58
  59         concat_opts = self._make_concat_opts(cuts, real_duration)
  60         self.write_debug('Concat spec = %s' % ', '.join(f'{c.get("inpoint", 0.0)}-{c.get("outpoint", "inf")}' for c in concat_opts))
  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:
  71             uncut_file = prepend_extension(in_file, 'uncut')
  72             os.replace(in_file, uncut_file)
  73             os.replace(out_file, in_file)
  74             files_to_remove.append(uncut_file)
  75
  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
 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
 111         return chapters, sponsor_chapters
 112
 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):
 131             assert 'remove' in c, 'Not a cut is appended to cuts'
 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
 158         def append_chapter(c):
 159             assert 'remove' not in c, 'Cut is appended to chapters'
 160             length = c['end_time'] - c['start_time'] - excess_duration(c)
 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)
 166             new_chapters.append(c)
 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:
 242                 assert '_categories' in c, 'Normal chapters overlap'
 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)
 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
 294
 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                 })
 306                 c['title'] = self._downloader.evaluate_outtmpl(self._sponsorblock_chapter_title, c.copy())
 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
 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:
 319             in_file = self.force_keyframes(in_file, (t for c in ranges_to_cut for t in (c['start_time'], c['end_time'])))
 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