app.py

"""
Video Hard-Subtitle OCR Application
Extract text from video hard subtitles with customizable region selection
"""

import gradio as gr
import cv2
import numpy as np
import json
import os
import tempfile
from datetime import datetime
from pathlib import Path

# Try to import optional dependencies
try:
    import moviepy.editor as mp
    MOVIEPY_AVAILABLE = True
except ImportError:
    MOVIEPY_AVAILABLE = False

try:
    import easyocr
    EASYOCR_AVAILABLE = True
except ImportError:
    EASYOCR_AVAILABLE = False


# Language to OCR language code mapping
LANGUAGE_CODES = {
    "English": "en",
    "Japanese": "ja",
    "Korean": "ko",
    "Chinese (Simplified)": "ch_sim",
    "Chinese (Traditional)": "ch_tra",
    "Thai": "th"
}

# OCR Reader cache
_reader_cache = {}


def get_ocr_reader(language: str):
    """Get or create OCR reader for specified language"""
    lang_code = LANGUAGE_CODES.get(language, "en")
    
    if lang_code in _reader_cache:
        return _reader_cache[lang_code]
    
    if EASYOCR_AVAILABLE:
        try:
            reader = easyocr.Reader([lang_code], gpu=False, verbose=False)
            _reader_cache[lang_code] = reader
            return reader
        except Exception as e:
            return None
    return None


def extract_frames_from_video(video_path: str, start_time: float, end_time: float):
    """Extract frames from video in the specified time range"""
    frames = []
    timestamps = []
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return frames, timestamps
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps > 0 else 0
    
    start_frame = int(start_time * fps)
    end_frame = int(end_time * fps)
    
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    
    current_frame = start_frame
    while current_frame <= end_frame:
        ret, frame = cap.read()
        if not ret:
            break
        
        timestamp = current_frame / fps
        frames.append(frame)
        timestamps.append(timestamp)
        current_frame += 1
    
    cap.release()
    return frames, timestamps


def extract_subtitle_from_frame(frame, x: int, y: int, width: int, height: int, reader):
    """Extract text from a specific region of a frame"""
    h, w = frame.shape[:2]
    
    # Ensure coordinates are within bounds
    x = max(0, min(x, w - 1))
    y = max(0, min(y, h - 1))
    width = min(width, w - x)
    height = min(height, h - y)
    
    if width <= 0 or height <= 0:
        return ""
    
    # Crop the subtitle region
    roi = frame[y:y+height, x:x+width]
    
    # Convert to RGB for OCR
    roi_rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
    
    try:
        results = reader.readtext(roi_rgb)
        if results:
            # Get the text with highest confidence
            text = " ".join([result[1] for result in results])
            return text.strip()
    except Exception:
        pass
    
    return ""


def merge_subtitles(subtitles: list, time_threshold: float = 1.0, iou_threshold: float = 0.5):
    """Merge consecutive subtitles that are similar or close in time"""
    if not subtitles:
        return []
    
    merged = []
    current = subtitles[0].copy()
    
    for i in range(1, len(subtitles)):
        next_sub = subtitles[i]
        
        # Check if they overlap in time
        time_gap = next_sub["start"] - current["end"]
        
        # Check if text is similar (simple check)
        text_similar = (
            next_sub["text"].strip().lower() == current["text"].strip().lower()
            or next_sub["text"].strip() in current["text"].strip()
        )
        
        if time_gap <= time_threshold and text_similar:
            # Extend current subtitle
            current["end"] = next_sub["end"]
        else:
            merged.append(current)
            current = next_sub.copy()
    
    merged.append(current)
    return merged


def process_video_ocr(
    video_path: str,
    x: int, y: int, width: int, height: int,
    language: str,
    start_time: float, end_time: float,
    progress=gr.Progress()
):
    """Main OCR processing function"""
    
    if not video_path:
        return None, "Please upload a video first.", ""
    
    if not EASYOCR_AVAILABLE:
        return None, "EasyOCR is not installed. Please install it with: pip install easyocr", ""
    
    reader = get_ocr_reader(language)
    if reader is None:
        return None, f"Could not initialize OCR for {language}", ""
    
    # Extract frames
    progress(0.1, "Extracting frames from video...")
    frames, timestamps = extract_frames_from_video(video_path, start_time, end_time)
    
    if not frames:
        return None, "Could not extract frames from video", ""
    
    # Process frames
    subtitles = []
    frame_interval = max(1, len(frames) // 100)  # Sample ~100 points
    
    for i, (frame, timestamp) in enumerate(frames):
        if i % frame_interval == 0:
            progress(0.1 + 0.7 * (i / len(frames)), f"Processing frame {i}/{len(frames)}...")
        
        text = extract_subtitle_from_frame(frame, x, y, width, height, reader)
        
        if text and text.strip():
            # Check if this is a new subtitle or continuation
            if subtitles and (timestamp - subtitles[-1]["end"]) < 0.5:
                # Might be continuation, check text similarity
                if text.strip() != subtitles[-1]["text"].strip():
                    subtitles.append({
                        "start": timestamp,
                        "end": timestamp + 0.1,
                        "text": text.strip()
                    })
            else:
                subtitles.append({
                    "start": timestamp,
                    "end": timestamp + 0.1,
                    "text": text.strip()
                })
    
    # Merge similar subtitles
    progress(0.85, "Merging subtitle entries...")
    subtitles = merge_subtitles(subtitles)
    
    # Assign sequential IDs
    for i, sub in enumerate(subtitles):
        sub["id"] = i + 1
    
    # Generate SRT
    progress(0.95, "Generating output files...")
    srt_output = generate_srt(subtitles)
    txt_output = generate_txt(subtitles)
    json_output = generate_json(subtitles, language, start_time, end_time)
    
    progress(1.0, "Complete!")
    
    return subtitles, srt_output, json_output


def generate_srt(subtitles: list) -> str:
    """Generate SRT format output"""
    srt_lines = []
    
    for sub in subtitles:
        sub_id = sub.get("id", 1)
        start_time = format_srt_time(sub["start"])
        end_time = format_srt_time(sub["end"])
        text = sub["text"]
        
        srt_lines.append(f"{sub_id}")
        srt_lines.append(f"{start_time} --> {end_time}")
        srt_lines.append(text)
        srt_lines.append("")
    
    return "\n".join(srt_lines)


def format_srt_time(seconds: float) -> str:
    """Format time for SRT (HH:MM:SS,mmm)"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


def generate_txt(subtitles: list) -> str:
    """Generate plain text output"""
    return "\n".join([sub["text"] for sub in subtitles])


def generate_json(subtitles: list, language: str, start: float, end: float) -> str:
    """Generate JSON output"""
    output = {
        "metadata": {
            "language": language,
            "start_time": start,
            "end_time": end,
            "total_subtitles": len(subtitles),
            "generated_at": datetime.now().isoformat()
        },
        "subtitles": subtitles
    }
    return json.dumps(output, indent=2, ensure_ascii=False)


def get_video_info(video_path: str):
    """Get video information"""
    if not video_path or not os.path.exists(video_path):
        return 0, 0, 0, 0
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return 0, 0, 0, 0
    
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / fps if fps > 0 else 0
    
    cap.release()
    
    return width, height, fps, duration


def create_demo():
    """Create the Gradio demo"""
    
    with gr.Blocks(title="Video Subtitle OCR") as demo:
        
        # Header with anycoder link
        gr.HTML("""
        <div style="text-align: center; padding: 10px; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); border-radius: 10px; margin-bottom: 20px;">
            <h1 style="color: white; margin: 0;">🎬 Video Subtitle OCR</h1>
            <p style="color: #e0e0e0; margin: 5px 0;">Extract hard-coded subtitles from any video</p>
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #ffd700; font-weight: bold;">Built with anycoder</a>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                # Video upload and preview
                video_input = gr.Video(
                    label="Upload Video",
                    sources=["upload"],
                    height=400
                )
                
                # Video info display
                video_info = gr.JSON(label="Video Information", visible=False)
                
            with gr.Column(scale=1):
                # Settings panel
                gr.Markdown("### ⚙️ Settings")
                
                # Language selection
                language = gr.Dropdown(
                    choices=list(LANGUAGE_CODES.keys()),
                    value="English",
                    label="OCR Language",
                    info="Select the primary language for OCR accuracy"
                )
                
                # Time range
                with gr.Row():
                    start_time = gr.Number(
                        label="Start Time (sec)",
                        value=0,
                        minimum=0,
                        step=0.1
                    )
                    end_time = gr.Number(
                        label="End Time (sec)",
                        value=10,
                        minimum=0,
                        step=0.1
                    )
        
        # Subtitle region adjustment
        gr.Markdown("### 📐 Adjust Subtitle Area")
        
        with gr.Row():
            with gr.Column():
                # Position sliders
                x_pos = gr.Slider(
                    minimum=0, maximum=100, value=10,
                    label="Horizontal Position (X %)",
                    info="X-axis position of subtitle box"
                )
                y_pos = gr.Slider(
                    minimum=0, maximum=100, value=80,
                    label="Vertical Position (Y %)",
                    info="Y-axis position of subtitle box"
                )
            
            with gr.Column():
                # Size sliders
                box_width = gr.Slider(
                    minimum=10, maximum=100, value=80,
                    label="Width (%)",
                    info="Width of subtitle box"
                )
                box_height = gr.Slider(
                    minimum=5, maximum=50, value=15,
                    label="Height (%)",
                    info="Height of subtitle box"
                )
        
        # Options
        with gr.Row():
            aspect_lock = gr.Checkbox(
                value=True,
                label="Lock Aspect Ratio",
                info="Maintain 16:9 aspect ratio for subtitle box"
            )
            auto_process = gr.Checkbox(
                value=False,
                label="Auto-process on change",
                info="Automatically process when parameters change"
            )
        
        # Reset button
        with gr.Row():
            reset_btn = gr.Button("↺ Reset to Defaults", variant="secondary")
        
        # Process button
        process_btn = gr.Button("🔍 Process Video", variant="primary", size="lg")
        
        # Progress indicator
        progress_bar = gr.Progress()
        
        # Results section
        gr.Markdown("### 📄 Results")
        
        with gr.Row():
            with gr.Column():
                # Subtitle list with status
                subtitle_list = gr.Dataframe(
                    headers=["ID", "Start", "End", "Status", "Text"],
                    datatype=["number", "str", "str", "str", "str"],
                    label="Extracted Subtitles",
                    wrap=True,
                    max_height=400
                )
            
            with gr.Column():
                # Export options
                with gr.Accordion("📥 Export Options", open=True):
                    srt_output = gr.Textbox(
                        label="SRT Format",
                        lines=10,
                        max_lines=20
                    )
                    
                    with gr.Row():
                        download_srt = gr.DownloadButton(
                            label="Download SRT",
                            variant="secondary"
                        )
                        download_txt = gr.DownloadButton(
                            label="Download TXT",
                            variant="secondary"
                        )
                        download_json = gr.DownloadButton(
                            label="Download JSON",
                            variant="secondary"
                        )
        
        # Preview section
        gr.Markdown("### 👁️ Active Subtitle Preview")
        
        subtitle_preview = gr.JSON(
            label="Active Subtitles Status",
            info="Shows subtitle status during video playback"
        )
        
        # Hidden state for processed data
        processed_data = gr.State()
        
        # Event handlers
        def update_video_info(video_path):
            """Update video information when video is uploaded"""
            if video_path:
                width, height, fps, duration = get_video_info(video_path)
                info = {
                    "Resolution": f"{width} x {height}",
                    "FPS": round(fps, 2) if fps else 0,
                    "Duration": f"{duration:.2f}s" if duration else "0s"
                }
                return info, {"visible": True}, duration
            return {}, {"visible": False}, 0
        
        def update_end_time_on_load(info, duration):
            """Update end time when video is loaded"""
            if duration > 0:
                return duration
            return 10
        
        def on_aspect_lock_change(lock, width, height):
            """Handle aspect ratio lock"""
            if lock and width:
                # Maintain 16:9 ratio
                new_height = width * 9 / 16
                return width, new_height
            return width, height
        
        def on_process(
            video_path, x, y, width, height,
            language, start, end, progress_bar
        ):
            """Handle video processing"""
            if not video_path:
                return (
                    gr.update(),
                    "Please upload a video first.",
                    "",
                    None
                )
            
            # Get video dimensions
            vw, vh, _, _ = get_video_info(video_path)
            
            # Convert percentages to pixels
            px = int(x / 100 * vw)
            py = int(y / 100 * vh)
            pw = int(width / 100 * vw)
            ph = int(height / 100 * vh)
            
            subtitles, srt, json_output = process_video_ocr(
                video_path, px, py, pw, ph,
                language, start, end,
                progress_bar
            )
            
            if subtitles is None:
                return (
                    gr.update(),
                    srt,  # Error message
                    "",
                    None
                )
            
            # Format for dataframe
            df_data = []
            for sub in subtitles:
                df_data.append([
                    sub.get("id", 0),
                    format_srt_time(sub["start"]).split(",")[0],
                    format_srt_time(sub["end"]).split(",")[0],
                    "entered",
                    sub["text"][:50] + "..." if len(sub["text"]) > 50 else sub["text"]
                ])
            
            # Create downloadable files
            with tempfile.NamedTemporaryFile(mode='w', suffix='.srt', delete=False) as f:
                f.write(srt)
                srt_path = f.name
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
                txt = generate_txt(subtitles)
                f.write(txt)
                txt_path = f.name
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                f.write(json_output)
                json_path = f.name
            
            return (
                df_data,
                srt,
                {"paths": [srt_path, txt_path, json_path]},
                json_output
            )
        
        def reset_defaults():
            """Reset all settings to defaults"""
            return (
                10,  # x
                80,  # y
                80,  # width
                15,  # height
                "English",
                0,
                10
            )
        
        def update_active_subtitles(subtitle_json, current_time):
            """Update active subtitle status based on current time"""
            if not subtitle_json:
                return {}
            
            try:
                data = json.loads(subtitle_json) if isinstance(subtitle_json, str) else subtitle_json
                subtitles = data.get("subtitles", [])
            except:
                return {}
            
            active = {}
            for sub in subtitles:
                start = sub.get("start", 0)
                end = sub.get("end", 0)
                sub_id = sub.get("id", 0)
                
                if current_time < start:
                    status = "entered"
                elif start <= current_time <= end:
                    status = "active"
                else:
                    status = "exited"
                
                active[f"Subtitle {sub_id}"] = {
                    "status": status,
                    "text": sub["text"][:30] + "..." if len(sub["text"]) > 30 else sub["text"],
                    "time": f"{start:.2f}s - {end:.2f}s"
                }
            
            return active
        
        # Bind events
        video_input.change(
            update_video_info,
            inputs=[video_input],
            outputs=[video_info, video_info, end_time]
        ).then(
            update_end_time_on_load,
            inputs=[video_info, end_time],
            outputs=[end_time]
        )
        
        aspect_lock.change(
            on_aspect_lock_change,
            inputs=[aspect_lock, box_width, box_height],
            outputs=[box_width, box_height]
        )
        
        reset_btn.click(
            reset_defaults,
            outputs=[x_pos, y_pos, box_width, box_height, language, start_time, end_time]
        )
        
        process_btn.click(
            on_process,
            inputs=[
                video_input, x_pos, y_pos, box_width, box_height,
                language, start_time, end_time, progress_bar
            ],
            outputs=[subtitle_list, srt_output, processed_data, subtitle_preview]
        )
        
        # Setup download buttons
        def prepare_downloads(data_state):
            if data_state and "paths" in data_state:
                paths = data_state["paths"]
                return (
                    paths[0],  # srt
                    paths[1],  # txt
                    paths[2]   # json
                )
            return None, None, None
        
        processed_data.change(
            prepare_downloads,
            inputs=[processed_data],
            outputs=[download_srt, download_txt, download_json]
        )
    
    return demo


# Create and launch the demo
if __name__ == "__main__":
    demo = create_demo()
    
    # Custom CSS for better UI
    custom_css = """
    .gradio-container {
        max-width: 1400px !important;
        margin: auto;
    }
    .subtitle-active {
        background-color: #4CAF50 !important;
        color: white !important;
        box-shadow: 0 0 10px #4CAF50;
    }
    .subtitle-entered {
        background-color: #2196F3 !important;
        color: white !important;
    }
    .subtitle-exited {
        background-color: #9E9E9E !important;
        color: white !important;
        opacity: 0.6;
    }
    """
    
    demo.launch(
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="purple",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter")
        ),
        css=custom_css,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
        ]
    )