fix of videos, and better code by gdr/daroche

app.py

@@ -1,114 +1,14 @@
-import torch
-from PIL import Image
-from RealESRGAN import RealESRGAN
 import gradio as gr
-import os
-from random import randint
-import shutil
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model2 = RealESRGAN(device, scale=2)
-model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
-model4 = RealESRGAN(device, scale=4)
-model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
-model8 = RealESRGAN(device, scale=8)
-model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
-
-
-def inference_image(image, size):
-    global model2
-    global model4
-    global model8
-    if image is None:
-        raise gr.Error("Image not uploaded")
-
-    width, height = image.size
-    if width >= 5000 or height >= 5000:
-        raise gr.Error("The image is too large.")
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    if size == '2x':
-        try:
-            result = model2.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model2 = RealESRGAN(device, scale=2)
-            model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
-            result = model2.predict(image.convert('RGB'))
-    elif size == '4x':
-        try:
-            result = model4.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model4 = RealESRGAN(device, scale=4)
-            model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
-            result = model2.predict(image.convert('RGB'))
-    else:
-        try:
-            result = model8.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model8 = RealESRGAN(device, scale=8)
-            model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
-            result = model2.predict(image.convert('RGB'))
-
-    print(f"Image size ({device}): {size} ... OK")
-    return result
-
-
-
-def inference_video(video, size):
-    _id = randint(1, 10000)
-    INPUT_DIR = "tmp"
-    os.makedirs(INPUT_DIR, exist_ok=True)
-    os.chdir(INPUT_DIR)
-    
-    upload_folder = 'upload'
-    result_folder = 'results'
-    video_folder = 'videos'
-    video_result_folder = 'results_videos'
-    video_mp4_result_folder = 'results_mp4_videos'
-    result_restored_imgs_folder = 'restored_imgs'
-    
-    os.makedirs(upload_folder, exist_ok=True)
-    
-    os.makedirs(video_folder, exist_ok=True)
-    
-    os.makedirs(video_result_folder, exist_ok=True)
-        
-    os.makedirs(video_mp4_result_folder, exist_ok=True)
-    
-    os.makedirs(result_folder, exist_ok=True)
-    
-    os.chdir("results")
-    os.makedirs(result_restored_imgs_folder, exist_ok=True)
-    os.chdir("..")
-    try:
-        # Specify the desired output file path with the custom name and ".mp4" extension
-        output_file_path = f"/{INPUT_DIR}/videos/input.mp4"
-
-        # Save the video input to the specified file path
-        with open(output_file_path, 'wb') as output_file:
-            output_file.write(video)
-        print(f"Video input saved as {output_file_path}")
-    except Exception as e:
-        print(f"Error saving video input: {str(e)}")
-
-    os.chdir("..")
-    os.system("python inference_video.py")
-    return os.path.join(f'/{INPUT_DIR}/results_mp4_videos/', 'input.mp4')
-    
 
+from infer import infer_image, infer_video
 
 input_image = gr.Image(type='pil', label='Input Image')
-input_model_image = gr.Radio(['2x', '4x', '8x'], type="value", value="4x", label="Model Upscale/Enhance Type")
+input_model_image = gr.Radio([('x2', 2), ('x4', 4), ('x8', 8)], type="value", value=4, label="Model Upscale/Enhance Type")
 submit_image_button = gr.Button('Submit')
 output_image = gr.Image(type="filepath", label="Output Image")
 
 tab_img = gr.Interface(
-    fn=inference_image,
+    fn=infer_image,
     inputs=[input_image, input_model_image],
     outputs=output_image,
     title="Real-ESRGAN Pytorch",
@@ -116,12 +16,12 @@ tab_img = gr.Interface(
 )
 
 input_video = gr.Video(label='Input Video')
-input_model_video = gr.Radio(['2x', '4x', '8x'], type="value", value="4x", label="Model Upscale/Enhance Type")
+input_model_video = gr.Radio([('x2', 2), ('x4', 4), ('x8', 8)], type="value", value=4, label="Model Upscale/Enhance Type")
 submit_video_button = gr.Button('Submit')
 output_video = gr.Video(label='Output Video')
 
 tab_vid = gr.Interface(
-    fn=inference_video,
+    fn=infer_video,
     inputs=[input_video, input_model_video],
     outputs=output_video,
     title="Real-ESRGAN Pytorch",
@@ -130,6 +30,4 @@ tab_vid = gr.Interface(
 
 demo = gr.TabbedInterface([tab_img, tab_vid], ["Image", "Video"])
 
-
-
 demo.launch(debug=True, show_error=True)

infer.py

@@ -0,0 +1,260 @@
+# import cv2
+# from os.path import isfile, join
+# import subprocess
+# import os
+# from RealESRGAN import RealESRGAN
+# import torch
+# import gradio as gr
+
+# IMAGE_FORMATS = ('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')
+
+# def inference_image(image, size):
+#     global model2
+#     global model4
+#     global model8
+#     if image is None:
+#         raise gr.Error("Image not uploaded")
+
+#     width, height = image.size
+#     if width >= 5000 or height >= 5000:
+#         raise gr.Error("The image is too large.")
+
+#     if torch.cuda.is_available():
+#         torch.cuda.empty_cache()
+
+#     if size == '2x':
+#         try:
+#             result = model2.predict(image.convert('RGB'))
+#         except torch.cuda.OutOfMemoryError as e:
+#             print(e)
+#             model2 = RealESRGAN(device, scale=2)
+#             model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
+#             result = model2.predict(image.convert('RGB'))
+#     elif size == '4x':
+#         try:
+#             result = model4.predict(image.convert('RGB'))
+#         except torch.cuda.OutOfMemoryError as e:
+#             print(e)
+#             model4 = RealESRGAN(device, scale=4)
+#             model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
+#             result = model2.predict(image.convert('RGB'))
+#     else:
+#         try:
+#             result = model8.predict(image.convert('RGB'))
+#         except torch.cuda.OutOfMemoryError as e:
+#             print(e)
+#             model8 = RealESRGAN(device, scale=8)
+#             model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
+#             result = model2.predict(image.convert('RGB'))
+
+#     print(f"Frame of the Video size ({device}): {size} ... OK")
+#     return result
+
+
+# # assign directory
+# directory = 'videos' #PATH_WITH_INPUT_VIDEOS
+# zee = 0
+
+# def convert_frames_to_video(pathIn,pathOut,fps):
+#     global INPUT_DIR
+#     cap = cv2.VideoCapture(f'/{INPUT_DIR}/videos/input.mp4')
+#     fps = cap.get(cv2.CAP_PROP_FPS)
+#     frame_array = []
+#     files = [f for f in os.listdir(pathIn) if isfile(join(pathIn, f))]
+#     #for sorting the file names properly
+#     files.sort(key = lambda x: int(x[5:-4]))
+#     size2 = (0,0)
+
+#     for i in range(len(files)):
+#         filename=pathIn + files[i]
+#         #reading each files
+#         img = cv2.imread(filename)
+#         height, width, layers = img.shape
+#         size = (width,height)
+#         size2 = size
+#         print(filename)
+#         #inserting the frames into an image array
+#         frame_array.append(img)
+#     out = cv2.VideoWriter(pathOut,cv2.VideoWriter_fourcc(*'DIVX'), fps, size2)
+#     for i in range(len(frame_array)):
+#         # writing to a image array
+#         out.write(frame_array[i])
+#     out.release()
+
+
+# for filename in os.listdir(directory):
+
+#     f = os.path.join(directory, filename)
+#     # checking if it is a file
+#     if os.path.isfile(f):
+
+
+#       print("PROCESSING :"+str(f)+"\n")
+#       # Read the video from specified path
+
+#       #video to frames
+#       cam = cv2.VideoCapture(str(f))
+
+#       try:
+
+#           # PATH TO STORE VIDEO FRAMES
+#           if not os.path.exists(f'/{INPUT_DIR}/upload/'):
+#               os.makedirs(f'/{INPUT_DIR}/upload/')
+
+#       # if not created then raise error
+#       except OSError:
+#           print ('Error: Creating directory of data')
+
+#       # frame
+#       currentframe = 0
+
+
+#       while(True):
+
+#           # reading from frame
+#           ret,frame = cam.read()
+
+#           if ret:
+#               # if video is still left continue creating images
+#               name = f'/{INPUT_DIR}/upload/frame' + str(currentframe) + '.jpg'
+
+#               # writing the extracted images
+#               cv2.imwrite(name, frame)
+
+
+#                 # increasing counter so that it will
+#                 # show how many frames are created
+#               currentframe += 1
+#               print(currentframe)
+#           else:
+#               #deletes all the videos you uploaded for upscaling
+#               #for f in os.listdir(video_folder):
+#               #  os.remove(os.path.join(video_folder, f))
+
+#               break
+
+#         # Release all space and windows once done
+#       cam.release()
+#       cv2.destroyAllWindows()
+
+#       #apply super-resolution on all frames of a video
+
+#       # Specify the directory path
+#       all_frames_path = f"/{INPUT_DIR}/upload/"
+
+#       # Get a list of all files in the directory
+#       file_names = os.listdir(all_frames_path)
+
+#       # process the files
+#       for file_name in file_names:
+#         inference_image(f"/{INPUT_DIR}/upload/{file_name}")
+
+
+#       #convert super res frames to .avi
+#       pathIn = f'/{INPUT_DIR}/results/restored_imgs/'
+
+#       zee = zee+1
+#       fName = "video"+str(zee)
+#       filenameVid = f"{fName}.avi"
+
+#       pathOut = f"/{INPUT_DIR}/results_videos/"+filenameVid
+
+#       convert_frames_to_video(pathIn, pathOut, fps)
+
+
+#       #convert .avi to .mp4
+#       src = f'/{INPUT_DIR}/results_videos/'
+#       dst = f'/{INPUT_DIR}/results_mp4_videos/'
+
+#       for root, dirs, filenames in os.walk(src, topdown=False):
+#           #print(filenames)
+#           for filename in filenames:
+#               print('[INFO] 1',filename)
+#               try:
+#                   _format = ''
+#                   if ".flv" in filename.lower():
+#                       _format=".flv"
+#                   if ".mp4" in filename.lower():
+#                       _format=".mp4"
+#                   if ".avi" in filename.lower():
+#                       _format=".avi"
+#                   if ".mov" in filename.lower():
+#                       _format=".mov"
+
+#                   inputfile = os.path.join(root, filename)
+#                   print('[INFO] 1',inputfile)
+#                   outputfile = os.path.join(dst, filename.lower().replace(_format, ".mp4"))
+#                   subprocess.call(['ffmpeg', '-i', inputfile, outputfile])
+#               except:
+#                   print("An exception occurred")
+
+from PIL import Image
+import cv2 as cv
+import torch
+from RealESRGAN import RealESRGAN
+import tempfile
+import numpy as np
+import tqdm
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def infer_image(img: Image.Image, size_modifier: int ) -> Image.Image:
+    if img is None:
+        raise Exception("Image not uploaded")
+    
+    width, height = img.size
+    
+    if width >= 5000 or height >= 5000:
+        raise Exception("The image is too large.")
+
+    model = RealESRGAN(device, scale=size_modifier)
+    model.load_weights(f'weights/RealESRGAN_x{size_modifier}.pth', download=False)
+
+    result = model.predict(img.convert('RGB'))
+    print(f"Image size ({device}): {size_modifier} ... OK")
+    return result
+
+def infer_video(video_filepath: str, size_modifier: int) -> str:
+    model = RealESRGAN(device, scale=size_modifier)
+    model.load_weights(f'weights/RealESRGAN_x{size_modifier}.pth', download=False)
+
+    cap = cv.VideoCapture(video_filepath)
+    
+    tmpfile = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    vid_output = tmpfile.name
+    tmpfile.close()
+
+    vid_writer = cv.VideoWriter(
+        vid_output,
+        fourcc=cv.VideoWriter.fourcc(*'mp4v'),
+        fps=cap.get(cv.CAP_PROP_FPS),
+        frameSize=(int(cap.get(cv.CAP_PROP_FRAME_WIDTH)) * size_modifier, int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)) * size_modifier)
+    )
+
+    n_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
+
+    # while cap.isOpened():
+    for _ in tqdm.tqdm(range(n_frames)):
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
+        frame = Image.fromarray(frame)
+
+        upscaled_frame = model.predict(frame.convert('RGB'))
+        
+        upscaled_frame = np.array(upscaled_frame)
+        upscaled_frame = cv.cvtColor(upscaled_frame, cv.COLOR_RGB2BGR)
+
+        print(upscaled_frame.shape)
+
+        vid_writer.write(upscaled_frame)
+
+    vid_writer.release()
+
+    print(f"Video file : {video_filepath}")
+
+    return vid_output
+                                
+