app.py

import shlex
import subprocess
import spaces
import torch
import os
import gradio as gr

# install packages for mamba
def install_mamba():
    #subprocess.run(shlex.split("pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118"))
    subprocess.run(shlex.split("pip install https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
    #subprocess.run(shlex.split("pip install numpy==1.26.4"))

def clone_github():
    subprocess.run([
        "git", "clone",
        f"https://RoyChao19477:{os.environ['GITHUB_TOKEN']}@github.com/RoyChao19477/for_HF_AVSEMamba.git"
    ])

install_mamba()
clone_github()

ABOUT = """
# SEMamba: Speech Enhancement
A Mamba-based model that denoises real-world audio.
Upload or record a noisy clip and click **Enhance** to hear + see its spectrogram.
"""


import torch
import yaml
import librosa
import librosa.display
import matplotlib
import numpy as np
import soundfile as sf
import matplotlib.pyplot as plt
from models.stfts    import mag_phase_stft, mag_phase_istft
from models.generator import SEMamba
from models.pcs400   import cal_pcs
from ultralytics import YOLO
import supervision as sv

import gradio as gr
import cv2
import os
import tempfile
from ultralytics import YOLO
from moviepy import ImageSequenceClip

# Load face detector
model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available

@spaces.GPU
def extract_faces(video_file):
    cap = cv2.VideoCapture(video_file)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frames = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Inference
        results = model(frame, verbose=False)[0]
        for box in results.boxes:
            # version 1
            # x1, y1, x2, y2 = map(int, box.xyxy[0])
            # version 2
            h, w, _ = frame.shape
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
            pad_ratio = 0.3  # 30% padding

            dx = (x2 - x1) * pad_ratio
            dy = (y2 - y1) * pad_ratio

            x1 = int(max(0, x1 - dx))
            y1 = int(max(0, y1 - dy))
            x2 = int(min(w, x2 + dx))
            y2 = int(min(h, y2 + dy))
            face_crop = frame[y1:y2, x1:x2]
            if face_crop.size != 0:
                resized = cv2.resize(face_crop, (224, 224))
                frames.append(resized)
            break  # only one face per frame

    cap.release()

    # Save as video
    tmpdir = tempfile.mkdtemp()
    output_path = os.path.join(tmpdir, "face_only_video.mp4")
    clip = ImageSequenceClip([cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in frames], fps=25)
    clip.write_videofile(output_path, codec="libx264", audio=False)

    return output_path

iface = gr.Interface(
    fn=extract_faces,
    inputs=gr.Video(label="Upload or record your video"),
    outputs=gr.Video(label="Detected Face Only Video"),
    title="Face Detector",
    description="Upload or record a video. We'll crop face regions and return a face-only video."
)

iface.launch()



ckpt = "ckpts/SEMamba_advanced.pth"
cfg_f = "recipes/SEMamba_advanced.yaml"

# load config
with open(cfg_f, 'r') as f:
    cfg = yaml.safe_load(f)


# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cuda"
model  = SEMamba(cfg).to(device)
#sdict  = torch.load(ckpt, map_location=device)
#model.load_state_dict(sdict["generator"])
#model.eval()

@spaces.GPU
def enhance(filepath, model_name):
    # Load model based on selection
    ckpt_path = {
        "VCTK-Demand": "ckpts/SEMamba_advanced.pth",
        "VCTK+DNS": "ckpts/vd.pth"
    }[model_name]

    print("Loading:", ckpt_path)
    model.load_state_dict(torch.load(ckpt_path, map_location=device)["generator"])
    model.eval()
    with torch.no_grad():
        # load & resample
        wav, orig_sr = librosa.load(filepath, sr=None)
        noisy_wav = wav.copy()
        if orig_sr != 16000:
            wav = librosa.resample(wav, orig_sr=orig_sr, target_sr=16000)
        x = torch.from_numpy(wav).float().to(device)
        norm = torch.sqrt(len(x)/torch.sum(x**2))
        #x = (x * norm).unsqueeze(0)
        x = (x * norm)

        # split into 4s segments (64000 samples)
        segment_len = 4 * 16000
        chunks = x.split(segment_len)
        enhanced_chunks = []

        for chunk in chunks:
            if len(chunk) < segment_len:
                #pad = torch.zeros(segment_len - len(chunk), device=chunk.device)
                pad = (torch.randn(segment_len - len(chunk), device=chunk.device) * 1e-4)
                chunk = torch.cat([chunk, pad])
            chunk = chunk.unsqueeze(0)

            amp, pha, _ = mag_phase_stft(chunk, 400, 100, 400, 0.3)
            amp2, pha2, _ = model(amp, pha)
            out = mag_phase_istft(amp2, pha2, 400, 100, 400, 0.3)
            out = (out / norm).squeeze(0)
            enhanced_chunks.append(out)

        out = torch.cat(enhanced_chunks)[:len(x)].cpu().numpy()  # trim padding

        # back to original rate
        if orig_sr != 16000:
            out = librosa.resample(out, orig_sr=16000, target_sr=orig_sr)

        # Normalize
        peak = np.max(np.abs(out))
        if peak > 0.05:
            out = out / peak * 0.85

        # write file
        sf.write("enhanced.wav", out, orig_sr)

        # spectrograms
        fig, axs = plt.subplots(1, 2, figsize=(16, 4))

        # noisy
        D_noisy = librosa.stft(noisy_wav, n_fft=512, hop_length=256)
        S_noisy = librosa.amplitude_to_db(np.abs(D_noisy), ref=np.max)
        librosa.display.specshow(S_noisy, sr=orig_sr, hop_length=256, x_axis="time", y_axis="hz", ax=axs[0], vmax=0)
        axs[0].set_title("Noisy Spectrogram")

        # enhanced
        D_clean = librosa.stft(out, n_fft=512, hop_length=256)
        S_clean = librosa.amplitude_to_db(np.abs(D_clean), ref=np.max)
        librosa.display.specshow(S_clean, sr=orig_sr, hop_length=256, x_axis="time", y_axis="hz", ax=axs[1], vmax=0)
        #librosa.display.specshow(S_clean, sr=16000, hop_length=512, x_axis="time", y_axis="hz", ax=axs[1], vmax=0)
        axs[1].set_title("Enhanced Spectrogram")

        plt.tight_layout()

    return "enhanced.wav", fig

#with gr.Blocks() as demo:
#    gr.Markdown(ABOUT)
#    input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
#    enhance_btn = gr.Button("Enhance")
#    output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
#    plot_output = gr.Plot(label="Spectrograms")
#
#    enhance_btn.click(fn=enhance, inputs=input_audio, outputs=[output_audio, plot_output])
#
#demo.queue().launch()

with gr.Blocks() as demo:
    gr.Markdown(ABOUT)
    input_audio = gr.Audio(label="Input Audio", type="filepath", interactive=True)
    model_choice = gr.Radio(
        label="Choose Model (The use of VCTK+DNS is recommended)",
        choices=["VCTK-Demand", "VCTK+DNS"],
        value="VCTK-Demand"
    )
    enhance_btn = gr.Button("Enhance")
    output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
    plot_output = gr.Plot(label="Spectrograms")

    enhance_btn.click(
        fn=enhance,
        inputs=[input_audio, model_choice],
        outputs=[output_audio, plot_output]
    )
    gr.Markdown("**Note**: The current models are trained on 16kHz audio. Therefore, any input audio not sampled at 16kHz will be automatically resampled before enhancement.")

demo.queue().launch()