Upload 11 files

.gitignore

@@ -0,0 +1,3 @@
+venv/
+__pycache__/
+flagged/

README.md

@@ -1,6 +1,6 @@
 ---
 title: Japanese Ero Voice Classifier
-emoji: 🌍
+emoji: 🥰
 colorFrom: red
 colorTo: yellow
 sdk: gradio

app.py

@@ -0,0 +1,51 @@
+import argparse
+import json
+from pathlib import Path
+import gradio as gr
+import torch
+
+from models import AudioClassifier
+from utils import logger
+
+
+ckpt_dir = Path("ckpt/")
+config_path = ckpt_dir / "config.json"
+assert config_path.exists(), f"config.json not found in {ckpt_dir}"
+config = json.loads((ckpt_dir / "config.json").read_text())
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AudioClassifier(device=device, **config["model"]).to(device)
+
+# Latest checkpoint
+if (ckpt_dir / "model_final.pth").exists():
+    ckpt = ckpt_dir / "model_final.pth"
+else:
+    ckpt = sorted(ckpt_dir.glob("*.pth"))[-1]
+logger.info(f"Loading {ckpt}...")
+model.load_state_dict(torch.load(ckpt))
+
+
+def classify_audio(audio_file: str):
+    logger.info(f"Classifying {audio_file}...")
+    output = model.infer_from_file(audio_file)
+    logger.success(f"Predicted: {output}")
+    return output
+
+
+desc = """
+# NSFW音声分類器
+
+出力は以下の3つのクラスの確率です。
+- usual: 通常の音声
+- aegi: 喘ぎ声
+- chupa: チュパ音(フェラやキス音声)
+"""
+
+
+with gr.Interface(
+    fn=classify_audio,
+    inputs=gr.Audio(label="Input audio", type="filepath"),
+    outputs=gr.Text(label="Classification"),
+    description=desc,
+    allow_flagging="never",
+) as iface:
+    iface.launch()

ckpt/config.json

@@ -0,0 +1,13 @@
+{
+    "model": {
+        "label2id": {
+            "usual": 0,
+            "aegi": 1,
+            "chupa": 2
+        },
+        "num_hidden_layers": 2,
+        "hidden_dim": 128
+    },
+    "lr": 0.001,
+    "lr_decay": 0.996
+}

ckpt/model_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67ffab6e224d9c7f9acbeab40892cfda200a88c9dc2ee2714621bc90eed7a4d5
+size 279357

losses.py

@@ -0,0 +1,176 @@
+import torch
+import torch.nn as nn
+
+
+class AsymmetricLoss(nn.Module):
+    def __init__(
+        self,
+        gamma_neg=4,
+        gamma_pos=1,
+        clip=0.05,
+        eps=1e-8,
+        disable_torch_grad_focal_loss=True,
+    ):
+        super(AsymmetricLoss, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+
+    def forward(self, x, y):
+        """ "
+        Parameters
+        ----------
+        x: input logits
+        y: targets (multi-label binarized vector)
+        """
+
+        # Calculating Probabilities
+        x_sigmoid = torch.sigmoid(x)
+        xs_pos = x_sigmoid
+        xs_neg = 1 - x_sigmoid
+
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            xs_neg = (xs_neg + self.clip).clamp(max=1)
+
+        # Basic CE calculation
+        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
+        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
+        loss = los_pos + los_neg
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            pt0 = xs_pos * y
+            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
+            pt = pt0 + pt1
+            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
+            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            loss *= one_sided_w
+
+        return -loss.sum()
+
+
+class AsymmetricLossOptimized(nn.Module):
+    """Notice - optimized version, minimizes memory allocation and gpu uploading,
+    favors inplace operations"""
+
+    def __init__(
+        self,
+        gamma_neg=4,
+        gamma_pos=1,
+        clip=0.05,
+        eps=1e-8,
+        disable_torch_grad_focal_loss=False,
+    ):
+        super(AsymmetricLossOptimized, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+
+        # prevent memory allocation and gpu uploading every iteration, and encourages inplace operations
+        self.targets = self.anti_targets = self.xs_pos = self.xs_neg = (
+            self.asymmetric_w
+        ) = self.loss = None
+
+    def forward(self, x, y):
+        """ "
+        Parameters
+        ----------
+        x: input logits
+        y: targets (multi-label binarized vector)
+        """
+
+        self.targets = y
+        self.anti_targets = 1 - y
+
+        # Calculating Probabilities
+        self.xs_pos = torch.sigmoid(x)
+        self.xs_neg = 1.0 - self.xs_pos
+
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            self.xs_neg.add_(self.clip).clamp_(max=1)
+
+        # Basic CE calculation
+        self.loss = self.targets * torch.log(self.xs_pos.clamp(min=self.eps))
+        self.loss.add_(self.anti_targets * torch.log(self.xs_neg.clamp(min=self.eps)))
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            self.xs_pos = self.xs_pos * self.targets
+            self.xs_neg = self.xs_neg * self.anti_targets
+            self.asymmetric_w = torch.pow(
+                1 - self.xs_pos - self.xs_neg,
+                self.gamma_pos * self.targets + self.gamma_neg * self.anti_targets,
+            )
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            self.loss *= self.asymmetric_w
+
+        return -self.loss.sum()
+
+
+class ASLSingleLabel(nn.Module):
+    """
+    This loss is intended for single-label classification problems
+    """
+
+    def __init__(self, gamma_pos=0, gamma_neg=4, eps: float = 0.1, reduction="mean"):
+        super(ASLSingleLabel, self).__init__()
+
+        self.eps = eps
+        self.logsoftmax = nn.LogSoftmax(dim=-1)
+        self.targets_classes = []
+        self.gamma_pos = gamma_pos
+        self.gamma_neg = gamma_neg
+        self.reduction = reduction
+
+    def forward(self, inputs, target):
+        """
+        "input" dimensions: - (batch_size,number_classes)
+        "target" dimensions: - (batch_size)
+        """
+        num_classes = inputs.size()[-1]
+        log_preds = self.logsoftmax(inputs)
+        self.targets_classes = torch.zeros_like(inputs).scatter_(
+            1, target.long().unsqueeze(1), 1
+        )
+
+        # ASL weights
+        targets = self.targets_classes
+        anti_targets = 1 - targets
+        xs_pos = torch.exp(log_preds)
+        xs_neg = 1 - xs_pos
+        xs_pos = xs_pos * targets
+        xs_neg = xs_neg * anti_targets
+        asymmetric_w = torch.pow(
+            1 - xs_pos - xs_neg,
+            self.gamma_pos * targets + self.gamma_neg * anti_targets,
+        )
+        log_preds = log_preds * asymmetric_w
+
+        if self.eps > 0:  # label smoothing
+            self.targets_classes = self.targets_classes.mul(1 - self.eps).add(
+                self.eps / num_classes
+            )
+
+        # loss calculation
+        loss = -self.targets_classes.mul(log_preds)
+
+        loss = loss.sum(dim=-1)
+        if self.reduction == "mean":
+            loss = loss.mean()
+
+        return loss

models.py

@@ -0,0 +1,88 @@
+import torch
+from torch import nn
+
+
+# モデルの定義
+class AudioClassifier(nn.Module):
+    def __init__(
+        self,
+        label2id: dict,
+        feature_dim=256,
+        hidden_dim=256,
+        device="cpu",
+        dropout_rate=0.5,
+        num_hidden_layers=2,
+    ):
+        super(AudioClassifier, self).__init__()
+        self.num_classes = len(label2id)
+        self.device = device
+        self.label2id = label2id
+        self.id2label = {v: k for k, v in self.label2id.items()}
+        # 最初の線形層と活性化層を追加
+        self.fc1 = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim),
+            nn.Mish(),
+            nn.Dropout(dropout_rate),
+        )
+        # 隠れ層の追加
+        self.hidden_layers = nn.ModuleList()
+        for _ in range(num_hidden_layers):
+            layer = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.Mish(),
+                nn.Dropout(dropout_rate),
+            )
+            self.hidden_layers.append(layer)
+        # 最後の層（クラス分類用）
+        self.fc_last = nn.Linear(hidden_dim, self.num_classes)
+
+    def forward(self, x):
+        # 最初の層を通過
+        x = self.fc1(x)
+
+        # 隠れ層を順に通過
+        for layer in self.hidden_layers:
+            x = layer(x)
+
+        # 最後の分類層
+        x = self.fc_last(x)
+        return x
+
+    def infer_from_features(self, features):
+        # 特徴量をテンソルに変換
+        features = (
+            torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(self.device)
+        )
+
+        # モデルを評価モードに設定
+        self.eval()
+
+        # モデルの出力を取得
+        with torch.no_grad():
+            output = self.forward(features)
+
+        # ソフトマックス関数を適用して確率を計算
+        probs = torch.softmax(output, dim=1)
+
+        # ラベルごとの確率を計算して大きい順に並べ替えて返す
+        probs, indices = torch.sort(probs, descending=True)
+        probs = probs.cpu().numpy().squeeze()
+        indices = indices.cpu().numpy().squeeze()
+        return [(self.id2label[i], p) for i, p in zip(indices, probs)]
+
+    def infer_from_file(self, file_path):
+        feature = extract_features(file_path, device=self.device)
+        return self.infer_from_features(feature)
+
+
+from pyannote.audio import Inference, Model
+
+emb_model = Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM")
+inference = Inference(emb_model, window="whole")
+
+
+def extract_features(file_path, device="cpu"):
+    inference.to(torch.device(device))
+    return inference(file_path)

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+loguru
+pyannote.audio
+torch

train.py

@@ -0,0 +1,243 @@
+import argparse
+import json
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+# import torch_optimizer as optim
+import transformers
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    f1_score,
+    precision_score,
+    recall_score,
+)
+from torch.optim.lr_scheduler import (
+    CosineAnnealingLR,
+    CosineAnnealingWarmRestarts,
+    ExponentialLR,
+)
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from models import AudioClassifier, extract_features
+from losses import AsymmetricLoss, ASLSingleLabel
+
+torch.manual_seed(42)
+
+label2id = {
+    "usual": 0,
+    "aegi": 1,
+    "chupa": 2,
+    # "cry": 3,
+    # "laugh": 4,
+    # "silent": 5,
+    # "unusual": 6,
+}
+id2label = {v: k for k, v in label2id.items()}
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--exp_dir", type=str, default="data")
+parser.add_argument("--ckpt_dir", type=str, required=True)
+parser.add_argument("--device", type=str, default="cuda")
+parser.add_argument("--epochs", type=int, default=1000)
+parser.add_argument("--save_every", type=int, default=100)
+
+args = parser.parse_args()
+device = args.device
+if not torch.cuda.is_available():
+    print("No GPU detected. Using CPU.")
+    device = "cpu"
+print(f"Using {device} for training.")
+
+
+# データセットの定義
+class AudioDataset(Dataset):
+    def __init__(self, file_paths, labels, features):
+        self.file_paths = file_paths
+        self.labels = labels
+        self.features = features
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        return self.features[idx], self.labels[idx]
+
+
+def prepare_dataset(directory):
+    file_paths = list(Path(directory).rglob("*.npy"))
+    if len(file_paths) == 0:
+        return [], [], []
+    # file_paths = [f for f in file_paths if f.parent.name in label2id]
+
+    def process(file_path: Path):
+        npy_feature = np.load(file_path)
+        id = int(label2id[file_path.parent.name])
+        return (
+            file_path,
+            torch.tensor(id, dtype=torch.long).to(device),
+            torch.tensor(npy_feature, dtype=torch.float32).to(device),
+        )
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        results = list(tqdm(executor.map(process, file_paths), total=len(file_paths)))
+
+    file_paths, labels, features = zip(*results)
+
+    return file_paths, labels, features
+
+
+print("Preparing dataset...")
+
+exp_dir = Path(args.exp_dir)
+train_file_paths, train_labels, train_feats = prepare_dataset(exp_dir / "train")
+val_file_paths, val_labels, val_feats = prepare_dataset(exp_dir / "val")
+
+print(f"Train: {len(train_file_paths)}, Val: {len(val_file_paths)}")
+
+# データセットとデータローダーの準備
+train_dataset = AudioDataset(train_file_paths, train_labels, train_feats)
+print("Train dataset prepared.")
+train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+print("Train loader prepared.")
+if len(val_file_paths) == 0:
+    val_dataset = None
+    val_loader = None
+    print("No validation dataset found.")
+else:
+    val_dataset = AudioDataset(val_file_paths, val_labels, val_feats)
+    print("Val dataset prepared.")
+    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)
+    print("Val loader prepared.")
+
+
+# モデル、損失関数、最適化アルゴリズムの設定
+config = {
+    "model": {
+        "label2id": label2id,
+        "num_hidden_layers": 2,
+        "hidden_dim": 128,
+    },
+    "lr": 1e-3,
+    "lr_decay": 0.996,
+}
+model = AudioClassifier(device="cuda", **config["model"]).to(device)
+model.to(device)
+# criterion = nn.CrossEntropyLoss()
+criterion = ASLSingleLabel(gamma_pos=1, gamma_neg=4)
+optimizer = optim.AdamW(model.parameters(), lr=config["lr"], weight_decay=1e-2)
+scheduler = ExponentialLR(optimizer, gamma=config["lr_decay"])
+# scheduler = transformers.optimization.AdafactorSchedule(optimizer)
+num_epochs = args.epochs
+# scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
+
+print("Start training...")
+current_time = datetime.now().strftime("%b%d_%H-%M-%S")
+ckpt_dir = Path(args.ckpt_dir) / current_time
+ckpt_dir.mkdir(parents=True, exist_ok=True)
+# Save config
+with open(ckpt_dir / "config.json", "w", encoding="utf-8") as f:
+    json.dump(config, f, indent=4)
+# 訓練ループ
+save_every = args.save_every
+val_interval = 1
+eval_interval = 1
+
+writer = SummaryWriter(ckpt_dir / "logs")
+for epoch in tqdm(range(1, num_epochs + 1)):
+    train_loss = 0.0
+    model.train()  # 訓練モードに設定
+    train_labels = []
+    train_preds = []
+    for inputs, labels in train_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+
+        # 順伝播、損失の計算、逆伝播、パラメータ更新
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs.squeeze(), labels)
+        loss.backward()
+        optimizer.step()
+        train_loss += loss.item()
+
+        # 評価指標の計算
+        if epoch % eval_interval == 0:
+            with torch.no_grad():
+                # 最も高い確率を持つクラスのインデックスを取得
+                _, predictions = torch.max(outputs, 1)
+
+                # 実際のラベルと予測値をリストに追加
+                train_labels.extend(labels.cpu().numpy())
+                train_preds.extend(predictions.cpu().numpy())
+
+    scheduler.step()
+    if epoch % eval_interval == 0:
+        # 訓練データに対する評価指標の計算
+        accuracy = accuracy_score(train_labels, train_preds)
+        precision = precision_score(train_labels, train_preds, average="macro")
+        recall = recall_score(train_labels, train_preds, average="macro")
+        f1 = f1_score(train_labels, train_preds, average="macro")
+        report = classification_report(
+            train_labels, train_preds, target_names=list(label2id.keys())
+        )
+
+        writer.add_scalar("train/Accuracy", accuracy, epoch)
+        writer.add_scalar("train/Precision", precision, epoch)
+        writer.add_scalar("train/Recall", recall, epoch)
+        writer.add_scalar("train/F1", f1, epoch)
+
+    writer.add_scalar("Loss/train", train_loss / len(train_loader), epoch)
+    writer.add_scalar("Learning Rate", optimizer.param_groups[0]["lr"], epoch)
+
+    if epoch % save_every == 0:
+        torch.save(model.state_dict(), ckpt_dir / f"model_{epoch}.pth")
+
+    if epoch % val_interval != 0 or val_loader is None:
+        tqdm.write(f"loss: {train_loss / len(train_loader):4f}\n{report}")
+        continue
+    model.eval()  # 評価モードに設定
+    val_labels = []
+    val_preds = []
+    val_loss = 0.0
+    with torch.no_grad():
+        for inputs, labels in val_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs)
+            # 最も高い確率を持つクラスのインデックスを取得
+            _, predictions = torch.max(outputs, 1)
+            val_labels.extend(labels.cpu().numpy())
+            val_preds.extend(predictions.cpu().numpy())
+            loss = criterion(outputs.squeeze(), labels)
+            val_loss += loss.item()
+
+    # 評価指標の計算
+    accuracy = accuracy_score(val_labels, val_preds)
+    precision = precision_score(val_labels, val_preds, average="macro")
+    recall = recall_score(val_labels, val_preds, average="macro")
+    f1 = f1_score(val_labels, val_preds, average="macro")
+    report = classification_report(
+        val_labels, val_preds, target_names=list(label2id.keys())
+    )
+
+    writer.add_scalar("Loss/val", val_loss / len(val_loader), epoch)
+    writer.add_scalar("val/Accuracy", accuracy, epoch)
+    writer.add_scalar("val/Precision", precision, epoch)
+    writer.add_scalar("val/Recall", recall, epoch)
+    writer.add_scalar("val/F1", f1, epoch)
+
+    tqdm.write(
+        f"loss: {train_loss / len(train_loader):4f}, val loss: {val_loss / len(val_loader):4f}, "
+        f"acc: {accuracy:4f}, f1: {f1:4f}, prec: {precision:4f}, recall: {recall:4f}\n{report}"
+    )
+    # tqdm.write(report)
+    # Save
+torch.save(model.state_dict(), ckpt_dir / "model_final.pth")

utils.py

@@ -0,0 +1,29 @@
+import sys
+from pathlib import Path
+import wave
+from pydub import AudioSegment
+import loguru
+
+
+def is_audio_file(file: Path):
+    return file.suffix.lower() in [".wav", ".mp3", ".ogg"]
+
+
+def get_audio_duration_ms(file_path):
+    try:
+        with wave.open(str(file_path), "r") as wav_file:
+            return wav_file.getnframes() / wav_file.getframerate() * 1000
+    except wave.Error as e:
+        audio = AudioSegment.from_file(file_path)
+        return len(audio)
+    except Exception as e:
+        raise e
+
+
+logger = loguru.logger
+logger.remove()
+
+log_format = (
+    "<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}"
+)
+logger.add(sys.stdout, format=log_format)