Create app.py

app.py

@@ -0,0 +1,207 @@
+import gradio as gr
+import pytesseract
+import cv2
+import multiprocessing
+from fuzzywuzzy import fuzz
+from dataclasses import dataclass
+from urllib.request import urlopen
+import shutil
+import pathlib
+import datetime
+import sys
+
+# Constants
+TESSDATA_DIR = pathlib.Path.home() / 'tessdata'
+TESSDATA_URL = 'https://github.com/tesseract-ocr/tessdata_fast/raw/master/{}.traineddata'
+TESSDATA_SCRIPT_URL = 'https://github.com/tesseract-ocr/tessdata_best/raw/master/script/{}.traineddata'
+
+
+# Download language data files if necessary
+def download_lang_data(lang: str):
+    TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
+    for lang_name in lang.split('+'):
+        filepath = TESSDATA_DIR / f'{lang_name}.traineddata'
+        if not filepath.is_file():
+            url = TESSDATA_SCRIPT_URL.format(lang_name) if lang_name[0].isupper() else TESSDATA_URL.format(lang_name)
+            with urlopen(url) as res, open(filepath, 'w+b') as f:
+                shutil.copyfileobj(res, f)
+
+
+# Helper functions for time and frame conversion
+def get_frame_index(time_str: str, fps: float):
+    t = list(map(float, time_str.split(':')))
+    if len(t) == 3:
+        td = datetime.timedelta(hours=t[0], minutes=t[1], seconds=t[2])
+    elif len(t) == 2:
+        td = datetime.timedelta(minutes=t[0], seconds=t[1])
+    else:
+        raise ValueError(f'Time data "{time_str}" does not match format "%H:%M:%S"')
+    return int(td.total_seconds() * fps)
+
+
+def get_srt_timestamp(frame_index: int, fps: float):
+    td = datetime.timedelta(seconds=frame_index / fps)
+    ms = td.microseconds // 1000
+    m, s = divmod(td.seconds, 60)
+    h, m = divmod(m, 60)
+    return f'{h:02d}:{m:02d}:{s:02d},{ms:03d}'
+
+
+# Video capture class using OpenCV
+class Capture:
+    def __init__(self, video_path):
+        self.path = video_path
+
+    def __enter__(self):
+        self.cap = cv2.VideoCapture(self.path)
+        if not self.cap.isOpened():
+            raise IOError(f'Cannot open video {self.path}.')
+        return self.cap
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.cap.release()
+
+
+@dataclass
+class PredictedWord:
+    confidence: int
+    text: str
+
+
+class PredictedFrame:
+    def __init__(self, index: int, pred_data: str, conf_threshold: int):
+        self.index = index
+        self.words = []
+        block = 0
+        for l in pred_data.splitlines()[1:]:
+            word_data = l.split()
+            if len(word_data) < 12:
+                continue
+            _, _, block_num, *_, conf, text = word_data
+            block_num, conf = int(block_num), int(conf)
+            if block < block_num:
+                block = block_num
+                if self.words and self.words[-1].text != '\n':
+                    self.words.append(PredictedWord(0, '\n'))
+            if conf >= conf_threshold:
+                self.words.append(PredictedWord(conf, text))
+        self.confidence = sum(word.confidence for word in self.words)
+        self.text = ' '.join(word.text for word in self.words).translate(str.maketrans('|', 'I', '<>{}[];`@#$%^*_=~\\')).replace(' \n ', '\n').strip()
+
+    def is_similar_to(self, other, threshold=70):
+        return fuzz.ratio(self.text, other.text) >= threshold
+
+
+class PredictedSubtitle:
+    def __init__(self, frames, sim_threshold):
+        self.frames = [f for f in frames if f.confidence > 0]
+        self.sim_threshold = sim_threshold
+        self.text = max(self.frames, key=lambda f: f.confidence).text if self.frames else ''
+
+    @property
+    def index_start(self):
+        return self.frames[0].index if self.frames else 0
+
+    @property
+    def index_end(self):
+        return self.frames[-1].index if self.frames else 0
+
+    def is_similar_to(self, other):
+        return fuzz.partial_ratio(self.text, other.text) >= self.sim_threshold
+
+
+class Video:
+    def __init__(self, path):
+        self.path = path
+        with Capture(path) as v:
+            self.num_frames = int(v.get(cv2.CAP_PROP_FRAME_COUNT))
+            self.fps = v.get(cv2.CAP_PROP_FPS)
+            self.height = int(v.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    def run_ocr(self, lang, time_start, time_end, conf_threshold, use_fullframe):
+        self.lang = lang
+        self.use_fullframe = use_fullframe
+        ocr_start = get_frame_index(time_start, self.fps) if time_start else 0
+        ocr_end = get_frame_index(time_end, self.fps) if time_end else self.num_frames
+        if ocr_end < ocr_start:
+            raise ValueError('time_start is later than time_end')
+
+        num_ocr_frames = ocr_end - ocr_start
+        with Capture(self.path) as v, multiprocessing.Pool() as pool:
+            v.set(cv2.CAP_PROP_POS_FRAMES, ocr_start)
+            frames = (v.read()[1] for _ in range(num_ocr_frames))
+            it_ocr = pool.imap(self._image_to_data, frames, chunksize=10)
+            self.pred_frames = [PredictedFrame(i + ocr_start, data, conf_threshold) for i, data in enumerate(it_ocr)]
+
+    def _image_to_data(self, img):
+        if not self.use_fullframe:
+            img = img[self.height // 2:, :]
+        config = f'--tessdata-dir "{TESSDATA_DIR}"'
+        try:
+            return pytesseract.image_to_data(img, lang=self.lang, config=config)
+        except Exception as e:
+            sys.exit(f'{e.__class__.__name__}: {e}')
+
+    def get_subtitles(self, sim_threshold):
+        self._generate_subtitles(sim_threshold)
+        return ''.join(f'{i}\n{get_srt_timestamp(sub.index_start, self.fps)} --> {get_srt_timestamp(sub.index_end, self.fps)}\n{sub.text}\n\n' for i, sub in enumerate(self.pred_subs))
+
+    def _generate_subtitles(self, sim_threshold):
+        self.pred_subs = []
+        if self.pred_frames is None:
+            raise AttributeError('Please call self.run_ocr() first to perform OCR on frames')
+
+        WIN_BOUND = int(self.fps // 2)
+        bound = WIN_BOUND
+        i = 0
+        j = 1
+        while j < len(self.pred_frames):
+            fi, fj = self.pred_frames[i], self.pred_frames[j]
+            if fi.is_similar_to(fj):
+                bound = WIN_BOUND
+            elif bound > 0:
+                bound -= 1
+            else:
+                para_new = j - WIN_BOUND
+                self._append_sub(PredictedSubtitle(self.pred_frames[i:para_new], sim_threshold))
+                i = para_new
+                j = i
+                bound = WIN_BOUND
+            j += 1
+        if i < len(self.pred_frames) - 1:
+            self._append_sub(PredictedSubtitle(self.pred_frames[i:], sim_threshold))
+
+    def _append_sub(self, sub):
+        if not sub.text:
+            return
+        while self.pred_subs and sub.is_similar_to(self.pred_subs[-1]):
+            ls = self.pred_subs.pop()
+            sub = PredictedSubtitle(ls.frames + sub.frames, sub.sim_threshold)
+        self.pred_subs.append(sub)
+
+
+# Gradio app
+def extract_subtitles(video_file, lang, time_start, time_end, conf_threshold, use_fullframe, sim_threshold):
+    video = Video(video_file.name)
+    video.run_ocr(lang, time_start, time_end, conf_threshold, use_fullframe)
+    subtitles = video.get_subtitles(sim_threshold)
+    return subtitles
+
+
+iface = gr.Interface(
+    fn=extract_subtitles,
+    inputs=[
+        gr.Video(label="Video File"),
+        gr.Textbox(value='eng', label="OCR Language"),
+        gr.Textbox(value='00:00:00', label="Start Time (HH:MM:SS)"),
+        gr.Textbox(value='', label="End Time (HH:MM:SS, leave empty for full video)"),
+        gr.Slider(0, 100, value=60, step=1, label="Confidence Threshold"),
+        gr.Checkbox(label="Use Full Frame for OCR", default=False),
+        gr.Slider(0, 100, value=70, step=1, label="Similarity Threshold")
+    ],
+    outputs=gr.Textbox(label="Extracted Subtitles"),
+    title="Video Subtitle Extractor",
+    description="Extract hardcoded subtitles from videos using machine learning.",
+)
+
+iface.launch()