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 FFmpegPostProcessor, FFmpegSubtitlesConvertorPP
   7 from .sponsorblock import SponsorBlockPP
   8 from ..utils import PostProcessingError, orderedSet, prepend_extension
   9
  10 _TINY_CHAPTER_DURATION = 1
  11 DEFAULT_SPONSORBLOCK_CHAPTER_TITLE = '[SponsorBlock]: %(category_names)l'
  12
  13
  14 class ModifyChaptersPP(FFmpegPostProcessor):
  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):
  17         FFmpegPostProcessor.__init__(self, downloader)
  18         self._remove_chapters_patterns = set(remove_chapters_patterns or [])
  19         self._remove_sponsor_segments = set(remove_sponsor_segments or []) - set(SponsorBlockPP.NON_SKIPPABLE_CATEGORIES.keys())
  20         self._ranges_to_remove = set(remove_ranges or [])
  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):
  26         self._fixup_chapters(info)
  27         # Chapters must be preserved intact when downloading multiple formats of the same video.
  28         chapters, sponsor_chapters = self._mark_chapters_to_remove(
  29             copy.deepcopy(info.get('chapters')) or [],
  30             copy.deepcopy(info.get('sponsorblock_chapters')) or [])
  31         if not chapters and not sponsor_chapters:
  32             return [], info
  33
  34         real_duration = self._get_real_video_duration(info['filepath'])
  35         if not chapters:
  36             chapters = [{'start_time': 0, 'end_time': info.get('duration') or real_duration, 'title': info['title']}]
  37
  38         info['chapters'], cuts = self._remove_marked_arrange_sponsors(chapters + sponsor_chapters)
  39         if not cuts:
  40             return [], info
  41         elif not info['chapters']:
  42             self.report_warning('You have requested to remove the entire video, which is not possible')
  43             return [], info
  44
  45         original_duration, info['duration'] = info.get('duration'), info['chapters'][-1]['end_time']
  46         if self._duration_mismatch(real_duration, original_duration, 1):
  47             if not self._duration_mismatch(real_duration, info['duration']):
  48                 self.to_screen(f'Skipping {self.pp_key()} since the video appears to be already cut')
  49                 return [], info
  50             if not info.get('__real_download'):
  51                 raise PostProcessingError('Cannot cut video since the real and expected durations mismatch. '
  52                                           'Different chapters may have already been removed')
  53             else:
  54                 self.write_debug('Expected and actual durations mismatch')
  55
  56         concat_opts = self._make_concat_opts(cuts, real_duration)
  57         self.write_debug('Concat spec = {}'.format(', '.join(f'{c.get("inpoint", 0.0)}-{c.get("outpoint", "inf")}' for c in concat_opts)))
  58
  59         def remove_chapters(file, is_sub):
  60             return file, self.remove_chapters(file, cuts, concat_opts, self._force_keyframes and not is_sub)
  61
  62         in_out_files = [remove_chapters(info['filepath'], False)]
  63         in_out_files.extend(remove_chapters(in_file, True) for in_file in self._get_supported_subs(info))
  64
  65         # Renaming should only happen after all files are processed
  66         files_to_remove = []
  67         for in_file, out_file in in_out_files:
  68             mtime = os.stat(in_file).st_mtime
  69             uncut_file = prepend_extension(in_file, 'uncut')
  70             os.replace(in_file, uncut_file)
  71             os.replace(out_file, in_file)
  72             self.try_utime(in_file, mtime, mtime)
  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
 102         sponsor_chapters.extend({
 103             'start_time': start,
 104             'end_time': end,
 105             'category': 'manually_removed',
 106             '_categories': [('manually_removed', start, end, 'Manually removed')],
 107             'remove': True,
 108         } for start, end in self._ranges_to_remove)
 109
 110         return chapters, sponsor_chapters
 111
 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):
 130             assert 'remove' in c, 'Not a cut is appended to cuts'
 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
 157         def append_chapter(c):
 158             assert 'remove' not in c, 'Cut is appended to chapters'
 159             length = c['end_time'] - c['start_time'] - excess_duration(c)
 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)
 165             new_chapters.append(c)
 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:
 241                 assert '_categories' in c, 'Normal chapters overlap'
 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)
 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
 293
 294             c.pop('_was_cut', None)
 295             cats = c.pop('_categories', None)
 296             if cats:
 297                 category, _, _, category_name = min(cats, key=lambda c: c[2] - c[1])
 298                 c.update({
 299                     'category': category,
 300                     'categories': orderedSet(x[0] for x in cats),
 301                     'name': category_name,
 302                     'category_names': orderedSet(x[3] for x in cats),
 303                 })
 304                 c['title'] = self._downloader.evaluate_outtmpl(self._sponsorblock_chapter_title, c.copy())
 305                 # Merge identically named sponsors.
 306                 if (new_chapters and 'categories' in new_chapters[-1]
 307                         and new_chapters[-1]['title'] == c['title']):
 308                     new_chapters[-1]['end_time'] = c['end_time']
 309                     continue
 310             new_chapters.append(c)
 311         return new_chapters
 312
 313     def remove_chapters(self, filename, ranges_to_cut, concat_opts, force_keyframes=False):
 314         in_file = filename
 315         out_file = prepend_extension(in_file, 'temp')
 316         if force_keyframes:
 317             in_file = self.force_keyframes(in_file, (t for c in ranges_to_cut for t in (c['start_time'], c['end_time'])))
 318         self.to_screen(f'Removing chapters from {filename}')
 319         self.concat_files([in_file] * len(concat_opts), out_file, concat_opts)
 320         if in_file != filename:
 321             self._delete_downloaded_files(in_file, msg=None)
 322         return out_file
 323
 324     @staticmethod
 325     def _make_concat_opts(chapters_to_remove, duration):
 326         opts = [{}]
 327         for s in chapters_to_remove:
 328             # Do not create 0 duration chunk at the beginning.
 329             if s['start_time'] == 0:
 330                 opts[-1]['inpoint'] = f'{s["end_time"]:.6f}'
 331                 continue
 332             opts[-1]['outpoint'] = f'{s["start_time"]:.6f}'
 333             # Do not create 0 duration chunk at the end.
 334             if s['end_time'] < duration:
 335                 opts.append({'inpoint': f'{s["end_time"]:.6f}'})
 336         return opts