Update app.py

app.py

@@ -1,106 +1,48 @@
 import gradio as gr
+import subprocess
 import os
-import cv2
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from realesrgan import RealESRGANer
-import shutil
 
-os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
-os.makedirs('/tmp/matplotlib', exist_ok=True)
-
-# Thư mục tạm thời để lưu trữ file tải lên
-UPLOAD_FOLDER = "uploads"
-OUTPUT_FOLDER = "outputs"
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs(OUTPUT_FOLDER, exist_ok=True)
-
-# Khởi tạo mô hình Real-ESRGAN
-model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-upsampler = RealESRGANer(
-    scale=4,
-    model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
-    model=model,
-    tile=0,
-    tile_pad=10,
-    pre_pad=0,
-    half=True
-)
-
-def enhance_video(input_video):
+def enhance_video_gradio(input_file):
     """
     Hàm xử lý video đầu vào và trả về video đã tăng cường.
     """
-    # Lưu file tải lên vào thư mục uploads
-    input_path = os.path.join(UPLOAD_FOLDER, "input.mp4")
-    output_path = os.path.join(OUTPUT_FOLDER, "enhanced_output.mp4")
-    input_video.save(input_path)
-
-    # Tách video thành khung hình
-    frames_folder = os.path.join(UPLOAD_FOLDER, "frames")
-    os.makedirs(frames_folder, exist_ok=True)
-    extract_frames(input_path, frames_folder)
-
-    # Xử lý từng khung hình với Real-ESRGAN
-    enhanced_frames_folder = os.path.join(OUTPUT_FOLDER, "enhanced_frames")
-    os.makedirs(enhanced_frames_folder, exist_ok=True)
-    enhance_frames(frames_folder, enhanced_frames_folder)
-
-    # Ghép các khung hình đã tăng cường thành video
-    combine_frames_to_video(enhanced_frames_folder, output_path)
+    # Tạo đường dẫn tạm thời cho file đầu vào và đầu ra
+    input_path = "input.mp4"
+    output_path = "enhanced_output.mp4"
 
-    # Xóa các thư mục tạm thời
-    shutil.rmtree(frames_folder)
-    shutil.rmtree(enhanced_frames_folder)
+    # Lưu file tải lên vào đường dẫn tạm thời
+    with open(input_path, "wb") as f:
+        f.write(input_file.read())
 
-    return output_path
-
-def extract_frames(video_path, output_folder):
-    """
-    Tách video thành các khung hình.
-    """
-    cap = cv2.VideoCapture(video_path)
-    frame_count = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame_path = os.path.join(output_folder, f"frame_{frame_count:04d}.png")
-        cv2.imwrite(frame_path, frame)
-        frame_count += 1
-    cap.release()
-
-def enhance_frames(input_folder, output_folder):
-    """
-    Tăng cường từng khung hình bằng Real-ESRGAN.
-    """
-    for frame_name in os.listdir(input_folder):
-        frame_path = os.path.join(input_folder, frame_name)
-        img = cv2.imread(frame_path, cv2.IMREAD_UNCHANGED)
-        enhanced_img, _ = upsampler.enhance(img, outscale=4)
-        output_path = os.path.join(output_folder, frame_name)
-        cv2.imwrite(output_path, enhanced_img)
-
-def combine_frames_to_video(frames_folder, output_path):
-    """
-    Ghép các khung hình thành video.
-    """
-    frame_rate = 30  # Số khung hình/giây
-    first_frame = os.path.join(frames_folder, os.listdir(frames_folder)[0])
-    height, width, _ = cv2.imread(first_frame).shape
-
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    video_writer = cv2.VideoWriter(output_path, fourcc, frame_rate, (width, height))
-
-    for frame_name in sorted(os.listdir(frames_folder)):
-        frame_path = os.path.join(frames_folder, frame_name)
-        frame = cv2.imread(frame_path)
-        video_writer.write(frame)
-
-    video_writer.release()
+    # Định nghĩa các bộ lọc ffmpeg
+    filters = [
+        "unsharp=5:5:1.0",  # Làm sắc nét video
+        "hqdn3d=4:3:6:4",   # Giảm nhiễu
+        "eq=contrast=1.2:brightness=0.1"  # Điều chỉnh tương phản và độ sáng
+    ]
+    
+    # Kết hợp các bộ lọc thành chuỗi
+    filter_chain = ",".join(filters)
+    
+    # Xây dựng lệnh ffmpeg
+    command = [
+        "ffmpeg",
+        "-i", input_path,          # File đầu vào
+        "-vf", filter_chain,       # Áp dụng các bộ lọc
+        "-c:a", "copy",            # Sao chép âm thanh nguyên bản
+        output_path                # File đầu ra
+    ]
+    
+    try:
+        # Thực thi lệnh ffmpeg
+        subprocess.run(command, check=True)
+        return output_path
+    except subprocess.CalledProcessError as e:
+        return f"Lỗi khi xử lý video: {e}"
 
 # Tạo giao diện Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Video Enhancement with Real-ESRGAN")
+    gr.Markdown("# Video Enhancement with FFmpeg")
     gr.Markdown("Upload a video and get an enhanced version with improved quality.")
     
     with gr.Row():
@@ -108,7 +50,7 @@ with gr.Blocks() as demo:
         output_video = gr.Video(label="Enhanced Video")
     
     submit_button = gr.Button("Enhance Video")
-    submit_button.click(enhance_video, inputs=input_video, outputs=output_video)
+    submit_button.click(enhance_video_gradio, inputs=input_video, outputs=output_video)
 
 # Khởi chạy ứng dụng
 demo.launch()