ffmpeg_subtitle_apply 新增内嵌字幕选项

* ffmpeg_subtitle_apply 新增内嵌字幕选项

---------

Merge request URL: https://g-ldyi2063.coding.net/p/dev/d/modalDeploy/git/merge/4860?initial=true
Co-authored-by: shuohigh@gmail.com

src/BowongModalFunctions/models/web_model.py

@@ -4,7 +4,8 @@ from enum import Enum
 from typing import List, Union, Optional, Dict, Any
 
 import pydantic
-from pydantic import BaseModel, Field, field_validator, ConfigDict, HttpUrl, computed_field, RootModel
+from pydantic import BaseModel, Field, field_validator, ConfigDict, HttpUrl, computed_field, RootModel, root_validator, \
+    model_validator
 
 from .ffmpeg_worker_model import FFMpegSliceSegment
 from .media_model import MediaSource, MediaSources, MediaProtocol
@@ -284,12 +285,13 @@ class FFMPEGOverlayGifTaskStatusResponse(BaseFFMPEGTaskStatusResponse):
 
 class FFMPEGSubtitleOverlayRequest(BaseFFMPEGTaskRequest):
     media: MediaSource = Field(description="需要处理的媒体源")
-    subtitle: MediaSource = Field(description="需要叠加的字幕文件")
+    subtitle: Optional[MediaSource] = Field(description="需要叠加的字幕文件")
-    fonts: List[MediaSource] = Field(description="字幕文件内使用到的字体文件")
+    embedded_subtitle: Optional[MediaSource] = Field(description="需要内嵌的字幕文件")
+    fonts: Optional[List[MediaSource]] = Field(description="字幕文件内使用到的字体文件")
 
     @field_validator('media', mode='before')
     @classmethod
-    def parse_media_input(cls, v: Union[str, MediaSource]):
+    def parse_media_input(cls, v: Union[str, MediaSource]) -> MediaSource:
         if isinstance(v, str):
             return MediaSource.from_str(v)
         elif isinstance(v, MediaSource):
@@ -299,7 +301,9 @@ class FFMPEGSubtitleOverlayRequest(BaseFFMPEGTaskRequest):
 
     @field_validator('subtitle', mode='before')
     @classmethod
-    def parse_subtitle_input(cls, v: Union[str, MediaSource]):
+    def parse_subtitle_input(cls, v: Union[str, MediaSource]) -> Optional[MediaSource]:
+        if not v:
+            return None
         if isinstance(v, str):
             media = MediaSource.from_str(v)
             if not media.file_extension in ['ass']:
@@ -310,11 +314,26 @@ class FFMPEGSubtitleOverlayRequest(BaseFFMPEGTaskRequest):
         else:
             raise TypeError(v)
 
+    @field_validator('embedded_subtitle', mode='before')
+    @classmethod
+    def parse_embedded_subtitle_input(cls, v: Union[str, MediaSource]) -> Optional[MediaSource]:
+        if not v:
+            return None
+        if isinstance(v, str):
+            media = MediaSource.from_str(v)
+            if not media.file_extension in ['vtt', 'srt', 'ass']:
+                raise pydantic.ValidationError("必须使用标准字幕文件, 如.vtt/.srt/.ass")
+            return media
+        elif isinstance(v, MediaSource):
+            return v
+        else:
+            raise TypeError(v)
+
     @field_validator('fonts', mode='before')
     @classmethod
-    def parse_font_input(cls, v: Union[str, MediaSource]) -> List[MediaSource]:
+    def parse_font_input(cls, v: Union[str, MediaSource]) -> Optional[List[MediaSource]]:
         if not v:
-            raise pydantic.ValidationError("fonts输入为空")
+            return None
         result = []
         for item in v:
             if isinstance(item, str):
@@ -325,6 +344,13 @@ class FFMPEGSubtitleOverlayRequest(BaseFFMPEGTaskRequest):
                 raise pydantic.ValidationError("fonts元素类型错误: 必须是字符串")
         return result
 
+    @model_validator(mode='after')
+    def check_at_least_one(self):
+        if self.subtitle is None and self.embedded_subtitle is None:
+            raise pydantic.ValidationError("至少需要提供一个有效字幕")
+        if self.subtitle is not None and self.fonts is None:
+            raise pydantic.ValidationError("使用叠加字幕时需要指定使用的字体文件")
+
 
 class FFMPEGSubtitleTaskStatusResponse(BaseFFMPEGTaskStatusResponse):
     result: Optional[str] = Field(default=None, description="生成结果的URN")
