Update infer.py
Browse files
infer.py
CHANGED
@@ -25,16 +25,16 @@ def infer_image(img: Image.Image, size_modifier: int ) -> Image.Image:
|
|
25 |
return result
|
26 |
|
27 |
def infer_video(video_filepath: str, size_modifier: int) -> str:
|
28 |
-
# Extract audio from the original video file
|
29 |
-
audio = cv.AudioCapture(video_filepath)
|
30 |
-
audio_data = np.frombuffer(audio.readAll(), dtype=np.int16)
|
31 |
-
audio_array = np.array(audio_data, dtype=np.int16)
|
32 |
-
|
33 |
model = RealESRGAN(device, scale=size_modifier)
|
34 |
model.load_weights(f'weights/RealESRGAN_x{size_modifier}.pth', download=False)
|
35 |
-
|
36 |
-
cap = cv.VideoCapture(video_filepath)
|
37 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
38 |
# Create a temporary file for the output video
|
39 |
tmpfile = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
|
40 |
vid_output = tmpfile.name
|
|
|
25 |
return result
|
26 |
|
27 |
def infer_video(video_filepath: str, size_modifier: int) -> str:
|
|
|
|
|
|
|
|
|
|
|
28 |
model = RealESRGAN(device, scale=size_modifier)
|
29 |
model.load_weights(f'weights/RealESRGAN_x{size_modifier}.pth', download=False)
|
|
|
|
|
30 |
|
31 |
+
# Extract audio from the original video file
|
32 |
+
audio = AudioSegment.from_file(video_filepath, format=video_filepath.split('.')[-1])
|
33 |
+
audio_array = np.array(audio.get_array_of_samples())
|
34 |
+
|
35 |
+
# Create a VideoCapture object for the video file
|
36 |
+
cap = cv2.VideoCapture(video_filepath)
|
37 |
+
|
38 |
# Create a temporary file for the output video
|
39 |
tmpfile = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
|
40 |
vid_output = tmpfile.name
|