nonoisefb / app.py
Harshitaraina's picture
Create app.py
1f7f82b verified
Raw
History Blame Contribute Delete
2.62 kB
import gradio as gr
import torch
import torchaudio
from torchaudio.transforms import Resample
import moviepy.editor as mp
import numpy as np
from denoiser.pretrained import master64 # Import Facebook denoiser pre-trained model
from denoiser.denoiser import Denoiser # Denoising wrapper
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device - ", DEVICE)
# Load Facebook denoiser model
model = master64()
den = Denoiser(model).to(DEVICE)
den.eval()
def identity(video_path):
print(video_path)
# Extract audio from the video
video = mp.VideoFileClip(video_path)
audio = video.audio
wav_file = "tmp.wav"
audio.write_audiofile(wav_file)
print("Wav stored.")
# Load audio
waveform, sr = torchaudio.load(wav_file)
waveform = waveform.to(DEVICE)
# Resample if necessary
target_sr = 48000
if sr != target_sr:
resampler = Resample(orig_freq=sr, new_freq=target_sr).to(DEVICE)
waveform = resampler(waveform)
sr = target_sr
# Process audio in chunks to avoid memory issues
chunk_duration = 10 # seconds
chunk_size = sr * chunk_duration
num_chunks = int(np.ceil(waveform.shape[1] / chunk_size))
enhanced_chunks = []
for i in range(num_chunks):
start = i * chunk_size
end = min((i + 1) * chunk_size, waveform.shape[1])
chunk = waveform[:, start:end]
enhanced_chunk = den(chunk.unsqueeze(0))
enhanced_chunks.append(enhanced_chunk.squeeze(0))
# Combine enhanced audio
enhanced_audio = torch.cat(enhanced_chunks, dim=1).cpu()
# Save enhanced audio
output_audio_path = "enhanced_aud.wav"
torchaudio.save(output_audio_path, enhanced_audio, sr)
# Replace audio in video
enhanced_audio_clip = mp.AudioFileClip(output_audio_path)
final_video = video.set_audio(enhanced_audio_clip)
output_video_path = "output_video.mp4"
final_video.write_videofile(output_video_path,
codec='libx264',
audio_codec='aac',
temp_audiofile='temp-audio.m4a',
remove_temp=True)
return output_video_path
demo = gr.Interface(
fn=identity,
title="NoNoise - THE BEST AUDIO DENOISER",
description="NoNoise is the only platform you need for removing all kinds of background noise from your videos!!",
examples=[['audiopure_og.mov'], ['example.mp4']],
cache_examples=True,
inputs=gr.Video(label="Input Video", source="upload"),
outputs=gr.Video(label="Output Video"),
)
demo.launch()