@@ -560,10 +586,11 @@ class GoogleUploadResponse(BaseModel):
     @computed_field(description="适用于Google内部服务的URN")
     @property
     def urn(self) -> str:
-        return ("gs://" + self.id).replace(f"/{self.generation}","")
+        return ("gs://" + self.id).replace(f"/{self.generation}", "")
 
     model_config = ConfigDict(populate_by_name=True)
 
+
 class VisualFeatures(BaseModel):
     """
     代表产品的视觉特征
@@ -579,6 +606,7 @@ class VisualFeatures(BaseModel):
         description="商品详细款式风格设计等有辨识度的款式特征"
     )
 
+
 class ProductInfo(BaseModel):
     """
     表示从图像中提取单个产品的信息
@@ -601,12 +629,14 @@ class ProductInfo(BaseModel):
         description="产品的详细视觉特征"
     )
 
+
 class GeminiFirstStageResponseModel(RootModel):
     """
     从图像中提取的产品信息列表
     """
     root: List[ProductInfo]
 
+
 class ProductTimeline(BaseModel):
     """
     表示单个商品及其在视频时间线中的出现信息。
@@ -616,8 +646,9 @@ class ProductTimeline(BaseModel):
         description="该商品在视频中出现的时间段列表。每个字符串表示一个时间段，格式为 '开始时间 - 结束时间 (描述)'"
     )
 
+
 class GeminiSecondStageResponseModel(RootModel):
     """
     表示一个包含多个商品及其时间线信息的列表。
     """
-    root: List[ProductTimeline]
+    root: List[ProductTimeline]

src/BowongModalFunctions/router/ffmpeg.py

@@ -161,8 +161,8 @@ async def subtitle_apply(body: FFMPEGSubtitleOverlayRequest,
     sentry_trace = None
     if headers.x_trace_id and headers.x_baggage:
         sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
-    fn_call = fn.spawn(media=body.media, subtitle=body.subtitle, fonts=body.fonts, sentry_trace=sentry_trace,
+    fn_call = fn.spawn(media=body.media, subtitle=body.subtitle, embed_subtitle=body.embedded_subtitle,
-                       webhook=body.webhook)
+                       fonts=body.fonts, sentry_trace=sentry_trace, webhook=body.webhook)
     return ModalTaskResponse(success=True, taskId=fn_call.object_id)
 
 
@@ -199,7 +199,8 @@ async def video_extract_frame(body: FFMPEGExtractFrameRequest,
     sentry_trace = None
     if headers.x_trace_id and headers.x_baggage:
         sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
-    fn_call = fn.spawn(media=body.video,seek_time=body.prased_seek_time, frame_index=body.frame_index, sentry_trace=sentry_trace, webhook=body.webhook)
+    fn_call = fn.spawn(media=body.video, seek_time=body.prased_seek_time, frame_index=body.frame_index,
+                       sentry_trace=sentry_trace, webhook=body.webhook)
     return ModalTaskResponse(success=True, taskId=fn_call.object_id)
 
 

src/BowongModalFunctions/utils/VideoUtils.py

@@ -143,6 +143,19 @@ class ImageStream(MediaStream):
     model_config = ConfigDict(extra='allow')
 
 
+class SubtitleStreamTags(BaseModel):
+    language: str = Field(description="内嵌字幕的语言")
+
+    model_config = ConfigDict(extra='allow')
+
+
+class SubtitleStream(MediaStream):
+    stream_type: str = "subtitle"
+    tags: SubtitleStreamTags
+
+    model_config = ConfigDict(extra='allow')
+
+
 class VideoFormat(BaseModel):
     filename: str
     format_name: str
@@ -153,7 +166,8 @@ class VideoFormat(BaseModel):
 
 
 class VideoMetadata(BaseModel):
-    streams: List[Union[ImageStream, AudioStream, VideoStream, HLSMediaAudioStream, HLSMediaVideoStream]] = Field(
+    streams: List[
+        Union[ImageStream, AudioStream, VideoStream, HLSMediaAudioStream, HLSMediaVideoStream, SubtitleStream]] = Field(
         description="媒体包含的数据轨道")
     format: Optional[VideoFormat] = Field(None)
 
@@ -187,6 +201,10 @@ class VideoMetadata(BaseModel):
                                 logger.info("Parsing video stream")
                                 video = VideoStream.model_validate(stream)
                                 streams.append(video)
+                    elif stream.get("codec_type") == 'subtitle':
+                        logger.info("Parsing subtitle stream")
+                        subtitle_stream = SubtitleStream.model_validate(stream)
+                        streams.append(subtitle_stream)
             return streams
         else:
             raise TypeError
@@ -558,7 +576,7 @@ class VideoUtils:
     @staticmethod
     async def convert_m3u8_to_local_source(media_stream_url: str,
                                            head: float = 0,
-                                           tail: float = 86400, # 使用24H时长替代♾️
+                                           tail: float = 86400,  # 使用24H时长替代♾️
                                            temp_dir: str = None) -> tuple[str, str, TimeDelta]:
         """
         转换m3u8为本地来源
