Create app.py

app.py

@@ -0,0 +1,111 @@
+import gradio as gr
+import os
+import cv2
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+import shutil
+
+# 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):
+    """
+    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)
+
+    # Xóa các thư mục tạm thời
+    shutil.rmtree(frames_folder)
+    shutil.rmtree(enhanced_frames_folder)
+
+    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()
+
+# Tạo giao diện Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# Video Enhancement with Real-ESRGAN")
+    gr.Markdown("Upload a video and get an enhanced version with improved quality.")
+    
+    with gr.Row():
+        input_video = gr.File(label="Upload Video")
+        output_video = gr.Video(label="Enhanced Video")
+    
+    submit_button = gr.Button("Enhance Video")
+    submit_button.click(enhance_video, inputs=input_video, outputs=output_video)
+
+# Khởi chạy ứng dụng
+demo.launch()