Upload README.md

README.md

@@ -1,3 +1,174 @@
----
-license: apache-2.0
----
+<h1 align="center">CoGenAV: Contrastive-Generative Audio-Visual Representation Learning</h1>
+
+<div align="center" style="display: flex; justify-content: center; align-items: center; gap: 10px;">
+  
+  <a href="https://arxiv.org/pdf/2505.03186" target="_blank">
+    <img src="https://img.shields.io/badge/arXiv-Paper-b31b1b.svg?logo=arXiv" alt="arXiv Paper">
+  </a>
+
+  <a href="https://huggingface.co/detao/CoGenAV" target="_blank">
+    <img src="https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow" alt="HuggingFace Model">
+  </a>
+
+  <a href="https://modelscope.cn/models/hongqi/cogenav" target="_blank">
+      <img src="https://img.shields.io/badge/🤖%20ModelScope-Model-yellow" alt="ModelScope Model">
+  </a>
+
+</div>
+
+
+---
+#### 🚀 Project Overview  
+CoGenAV is a framework for audio-visual representation learning based on **Contrastive-Generative Synchronization**, designed to learn efficient and generalizable audio-visual representations through multimodal alignment of speech, lip movements, and text. The model performs exceptionally well across multiple audio-visual tasks, including:  
+- **Audio-Visual Speech Recognition (AVSR)**  
+- **Visual Speech Recognition (VSR)**  
+- **Audio-Visual Speech Enhancement and Separation (AVSE/AVSS)**  
+- **Active Speaker Detection (ASD)**  
+
+---
+
+## 🏗️ Framework
+
+<p align="center">
+<img src="https://huggingface.co/detao/CoGenAV/resolve/main/cogenav_arch/cogen_arch.png" width=100%>
+<p>
+
+The left panel depicts the Audio-Visual Feature Representation framework and the Contrastive-Generative Synchronization Training methodology. For generative synchronization, we design a Feature Adaptation Module and employ a [frozen pre-trained ASR model](https://github.com/openai/whisper) as the Speech Recognition (SR) head. The right panel demonstrates the application of CoGenAV to diverse downstream tasks, including Visual Speech Recognition (VSR), Audio-Visual Speech Recognition (AVSR), Audio-Visual Speech Separation (AVSS), Audio-Visual Speech Enhancement (AVSE), and Active Speaker Detection (ASD).
+
+---
+
+#### 🌟 Key Advantages  
+1. **Efficient Learning**: High-performance models can be trained with only **223 hours of labeled data** (from the LRS2 dataset).  
+2. **Cross-Task Generalizability**: Unified representation learning allows direct adaptation to various downstream tasks without task-specific architectural adjustments.  
+3. **Robustness**: Performance improves by **70%+** in noisy environments (0 dB SNR), significantly outperforming traditional audio-only models.  
+
+
+---
+#### Usage
+1. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   #Need to ensure that whisper and fairseq is installed
+   pip install -U openai-whisper
+   git clone https://github.com/pytorch/fairseq
+   cd fairseq
+   pip install --editable ./
+
+2. Infer CoGenAV for VSR/AVSR :
+   ```python
+    import whisper
+    from whisper.model import AudioEncoder
+    from infer_vsr_avsr import cogenav_forward
+    from models.cogenav import CoGenAV
+    # Override the Whisper encoder's forward function
+    AudioEncoder.forward = cogenav_forward
+    # Load CoGenAV model
+    cogenav = CoGenAV(cfg_file="config/base.yaml", model_tensor="weights/base_cogenav.pt")
+    # Load Whisper model as SR_Head
+    SR_Head = whisper.load_model("small", download_root="weights/whisper/")
+    SR_Head.encoder.adapter = cogenav.adapter.half()
+    # Prepare input using CoGenAV
+    input_ids = cogenav(video, audio).permute(0, 2, 1)  # For cogenav_av
+    # input_ids = cogenav(video, None).permute(0, 2, 1)  # For cogenav_v
+    # input_ids = cogenav(None, audio).permute(0, 2, 1)  # For cogenav_a
+    # input_ids = audio  # For whisper_a
+    # Decode using Whisper model
+    result = whisper.decode(SR_Head, input_ids, options)[0]
+
+3. Infer CoGenAV for AVSS/AVSE :
+   ```python
+    from models.cogenav import CoGenAV
+    from  models.sepformer import build_Sepformer
+    # Load CoGenAV model
+    cogenav = CoGenAV(cfg_file="config/base.yaml", model_tensor="weights/base_cogenav.pt")
+    # Load sepformer model as avss/avse head
+    sepformer_head = build_Sepformer().cuda()
+    # sep speech with lip feature from mix wav
+    lip_feature = cogenav(video, None,use_upsampler=False)
+    sep_wav = sepformer_head.forward(audio_mix, lip_feature)
+   
+4. Infer script:
+   ```bash
+    python infer_vsr_avsr.py --input_type cogenav_av --model_size large  --cogenav_ckpt weights/large_cogenav.pt
+    python infer_avse_avss.py --task_type avse
+
+## 🎬 Demo
+### Demo For AVSR/VSR
+<table class="center">
+  <tr>
+    <td colspan="2" style="text-align: center; font-weight: bold;">
+      AVSR/VSR
+    </td>
+  </tr>
+  <tr>
+    <td colspan="2" style="text-align: center;">
+      <video src="https://github.com/user-attachments/assets/e44e4606-9ef0-4fc7-a1e0-0add000f8e5f" controls preload></video>
+      <video src="https://github.com/user-attachments/assets/6c0cfe05-e82e-4b05-bd07-f4e0ebf2375f" controls preload></video>
+      <video src="https://github.com/user-attachments/assets/d1190323-dd31-4a74-b2f7-25ce3ec72c35" controls preload></video>
+    </td>
+  </tr>
+</table>
+
+### Demo For AVSS/AVSE
+
+<table style="width:100%; text-align:center;">
+  <tr>
+    <td colspan="2" style="font-weight: bold; font-size: 1.5em; text-align: center;">
+      AVSS(Audio-Visual Speech Separation)
+    </td>
+  </tr>
+  <tr>
+    <td width="50%">
+      <video src="https://github.com/user-attachments/assets/13181ace-bb1e-4a6a-97b5-440caa1c93ef" controls preload></video>
+    </td>
+    <td width="50%">
+      <video src="https://github.com/user-attachments/assets/24a128fb-9686-4c48-955c-8f48c98847a8" controls preload></video>
+    </td>
+  </tr>
+</table>
+
+<table style="width:100%; text-align:center;">
+  <tr>
+    <td colspan="4" style="font-weight: bold; font-size: 1.5em; text-align: center;">
+      AVSE(Audio-Visual Speech Enhancement)
+    </td>
+  </tr>
+  <tr>
+    <td width="25%">
+      <video src="https://github.com/user-attachments/assets/bd7205e8-4eac-4f24-b5a3-251c35b35429" controls preload></video>
+    </td>
+    <td width="25%">
+      <video src="https://github.com/user-attachments/assets/3101da59-b535-43dc-b58f-8d62625a4b8b" controls preload></video>
+    </td>
+    <td width="25%">
+      <video src="https://github.com/user-attachments/assets/7f2011bf-ad67-4a67-b7b9-619e3bf04692" controls preload></video>
+    </td>
+    <td width="25%">
+      <video src="https://github.com/user-attachments/assets/e37e19d6-9a63-422b-b200-d827b4e9b317" controls preload></video>
+    </td>
+  </tr>
+</table>
+
+
+---
+## Result
+### CoGenAV Base for VSR/AVSR
+| Size        | SR Head        | Modalities | VSR  | AVSR@noise | AVSR@clean | AVSR with sft whisper @clean |
+|-------------|----------------|------------|------|------------|------------|------------|
+|     -        | Whisper medium  | A          | -    | 34.2       | 6.4        | 1.5                          |
+| **Base**    | Whisper small   | AV         | 24.8 | 5.2        | 2.5        | -                            |
+| **Large**   | Whisper medium  | AV         | 20.4 | 2.6        | 1.8        | **1.27**                     |
+> **Note:** VSR/AVSR results on LRS2. The evaluation metric used is WER, and the results are obtained from training conducted solely on the LRS2 dataset.
+
+### CoGenAV Base for AVSS/AVSE
+| Task        | SS Head       | Test Dataset    | SI-SNRi | SDRi | PESQ  |
+|-------------|----------------|------------------|---------|------|-------|
+| **AVSS**    | AV-Sepformer   | mix_2_spk_tt     | 15.7    | 16.0 | 3.23  |
+| **AVSE**    | AV-Sepformer   | lrs2_test+noise  | 8.3     | 9.0  | 2.56  |
+
+> **Note:** AVSS/AVSE results on LRS2. These metrics represent the average values for all speakers in each test set, where larger SI-SNRi, SDRi, and PESQ are better.
+
+### CoGenAV Base for ASD
+| Task        | SD Head       | Test Dataset    | mAP |
+|-------------|----------------|------------------|---------|
+| **ASD**    | LRASD   | Talkies     | 96.3    |