Update readme.

README.md

@@ -1,5 +1,5 @@
 ---
-title: Docker Fastapi
+title: Ytdlp Subtitle
 emoji: 🐢
 colorFrom: purple
 colorTo: blue

fetchYoutubeSubtitle.py

@@ -1,97 +1,77 @@
 import json
-import os
-import requests
-import logging
+from typing import Optional
+import yt_dlp
 
-logger = logging.getLogger(__name__)
+def getVttUrlFromSubtitles(item, lang='en', vttType="vtt"):
+    langs = item.keys()
+    key =  lang if langs.get(lang) else ('en' if langs.get('en') else langs[0] )
+    for item in langs[key]:
+        if(item.get("ext") == type):
+            return item.get("url")
+    return None
 
-from typing import List, Dict
-
-SUBTITLE_DOWNLOADER_URL = 'https://savesubs.com'
-
-def fetchYoutubeSubtitleUrls(video_id):
-    headers = {
-        'accept': 'application/json, text/plain, */*',
-        'accept-language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7',
-        'cache-control': 'no-cache',
-        'Content-Type': 'application/json; charset=UTF-8',
-        'pragma': 'no-cache',
-        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-        'x-auth-token': os.environ.get('SAVESUBS_X_AUTH_TOKEN', ''),
-        'x-requested-domain': 'savesubs.com',
-        'X-requested-with': 'xmlhttprequest',
-        'sec-ch-ua': '"Google Chrome";v="111", "Not(A:Brand";v="8", "Chromium";v="111"',
-        'sec-ch-ua-mobile': '?0',
-        'sec-ch-ua-platform': 'Linux',
-        'sec-fetch-dest': 'empty',
-        'sec-fetch-mode': 'cors',
-        'sec-fetch-site': 'same-origin',
-        'authority': 'savesubs.com',
-        'origin': 'https://savesubs.com',
-        'referer': f'https://savesubs.com/process?url=https%3A%2F%2Fwww.youtube.com%2Fwatch%3Fv%3D{video_id}'
+def getSubtitle(url: str, lang: Optional[str] = 'en', vttType="vtt") -> Optional[str]:
+    ydl_opts = {
+        "writesubtitles": True,
+        "allsubtitles": True,
+        "subtitleslangs": [lang] if lang else [],
+        "skip_download": True,
     }
-    data = {
-        'data': {'url': f'https://www.youtube.com/watch?v={video_id}'}
-    }
-    session = requests.Session()
-
-    proxy = os.environ.get('PROXY', None)
-    if proxy:
-        session.proxies = {
-            "http": proxy,
-            "https": proxy,
-        }
-
-    response = session.post(SUBTITLE_DOWNLOADER_URL + '/action/extract', json=data, headers=headers)
-    
-    
-    if response.status_code != 200:
-        logger.error("response.status_code: {}".format(response.status_code))
-        return {'title': None, 'subtitleList': None, 'error': response.reason}
-    else:
-        try:
-            json = response.json().get('response', {})
-            logger.info('subtitle url json: {}'.format(json))
-            return {'title': json.get('title'), 'subtitleList': json.get('formats')}
-        except Exception as error:
-            logger.error(error)
-            return {'title': None, 'subtitleList': None, 'error': str(error)}
 
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info_dict = ydl.extract_info(url, download=False)
+        if info_dict.get("subtitles"):
+            # get first available subtitle
+            subtitle_url = getVttUrlFromSubtitles(info_dict.get("subtitles"), lang, vttType)
+            with ydl.urlopen(subtitle_url) as subtitle:
+                return subtitle.read().decode()
+        if info_dict.get("automatic_captions"):
+            subtitle_url = getVttUrlFromSubtitles(info_dict.get("automatic_captions"), , lang, vttType)
+            with ydl.urlopen(subtitle_url) as subtitle:
+                return subtitle.read().decode()
+    return None
 
-async def find(subtitleList: List[Dict], args: Dict) -> Dict:
-    key = list(args.keys())[0]
-    return next((item for item in subtitleList if item.get(key, None) == args[key]), None)
-
-async def fetchYoutubeSubtitle(videoId: str) -> Dict:
-    subtitle_url = ""
-    betterSubtitle = {}
-    subtitleList = []
-    title = ""
-    error = ""
-
-    result = fetchYoutubeSubtitleUrls(videoId)
-    title, subtitleList, error = result["title"], result["subtitleList"], result.get("error", None)
-
-    if not subtitleList or len(subtitleList) <= 0:
-        return {"title": title, "subtitle": None, "error": error}
-
-    betterSubtitle = (
-        await find(subtitleList, {"quality": "English"})
-        or await find(subtitleList, {"quality": "English (auto"})
-        or await find(subtitleList, {"quality": "zh-CN"})
-        or subtitleList[0]
-    )
+def fetchSubtitleUrls(url: str) -> json:
+    ydl_opts = {
+        "writesubtitles": True,
+        "allsubtitles": True,
+        "skip_download": True,
+    }
 
-    subtitleUrl = f"{SUBTITLE_DOWNLOADER_URL}{betterSubtitle['url']}?ext=srt"
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info_dict = ydl.extract_info(url, download=False)
+        if info_dict.get("subtitles"):
+            langs = info_dict.get("subtitles").keys()
+            if not (langs.length == 1 and "live_chat" in langs):
+                return info_dict.get("subtitles")
+            if info_dict.get("automatic_captions"):
+                return info_dict.get("automatic_captions")
+    return None
 
-    session = requests.Session()
+def get_subtitle(url, lang='en'):
+    if lang is None:
+        lang = 'en'
+    # Download subtitles if available
+    ydl_opts = {
+        'writesubtitles': True,
+        'outtmpl': '%(id)s.%(ext)s',
+        'subtitleslangs': [lang],
+        'skip_download': True,
+    }
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info_dict = ydl.extract_info(url, download=True)
+            video_id = info_dict.get("id", None)
+            if video_id is None:
+                return None
 
-    proxy = os.environ.get('PROXY', None)
-    if proxy:
-        session.proxies = {
-            "http": proxy,
-            "https": proxy,
-        }
-    response = session.get(subtitleUrl)
+            subtitle_file = f"{video_id}.{lang}.vtt"
+            with open(subtitle_file, 'r') as f:
+                subtitle_content = f.read()
+                subtitle_content = re.sub(r"<[^>]+>", "", subtitle_content)
+                return subtitle_content
+    except error:
+        print(error)
+        return None
     
-    return {"title": title, "subtitle": response.text}
+    return None

main.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
-from fetchYoutubeSubtitle import fetchYoutubeSubtitle
+from fetchYoutubeSubtitle import getSubtitle, fetchSubtitleUrls
 
 app = FastAPI()
 
@@ -15,6 +15,12 @@ def read_json():
 
 
 @app.get("/subtitle/")
-async def get_subtitle(vid: str):
-    subtitle = await fetchYoutubeSubtitle(vid)
-    return JSONResponse(content=subtitle)
+async def get_subtitle(url: str):
+    subtitle = await fetchSubtitleUrls(url)
+    return JSONResponse(content=subtitle)
+
+
+@app.get("/subtitle-urls/")
+async def get_subtitleUrls(url: str):
+    subtitles = await fetchSubtitleUrls(url)
+    return JSONResponse(content=subtitles)

requirements.txt

@@ -4,4 +4,4 @@ sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
 uvicorn[standard]==0.17.*
-langchain
+yt-dlp