Update real_esrgan.py

real_esrgan.py

@@ -1,22 +1,60 @@
+import gradio as gr
 import os
 import tempfile
-import subprocess
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+import torch
+import cv2
 
-def upscale_video(video):
-    with tempfile.TemporaryDirectory() as tempdir:
-        input_path = os.path.join(tempdir, "input.mp4")
-        output_path = os.path.join(tempdir, "output.mp4")
-        
+def upscale_video(video_file):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        input_path = os.path.join(tmpdir, "input.mp4")
+        output_path = os.path.join(tmpdir, "output.mp4")
+
+        # Save uploaded video
         with open(input_path, "wb") as f:
-            f.write(video.read())
-        
-        cmd = [
-            "python", "inference_realesrgan.py",
-            "-i", input_path,
-            "-o", output_path,
-            "-n", "RealESRGAN_x4plus",
-            "--outscale", "2"
-        ]
-        subprocess.run(cmd, check=True)
-        
-        return output_path
+            f.write(video_file.read())
+
+        cap = cv2.VideoCapture(input_path)
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width*2, height*2))
+
+        # Load model
+        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='RealESRGAN_x4plus.pth',
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=True
+        )
+
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            try:
+                output, _ = upsampler.enhance(frame)
+            except RuntimeError as error:
+                print(f"Error: {error}")
+                break
+            output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+            out.write(output)
+
+        cap.release()
+        out.release()
+
+        return output_path
+
+demo = gr.Interface(fn=upscale_video, inputs=gr.Video(), outputs=gr.Video(), title="Video Upscaler (Real-ESRGAN)")
+
+if __name__ == "__main__":
+    demo.launch()