File size: 3,741 Bytes
7413d5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 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() |