Browse Source

[SponsorBlock] Support `chapter` category (#5260)

Authored by: ajayyy, pukkandan
Ajay Ramachandran 2 years ago
parent
commit
63c547d71c

+ 3 - 3
README.md

@@ -1042,7 +1042,7 @@ Make chapter entries for, or remove various segments (sponsor,
                                     for, separated by commas. Available
                                     for, separated by commas. Available
                                     categories are sponsor, intro, outro,
                                     categories are sponsor, intro, outro,
                                     selfpromo, preview, filler, interaction,
                                     selfpromo, preview, filler, interaction,
-                                    music_offtopic, poi_highlight, all and
+                                    music_offtopic, poi_highlight, chapter, all and
                                     default (=all). You can prefix the category
                                     default (=all). You can prefix the category
                                     with a "-" to exclude it. See [1] for
                                     with a "-" to exclude it. See [1] for
                                     description of the categories. E.g.
                                     description of the categories. E.g.
@@ -1054,8 +1054,8 @@ Make chapter entries for, or remove various segments (sponsor,
                                     remove takes precedence. The syntax and
                                     remove takes precedence. The syntax and
                                     available categories are the same as for
                                     available categories are the same as for
                                     --sponsorblock-mark except that "default"
                                     --sponsorblock-mark except that "default"
-                                    refers to "all,-filler" and poi_highlight is
-                                    not available
+                                    refers to "all,-filler" and poi_highlight and
+                                    chapter are not available
     --sponsorblock-chapter-title TEMPLATE
     --sponsorblock-chapter-title TEMPLATE
                                     An output template for the title of the
                                     An output template for the title of the
                                     SponsorBlock chapters created by
                                     SponsorBlock chapters created by

+ 26 - 8
test/test_postprocessors.py

@@ -16,6 +16,7 @@ from yt_dlp.postprocessor import (
     MetadataFromFieldPP,
     MetadataFromFieldPP,
     MetadataParserPP,
     MetadataParserPP,
     ModifyChaptersPP,
     ModifyChaptersPP,
+    SponsorBlockPP,
 )
 )
 
 
 
 
@@ -76,11 +77,15 @@ class TestModifyChaptersPP(unittest.TestCase):
         self._pp = ModifyChaptersPP(YoutubeDL())
         self._pp = ModifyChaptersPP(YoutubeDL())
 
 
     @staticmethod
     @staticmethod
-    def _sponsor_chapter(start, end, cat, remove=False):
-        c = {'start_time': start, 'end_time': end, '_categories': [(cat, start, end)]}
-        if remove:
-            c['remove'] = True
-        return c
+    def _sponsor_chapter(start, end, cat, remove=False, title=None):
+        if title is None:
+            title = SponsorBlockPP.CATEGORIES[cat]
+        return {
+            'start_time': start,
+            'end_time': end,
+            '_categories': [(cat, start, end, title)],
+            **({'remove': True} if remove else {}),
+        }
 
 
     @staticmethod
     @staticmethod
     def _chapter(start, end, title=None, remove=False):
     def _chapter(start, end, title=None, remove=False):
@@ -130,6 +135,19 @@ class TestModifyChaptersPP(unittest.TestCase):
              'c', '[SponsorBlock]: Filler Tangent', 'c'])
              'c', '[SponsorBlock]: Filler Tangent', 'c'])
         self._remove_marked_arrange_sponsors_test_impl(chapters, expected, [])
         self._remove_marked_arrange_sponsors_test_impl(chapters, expected, [])
 
 