@@ -628,7 +646,8 @@ class VideoUtils:
         os.makedirs(os.path.dirname(output_path), exist_ok=True)
 
         try:
-            local_m3u8_path, temp_dir, diff = await VideoUtils.convert_m3u8_to_local_source(media_stream_url=media_stream_url)
+            local_m3u8_path, temp_dir, diff = await VideoUtils.convert_m3u8_to_local_source(
+                media_stream_url=media_stream_url)
             # 使用ffmpeg合并TS片段
             ffmpeg_cmd = VideoUtils.async_ffmpeg_init()
             ffmpeg_cmd.input(local_m3u8_path,
@@ -974,12 +993,14 @@ class VideoUtils:
         return output_path, video_metadata
 
     @staticmethod
-    async def ffmpeg_subtitle_apply(media_path: str, subtitle_path: str,
+    async def ffmpeg_subtitle_apply(media_path: str, subtitle_path: Optional[str], embed_subtitle_path: Optional[str],
-                                    font_dir: str, output_path: Optional[str] = None) -> Tuple[str, VideoMetadata]:
+                                    font_dir: Optional[str], output_path: Optional[str] = None) -> Tuple[
+        str, VideoMetadata]:
         """
         给视频画面叠加字幕，需要确保字幕文件为ass字幕，并且subtitle文件内设置的字体存在与font_dir文件夹内
         :param media_path: 待处理的源视频
-        :param subtitle_path: ass字幕文件路径
+        :param subtitle_path: ass渲染字幕文件路径
+        :param embed_subtitle_path: ass/vtt/srt内嵌字幕文件路径
         :param font_dir: 字体文件目录路径
         :param output_path: 指定输出文件路径（可选）
         :return: 返回最终输出视频路径， 最终输出视频时长
@@ -990,14 +1011,23 @@ class VideoUtils:
         video_metadata = VideoUtils.ffprobe_video_format(media_path)
         ffmpeg_cmd = VideoUtils.async_ffmpeg_init()
         ffmpeg_cmd.input(media_path)
-        ffmpeg_cmd.output(output_path,
+        if embed_subtitle_path:
-                          vf=f"subtitles=filename={subtitle_path}:fontsdir={font_dir}",
+            ffmpeg_cmd.input(embed_subtitle_path)
-                          vcodec="libx264",
+        ffmpeg_options = {
-                          crf=16,
+            "vcodec": "libx264",
-                          b=video_metadata.video_bitrate,  # 视频码率
+            "acodec": "copy",
-                          r=video_metadata.video_frame_rate,  # 帧率
+            "crf": 16,
-                          acodec="copy",
+            "b": video_metadata.video_bitrate,
-                          )
+            "r": video_metadata.video_frame_rate,
+            "map": ["0:v", "0:a"]
+        }
+        if subtitle_path and font_dir:
+            ffmpeg_options["vf"] = f"subtitles=filename={subtitle_path}:fontsdir={font_dir}"
+        if embed_subtitle_path:
+            ffmpeg_options['map'].append("1")
+            ffmpeg_options['c:s'] = "mov_text"
+            ffmpeg_options['metadata:s:s:0'] = "language=chi"
+        ffmpeg_cmd.output(output_path, options=ffmpeg_options)
         await ffmpeg_cmd.execute()
         video_metadata = VideoUtils.ffprobe_media_metadata(output_path)
         return output_path, video_metadata

src/cluster/ffmpeg_apps/subtitle_apply.py

@@ -25,8 +25,10 @@ with ffmpeg_worker_image.imports():
                       ),
                   }, )
     @modal.concurrent(max_inputs=1)
-    async def ffmpeg_subtitle_apply(media: MediaSource, subtitle: MediaSource,
+    async def ffmpeg_subtitle_apply(media: MediaSource, subtitle: Optional[MediaSource],
-                                    fonts: List[MediaSource], sentry_trace: Optional[SentryTransactionInfo] = None,
+                                    embed_subtitle: Optional[MediaSource],
+                                    fonts: Optional[List[MediaSource]],
+                                    sentry_trace: Optional[SentryTransactionInfo] = None,
                                     webhook: Optional[WebhookNotify] = None) -> Tuple[
         FFMPEGResult, Optional[SentryTransactionInfo]]:
         fn_id = current_function_call_id()
@@ -35,25 +37,32 @@ with ffmpeg_worker_image.imports():
                                     sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None,
                                     sentry_baggage=sentry_trace.x_baggage if sentry_trace else None)
         @SentryUtils.webhook_handler(webhook=webhook, func_id=fn_id)
-        async def ffmpeg_process(video: MediaSource, subtitle: MediaSource,
+        async def ffmpeg_process(video: MediaSource, subtitle: Optional[MediaSource],
-                                 fonts: List[MediaSource], func_id: str) -> FFMPEGResult:
+                                 embed_subtitle: Optional[MediaSource], func_id: str,
+                                 fonts: Optional[List[MediaSource]], ) -> FFMPEGResult:
             media_path = f"{s3_mount}/{video.cache_filepath}"
-            subtitle_path = f"{s3_mount}/{subtitle.cache_filepath}"
+            subtitle_path = f"{s3_mount}/{subtitle.cache_filepath}" if subtitle else None
+            embed_subtitle_path = f"{s3_mount}/{embed_subtitle.cache_filepath}" if embed_subtitle else None
             output_path = f"{output_path_prefix}/{config.modal_environment}/subtitle_apply/{func_id}/output.mp4"
-            font_dir = f"{output_path_prefix}/{config.modal_environment}/subtitle_apply/{func_id}/fonts"
+            if fonts:
-            os.makedirs(font_dir, exist_ok=True)
+                font_dir = f"{output_path_prefix}/{config.modal_environment}/subtitle_apply/{func_id}/fonts"
-            local_fonts = [f"{s3_mount}/{font.cache_filepath}" for font in fonts]
+                os.makedirs(font_dir, exist_ok=True)
-            for font in local_fonts:
+                local_fonts = [f"{s3_mount}/{font.cache_filepath}" for font in fonts]
-                font_filename = os.path.basename(font)
+                for font in local_fonts:
-                shutil.copy(font, f"{font_dir}/{font_filename}")
+                    font_filename = os.path.basename(font)
+                    shutil.copy(font, f"{font_dir}/{font_filename}")
+            else:
+                font_dir = None
             local_output, metadata = await VideoUtils.ffmpeg_subtitle_apply(media_path=media_path,
+                                                                            embed_subtitle_path=embed_subtitle_path,
                                                                             subtitle_path=subtitle_path,
                                                                             font_dir=font_dir, output_path=output_path)
             s3_outputs = local_copy_to_s3([local_output])
             return FFMPEGResult(urn=s3_outputs[0], metadata=metadata,
                                 content_length=FileUtils.get_file_size(local_output), )
 
-        result = await ffmpeg_process(video=media, subtitle=subtitle, fonts=fonts, func_id=fn_id)
+        result = await ffmpeg_process(video=media, subtitle=subtitle, embed_subtitle=embed_subtitle, fonts=fonts,
+                                      func_id=fn_id)
         if not sentry_trace:
             sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
                                                  x_baggage=sentry_sdk.get_baggage())

tests/ffmpeg_test_case.py

@@ -257,5 +257,14 @@ class FFMPEGTestCase(unittest.IsolatedAsyncioTestCase):
         ffmpeg_cmd.output("./videos/output_large_h264.mp4", options=h264_options)
         await ffmpeg_cmd.execute()
 
+    async def test_ffmpeg_subtitle_apply(self):
+        video = "./videos/subtitle_test/A美洋MEIYANG 方圆牛仔马甲 率性哲学！八支精棉牛仔无袖上衣-周二.mp4"
+        subtitle = "./videos/subtitle_test/A美洋MEIYANG 方圆牛仔马甲 率性哲学！八支精棉牛仔无袖上衣-周二.vtt"
+        output = "./videos/subtitle_test/output.mp4"
+        output_video, metadata = await VideoUtils.ffmpeg_subtitle_apply(media_path=video, subtitle_path=None,
+                                                                        embed_subtitle_path=subtitle, font_dir=None,
+                                                                        output_path=output)
+        self.assertEqual(len(metadata.streams), 3)
+
     if __name__ == '__main__':
         unittest.main()