Delete app.py

app.py

@@ -1,207 +0,0 @@
-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()