+    def test_remove_marked_arrange_sponsors_SponsorBlockChapters(self):
+        chapters = self._chapters([70], ['c']) + [
+            self._sponsor_chapter(10, 20, 'chapter', title='sb c1'),
+            self._sponsor_chapter(15, 16, 'chapter', title='sb c2'),
+            self._sponsor_chapter(30, 40, 'preview'),
+            self._sponsor_chapter(50, 60, 'filler')]
+        expected = self._chapters(
+            [10, 15, 16, 20, 30, 40, 50, 60, 70],
+            ['c', '[SponsorBlock]: sb c1', '[SponsorBlock]: sb c1, sb c2', '[SponsorBlock]: sb c1',
+             'c', '[SponsorBlock]: Preview/Recap',
+             'c', '[SponsorBlock]: Filler Tangent', 'c'])
+        self._remove_marked_arrange_sponsors_test_impl(chapters, expected, [])
+
     def test_remove_marked_arrange_sponsors_UniqueNamesForOverlappingSponsors(self):
     def test_remove_marked_arrange_sponsors_UniqueNamesForOverlappingSponsors(self):
         chapters = self._chapters([120], ['c']) + [
         chapters = self._chapters([120], ['c']) + [
             self._sponsor_chapter(10, 45, 'sponsor'), self._sponsor_chapter(20, 40, 'selfpromo'),
             self._sponsor_chapter(10, 45, 'sponsor'), self._sponsor_chapter(20, 40, 'selfpromo'),
@@ -173,7 +191,7 @@ class TestModifyChaptersPP(unittest.TestCase):
         self._remove_marked_arrange_sponsors_test_impl(chapters, expected, cuts)
         self._remove_marked_arrange_sponsors_test_impl(chapters, expected, cuts)
 
 
     def test_remove_marked_arrange_sponsors_ChapterWithCutHidingSponsor(self):
     def test_remove_marked_arrange_sponsors_ChapterWithCutHidingSponsor(self):
-        cuts = [self._sponsor_chapter(20, 50, 'selpromo', remove=True)]
+        cuts = [self._sponsor_chapter(20, 50, 'selfpromo', remove=True)]
         chapters = self._chapters([60], ['c']) + [
         chapters = self._chapters([60], ['c']) + [
             self._sponsor_chapter(10, 20, 'intro'),
             self._sponsor_chapter(10, 20, 'intro'),
             self._sponsor_chapter(30, 40, 'sponsor'),
             self._sponsor_chapter(30, 40, 'sponsor'),
@@ -199,7 +217,7 @@ class TestModifyChaptersPP(unittest.TestCase):
             self._sponsor_chapter(10, 20, 'sponsor'),
             self._sponsor_chapter(10, 20, 'sponsor'),
             self._sponsor_chapter(20, 30, 'interaction', remove=True),
             self._sponsor_chapter(20, 30, 'interaction', remove=True),
             self._chapter(30, 40, remove=True),
             self._chapter(30, 40, remove=True),
-            self._sponsor_chapter(40, 50, 'selpromo', remove=True),
+            self._sponsor_chapter(40, 50, 'selfpromo', remove=True),
             self._sponsor_chapter(50, 60, 'interaction')]
             self._sponsor_chapter(50, 60, 'interaction')]
         expected = self._chapters([10, 20, 30, 40],
         expected = self._chapters([10, 20, 30, 40],
                                   ['c', '[SponsorBlock]: Sponsor',
                                   ['c', '[SponsorBlock]: Sponsor',
@@ -282,7 +300,7 @@ class TestModifyChaptersPP(unittest.TestCase):
         chapters = self._chapters([70], ['c']) + [
         chapters = self._chapters([70], ['c']) + [
             self._sponsor_chapter(10, 30, 'sponsor'),
             self._sponsor_chapter(10, 30, 'sponsor'),
             self._sponsor_chapter(20, 50, 'interaction'),
             self._sponsor_chapter(20, 50, 'interaction'),
-            self._sponsor_chapter(30, 50, 'selpromo', remove=True),
+            self._sponsor_chapter(30, 50, 'selfpromo', remove=True),
             self._sponsor_chapter(40, 60, 'sponsor'),
             self._sponsor_chapter(40, 60, 'sponsor'),
             self._sponsor_chapter(50, 60, 'interaction')]
             self._sponsor_chapter(50, 60, 'interaction')]
         expected = self._chapters(
         expected = self._chapters(

+ 2 - 2
yt_dlp/options.py

@@ -1737,7 +1737,7 @@ def create_parser():
         '--sponsorblock-remove', metavar='CATS',
         '--sponsorblock-remove', metavar='CATS',
         dest='sponsorblock_remove', default=set(), action='callback', type='str',
         dest='sponsorblock_remove', default=set(), action='callback', type='str',
         callback=_set_from_options_callback, callback_kwargs={
         callback=_set_from_options_callback, callback_kwargs={
-            'allowed_values': set(SponsorBlockPP.CATEGORIES.keys()) - set(SponsorBlockPP.POI_CATEGORIES.keys()),
+            'allowed_values': set(SponsorBlockPP.CATEGORIES.keys()) - set(SponsorBlockPP.NON_SKIPPABLE_CATEGORIES.keys()),
             # Note: From https://wiki.sponsor.ajay.app/w/Types:
             # Note: From https://wiki.sponsor.ajay.app/w/Types:
             # The filler category is very aggressive.
             # The filler category is very aggressive.
             # It is strongly recommended to not use this in a client by default.
             # It is strongly recommended to not use this in a client by default.
@@ -1747,7 +1747,7 @@ def create_parser():
             'If a category is present in both mark and remove, remove takes precedence. '
             'If a category is present in both mark and remove, remove takes precedence. '
             'The syntax and available categories are the same as for --sponsorblock-mark '
             'The syntax and available categories are the same as for --sponsorblock-mark '
             'except that "default" refers to "all,-filler" '
             'except that "default" refers to "all,-filler" '
-            f'and {", ".join(SponsorBlockPP.POI_CATEGORIES.keys())} is not available'))
+            f'and {", ".join(SponsorBlockPP.NON_SKIPPABLE_CATEGORIES.keys())} are not available'))
     sponsorblock.add_option(
     sponsorblock.add_option(
         '--sponsorblock-chapter-title', metavar='TEMPLATE',
         '--sponsorblock-chapter-title', metavar='TEMPLATE',
         default=DEFAULT_SPONSORBLOCK_CHAPTER_TITLE, dest='sponsorblock_chapter_title',
         default=DEFAULT_SPONSORBLOCK_CHAPTER_TITLE, dest='sponsorblock_chapter_title',

+ 6 - 7
yt_dlp/postprocessor/modify_chapters.py

@@ -16,7 +16,7 @@ class ModifyChaptersPP(FFmpegPostProcessor):
                  *, sponsorblock_chapter_title=DEFAULT_SPONSORBLOCK_CHAPTER_TITLE, force_keyframes=False):
                  *, sponsorblock_chapter_title=DEFAULT_SPONSORBLOCK_CHAPTER_TITLE, force_keyframes=False):
         FFmpegPostProcessor.__init__(self, downloader)
         FFmpegPostProcessor.__init__(self, downloader)
         self._remove_chapters_patterns = set(remove_chapters_patterns or [])
         self._remove_chapters_patterns = set(remove_chapters_patterns or [])
-        self._remove_sponsor_segments = set(remove_sponsor_segments or []) - set(SponsorBlockPP.POI_CATEGORIES.keys())
+        self._remove_sponsor_segments = set(remove_sponsor_segments or []) - set(SponsorBlockPP.NON_SKIPPABLE_CATEGORIES.keys())
         self._ranges_to_remove = set(remove_ranges or [])
         self._ranges_to_remove = set(remove_ranges or [])
         self._sponsorblock_chapter_title = sponsorblock_chapter_title
         self._sponsorblock_chapter_title = sponsorblock_chapter_title
         self._force_keyframes = force_keyframes
         self._force_keyframes = force_keyframes
@@ -99,7 +99,7 @@ class ModifyChaptersPP(FFmpegPostProcessor):
             'start_time': start,
             'start_time': start,
             'end_time': end,
             'end_time': end,
             'category': 'manually_removed',
             'category': 'manually_removed',
-            '_categories': [('manually_removed', start, end)],
+            '_categories': [('manually_removed', start, end, 'Manually removed')],
             'remove': True,
             'remove': True,
         } for start, end in self._ranges_to_remove)
         } for start, end in self._ranges_to_remove)
 
 
@@ -290,13 +290,12 @@ class ModifyChaptersPP(FFmpegPostProcessor):
             c.pop('_was_cut', None)
             c.pop('_was_cut', None)
             cats = c.pop('_categories', None)
             cats = c.pop('_categories', None)
             if cats:
             if cats:
-                category = min(cats, key=lambda c: c[2] - c[1])[0]
-                cats = orderedSet(x[0] for x in cats)
+                category, _, _, category_name = min(cats, key=lambda c: c[2] - c[1])
                 c.update({
                 c.update({
                     'category': category,
                     'category': category,
-                    'categories': cats,
-                    'name': SponsorBlockPP.CATEGORIES[category],
-                    'category_names': [SponsorBlockPP.CATEGORIES[c] for c in cats]
+                    'categories': orderedSet(x[0] for x in cats),
+                    'name': category_name,
+                    'category_names': orderedSet(x[3] for x in cats),
                 })
                 })
                 c['title'] = self._downloader.evaluate_outtmpl(self._sponsorblock_chapter_title, c.copy())
                 c['title'] = self._downloader.evaluate_outtmpl(self._sponsorblock_chapter_title, c.copy())
                 # Merge identically named sponsors.
                 # Merge identically named sponsors.

+ 9 - 4
yt_dlp/postprocessor/sponsorblock.py

@@ -14,6 +14,10 @@ class SponsorBlockPP(FFmpegPostProcessor):
     POI_CATEGORIES = {
     POI_CATEGORIES = {
         'poi_highlight': 'Highlight',
         'poi_highlight': 'Highlight',
     }
     }
+    NON_SKIPPABLE_CATEGORIES = {
+        **POI_CATEGORIES,
+        'chapter': 'Chapter',
+    }
     CATEGORIES = {
     CATEGORIES = {
         'sponsor': 'Sponsor',
         'sponsor': 'Sponsor',
         'intro': 'Intermission/Intro Animation',
         'intro': 'Intermission/Intro Animation',
@@ -23,7 +27,7 @@ class SponsorBlockPP(FFmpegPostProcessor):
         'filler': 'Filler Tangent',
         'filler': 'Filler Tangent',
         'interaction': 'Interaction Reminder',
         'interaction': 'Interaction Reminder',
         'music_offtopic': 'Non-Music Section',
         'music_offtopic': 'Non-Music Section',
-        **POI_CATEGORIES,
+        **NON_SKIPPABLE_CATEGORIES
     }
     }
 
 
     def __init__(self, downloader, categories=None, api='https://sponsor.ajay.app'):
     def __init__(self, downloader, categories=None, api='https://sponsor.ajay.app'):
@@ -68,12 +72,13 @@ class SponsorBlockPP(FFmpegPostProcessor):
 
 
         def to_chapter(s):
         def to_chapter(s):
             (start, end), cat = s['segment'], s['category']
             (start, end), cat = s['segment'], s['category']
+            title = s['description'] if cat == 'chapter' else self.CATEGORIES[cat]
             return {
             return {
                 'start_time': start,
                 'start_time': start,
                 'end_time': end,
                 'end_time': end,
                 'category': cat,
                 'category': cat,
-                'title': self.CATEGORIES[cat],
-                '_categories': [(cat, start, end)]
+                'title': title,
+                '_categories': [(cat, start, end, title)],
             }
             }
 
 
         sponsor_chapters = [to_chapter(s) for s in duration_match]
         sponsor_chapters = [to_chapter(s) for s in duration_match]
@@ -89,7 +94,7 @@ class SponsorBlockPP(FFmpegPostProcessor):
         url = f'{self._API_URL}/api/skipSegments/{hash[:4]}?' + urllib.parse.urlencode({
         url = f'{self._API_URL}/api/skipSegments/{hash[:4]}?' + urllib.parse.urlencode({
             'service': service,
             'service': service,
             'categories': json.dumps(self._categories),
             'categories': json.dumps(self._categories),
-            'actionTypes': json.dumps(['skip', 'poi'])
+            'actionTypes': json.dumps(['skip', 'poi', 'chapter'])
         })
         })
         for d in self._download_json(url) or []:
         for d in self._download_json(url) or []:
             if d['videoID'] == video_id:
             if d['videoID'] == video_id: