Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import cv2 | |
| from paddleocr import PaddleOCR | |
| import datetime | |
| import difflib | |
| import os | |
| # Khởi tạo mô hình OCR tiếng Việt với các tham số tương thích phiên bản mới | |
| ocr = PaddleOCR(lang='vi') | |
| def format_timestamp(ms): | |
| """Chuyển đổi mili-giây sang định dạng chuẩn SRT: HH:MM:SS,mmm""" | |
| seconds, milliseconds = divmod(ms, 1000) | |
| minutes, seconds = divmod(seconds, 60) | |
| hours, minutes = divmod(minutes, 60) | |
| return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d},{int(milliseconds):03d}" | |
| def is_similar(text1, text2, threshold=0.85): | |
| """So sánh độ giống nhau giữa 2 chuỗi để gộp khung hình (tránh nhiễu OCR)""" | |
| return difflib.SequenceMatcher(None, text1, text2).ratio() > threshold | |
| def process_video(video_path, progress=gr.Progress()): | |
| if not video_path: | |
| return None, "Vui lòng tải lên một video." | |
| cap = cv2.VideoCapture(video_path) | |
| # Lấy thông số video | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| # Thiết lập quét 2 khung hình mỗi giây để tối ưu tốc độ | |
| frames_to_skip = int(fps / 2) if fps > 0 else 15 | |
| interval_ms = (frames_to_skip / fps) * 1000 if fps > 0 else 500 | |
| srt_blocks = [] | |
| current_text = "" | |
| start_time_ms = 0 | |
| last_seen_ms = 0 | |
| frame_count = 0 | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Chỉ quét các khung hình được chỉ định | |
| if frame_count % frames_to_skip == 0: | |
| current_ms = cap.get(cv2.CAP_PROP_POS_MSEC) | |
| # Báo cáo tiến độ lên giao diện Gradio | |
| progress(frame_count / total_frames, desc=f"Đang xử lý thời gian: {format_timestamp(current_ms)}") | |
| # TỐI ƯU HÓA: Chỉ cắt lấy 30% khu vực bên dưới của video (nơi đặt phụ đề) | |
| h, w, _ = frame.shape | |
| cropped_frame = frame[int(h * 0.7):h, 0:w] | |
| # Chạy OCR trên ảnh đã cắt | |
| result = ocr.ocr(cropped_frame, cls=False) | |
| text_lines = [] | |
| if result and result[0]: | |
| text_lines = [line[1][0] for line in result[0]] | |
| frame_text = " ".join(text_lines).strip() | |
| # THUẬT TOÁN GỘP TIMECODE | |
| if frame_text: | |
| if not current_text: | |
| # Bắt đầu một câu phụ đề mới | |
| current_text = frame_text | |
| start_time_ms = current_ms | |
| last_seen_ms = current_ms | |
| elif is_similar(current_text, frame_text): | |
| # Kéo dài thời gian hiển thị nếu câu phụ đề vẫn giống nhau | |
| last_seen_ms = current_ms | |
| else: | |
| # Phụ đề đã thay đổi -> Lưu câu cũ vào danh sách SRT | |
| srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text)) | |
| # Cập nhật câu mới | |
| current_text = frame_text | |
| start_time_ms = current_ms | |
| last_seen_ms = current_ms | |
| else: | |
| # Không có text -> Lưu câu cũ nếu đang có | |
| if current_text: | |
| srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text)) | |
| current_text = "" | |
| frame_count += 1 | |
| # Lưu câu cuối cùng nếu video kết thúc đột ngột | |
| if current_text: | |
| srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text)) | |
| cap.release() | |
| # Định dạng thành cấu trúc file SRT chuẩn | |
| srt_content = "" | |
| for i, (start_ms, end_ms, text) in enumerate(srt_blocks, 1): | |
| srt_content += f"{i}\n" | |
| srt_content += f"{format_timestamp(start_ms)} --> {format_timestamp(end_ms)}\n" | |
| srt_content += f"{text}\n\n" | |
| # Xuất file | |
| output_filename = "output_subtitles.srt" | |
| with open(output_filename, "w", encoding="utf-8") as f: | |
| f.write(srt_content) | |
| return output_filename, f"Hoàn tất! Trích xuất được {len(srt_blocks)} dòng phụ đề." | |
| # Giao diện Gradio | |
| with gr.Blocks(title="Video OCR to SRT") as demo: | |
| gr.Markdown("# Trích xuất Phụ đề Video (OCR -> SRT)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_input = gr.Video(label="Tải video lên") | |
| process_btn = gr.Button("Bắt đầu trích xuất", variant="primary") | |
| with gr.Column(): | |
| file_output = gr.File(label="Tải xuống file SRT") | |
| status_output = gr.Textbox(label="Trạng thái", interactive=False) | |
| process_btn.click( | |
| fn=process_video, | |
| inputs=[video_input], | |
| outputs=[file_output, status_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |