first commit for takway.ai

This commit is contained in:
IrvingGao committed 2024-05-18 15:50:56 +08:00
commit d7a6de771c
3256 files changed
+795173

No files matched your search

View File
Whitespace-only changes.
+191
View File
@@ -0,0 +1,191 @@
#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
# MIT License (https://opensource.org/licenses/MIT)
# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
import scipy
import torch
import sklearn
import numpy as np
from sklearn.cluster._kmeans import k_means
from sklearn.cluster import HDBSCAN
class SpectralCluster:
r"""A spectral clustering mehtod using unnormalized Laplacian of affinity matrix.
This implementation is adapted from https://github.com/speechbrain/speechbrain.
"""
def __init__(self, min_num_spks=1, max_num_spks=15, pval=0.022):
self.min_num_spks = min_num_spks
self.max_num_spks = max_num_spks
self.pval = pval
def __call__(self, X, oracle_num=None):
# Similarity matrix computation
sim_mat = self.get_sim_mat(X)
# Refining similarity matrix with pval
prunned_sim_mat = self.p_pruning(sim_mat)
# Symmetrization
sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)
# Laplacian calculation
laplacian = self.get_laplacian(sym_prund_sim_mat)
# Get Spectral Embeddings
emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)
# Perform clustering
labels = self.cluster_embs(emb, num_of_spk)
return labels
def get_sim_mat(self, X):
# Cosine similarities
M = sklearn.metrics.pairwise.cosine_similarity(X, X)
return M
def p_pruning(self, A):
if A.shape[0] * self.pval < 6:
pval = 6.0 / A.shape[0]
else:
pval = self.pval
n_elems = int((1 - pval) * A.shape[0])
# For each row in a affinity matrix
for i in range(A.shape[0]):
low_indexes = np.argsort(A[i, :])
low_indexes = low_indexes[0:n_elems]
# Replace smaller similarity values by 0s
A[i, low_indexes] = 0
return A
def get_laplacian(self, M):
M[np.diag_indices(M.shape[0])] = 0
D = np.sum(np.abs(M), axis=1)
D = np.diag(D)
L = D - M
return L
def get_spec_embs(self, L, k_oracle=None):
lambdas, eig_vecs = scipy.linalg.eigh(L)
if k_oracle is not None:
num_of_spk = k_oracle
else:
lambda_gap_list = self.getEigenGaps(
lambdas[self.min_num_spks - 1 : self.max_num_spks + 1]
)
num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks
emb = eig_vecs[:, :num_of_spk]
return emb, num_of_spk
def cluster_embs(self, emb, k):
_, labels, _ = k_means(emb, k)
return labels
def getEigenGaps(self, eig_vals):
eig_vals_gap_list = []
for i in range(len(eig_vals) - 1):
gap = float(eig_vals[i + 1]) - float(eig_vals[i])
eig_vals_gap_list.append(gap)
return eig_vals_gap_list
class UmapHdbscan:
r"""
Reference:
- Siqi Zheng, Hongbin Suo. Reformulating Speaker Diarization as Community Detection With
Emphasis On Topological Structure. ICASSP2022
"""
def __init__(
self, n_neighbors=20, n_components=60, min_samples=10, min_cluster_size=10, metric="cosine"
):
self.n_neighbors = n_neighbors
self.n_components = n_components
self.min_samples = min_samples
self.min_cluster_size = min_cluster_size
self.metric = metric
def __call__(self, X):
import umap.umap_ as umap
umap_X = umap.UMAP(
n_neighbors=self.n_neighbors,
min_dist=0.0,
n_components=min(self.n_components, X.shape[0] - 2),
metric=self.metric,
).fit_transform(X)
labels = HDBSCAN(
min_samples=self.min_samples,
min_cluster_size=self.min_cluster_size,
allow_single_cluster=True,
).fit_predict(umap_X)
return labels
class ClusterBackend(torch.nn.Module):
r"""Perfom clustering for input embeddings and output the labels.
Args:
model_dir: A model dir.
model_config: The model config.
"""
def __init__(self):
super().__init__()
self.model_config = {"merge_thr": 0.78}
# self.other_config = kwargs
self.spectral_cluster = SpectralCluster()
self.umap_hdbscan_cluster = UmapHdbscan()
def forward(self, X, **params):
# clustering and return the labels
k = params["oracle_num"] if "oracle_num" in params else None
assert len(X.shape) == 2, "modelscope error: the shape of input should be [N, C]"
if X.shape[0] < 20:
return np.zeros(X.shape[0], dtype="int")
if X.shape[0] < 2048 or k is not None:
# unexpected corner case
labels = self.spectral_cluster(X, k)
else:
labels = self.umap_hdbscan_cluster(X)
if k is None and "merge_thr" in self.model_config:
labels = self.merge_by_cos(labels, X, self.model_config["merge_thr"])
return labels
def merge_by_cos(self, labels, embs, cos_thr):
# merge the similar speakers by cosine similarity
assert cos_thr > 0 and cos_thr <= 1
while True:
spk_num = labels.max() + 1
if spk_num == 1:
break
spk_center = []
for i in range(spk_num):
spk_emb = embs[labels == i].mean(0)
spk_center.append(spk_emb)
assert len(spk_center) > 0
spk_center = np.stack(spk_center, axis=0)
norm_spk_center = spk_center / np.linalg.norm(spk_center, axis=1, keepdims=True)
affinity = np.matmul(norm_spk_center, norm_spk_center.T)
affinity = np.triu(affinity, 1)
spks = np.unravel_index(np.argmax(affinity), affinity.shape)
if affinity[spks] < cos_thr:
break
for i in range(len(labels)):
if labels[i] == spks[1]:
labels[i] = spks[0]
elif labels[i] > spks[1]:
labels[i] -= 1
return labels
+288
View File
@@ -0,0 +1,288 @@
#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
# MIT License (https://opensource.org/licenses/MIT)
# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
import torch
import torch.nn.functional as F
import torch.utils.checkpoint as cp
class BasicResBlock(torch.nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1):
super(BasicResBlock, self).__init__()
self.conv1 = torch.nn.Conv2d(
in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False
)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = torch.nn.BatchNorm2d(planes)
self.shortcut = torch.nn.Sequential()
if stride != 1 or in_planes != self.expansion * planes:
self.shortcut = torch.nn.Sequential(
torch.nn.Conv2d(
in_planes,
self.expansion * planes,
kernel_size=1,
stride=(stride, 1),
bias=False,
),
torch.nn.BatchNorm2d(self.expansion * planes),
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class FCM(torch.nn.Module):
def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
super(FCM, self).__init__()
self.in_planes = m_channels
self.conv1 = torch.nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = torch.nn.BatchNorm2d(m_channels)
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
self.conv2 = torch.nn.Conv2d(
m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False
)
self.bn2 = torch.nn.BatchNorm2d(m_channels)
self.out_channels = m_channels * (feat_dim // 8)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return torch.nn.Sequential(*layers)
def forward(self, x):
x = x.unsqueeze(1)
out = F.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = F.relu(self.bn2(self.conv2(out)))
shape = out.shape
out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
return out
def get_nonlinear(config_str, channels):
nonlinear = torch.nn.Sequential()
for name in config_str.split("-"):
if name == "relu":
nonlinear.add_module("relu", torch.nn.ReLU(inplace=True))
elif name == "prelu":
nonlinear.add_module("prelu", torch.nn.PReLU(channels))
elif name == "batchnorm":
nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels))
elif name == "batchnorm_":
nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels, affine=False))
else:
raise ValueError("Unexpected module ({}).".format(name))
return nonlinear
def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
mean = x.mean(dim=dim)
std = x.std(dim=dim, unbiased=unbiased)
stats = torch.cat([mean, std], dim=-1)
if keepdim:
stats = stats.unsqueeze(dim=dim)
return stats
class StatsPool(torch.nn.Module):
def forward(self, x):
return statistics_pooling(x)
class TDNNLayer(torch.nn.Module):
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
bias=False,
config_str="batchnorm-relu",
):
super(TDNNLayer, self).__init__()
if padding < 0:
assert (
kernel_size % 2 == 1
), "Expect equal paddings, but got even kernel size ({})".format(kernel_size)
padding = (kernel_size - 1) // 2 * dilation
self.linear = torch.nn.Conv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=bias,
)
self.nonlinear = get_nonlinear(config_str, out_channels)
def forward(self, x):
x = self.linear(x)
x = self.nonlinear(x)
return x
class CAMLayer(torch.nn.Module):
def __init__(
self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2
):
super(CAMLayer, self).__init__()
self.linear_local = torch.nn.Conv1d(
bn_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=bias,
)
self.linear1 = torch.nn.Conv1d(bn_channels, bn_channels // reduction, 1)
self.relu = torch.nn.ReLU(inplace=True)
self.linear2 = torch.nn.Conv1d(bn_channels // reduction, out_channels, 1)
self.sigmoid = torch.nn.Sigmoid()
def forward(self, x):
y = self.linear_local(x)
context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
context = self.relu(self.linear1(context))
m = self.sigmoid(self.linear2(context))
return y * m
def seg_pooling(self, x, seg_len=100, stype="avg"):
if stype == "avg":
seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
elif stype == "max":
seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
else:
raise ValueError("Wrong segment pooling type.")
shape = seg.shape
seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
seg = seg[..., : x.shape[-1]]
return seg
class CAMDenseTDNNLayer(torch.nn.Module):
def __init__(
self,
in_channels,
out_channels,
bn_channels,
kernel_size,
stride=1,
dilation=1,
bias=False,
config_str="batchnorm-relu",
memory_efficient=False,
):
super(CAMDenseTDNNLayer, self).__init__()
assert kernel_size % 2 == 1, "Expect equal paddings, but got even kernel size ({})".format(
kernel_size
)
padding = (kernel_size - 1) // 2 * dilation
self.memory_efficient = memory_efficient
self.nonlinear1 = get_nonlinear(config_str, in_channels)
self.linear1 = torch.nn.Conv1d(in_channels, bn_channels, 1, bias=False)
self.nonlinear2 = get_nonlinear(config_str, bn_channels)
self.cam_layer = CAMLayer(
bn_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=bias,
)
def bn_function(self, x):
return self.linear1(self.nonlinear1(x))
def forward(self, x):
if self.training and self.memory_efficient:
x = cp.checkpoint(self.bn_function, x)
else:
x = self.bn_function(x)
x = self.cam_layer(self.nonlinear2(x))
return x
class CAMDenseTDNNBlock(torch.nn.ModuleList):
def __init__(
self,
num_layers,
in_channels,
out_channels,
bn_channels,
kernel_size,
stride=1,
dilation=1,
bias=False,
config_str="batchnorm-relu",
memory_efficient=False,
):
super(CAMDenseTDNNBlock, self).__init__()
for i in range(num_layers):
layer = CAMDenseTDNNLayer(
in_channels=in_channels + i * out_channels,
out_channels=out_channels,
bn_channels=bn_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
bias=bias,
config_str=config_str,
memory_efficient=memory_efficient,
)
self.add_module("tdnnd%d" % (i + 1), layer)
def forward(self, x):
for layer in self:
x = torch.cat([x, layer(x)], dim=1)
return x
class TransitLayer(torch.nn.Module):
def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
super(TransitLayer, self).__init__()
self.nonlinear = get_nonlinear(config_str, in_channels)
self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
def forward(self, x):
x = self.nonlinear(x)
x = self.linear(x)
return x
class DenseLayer(torch.nn.Module):
def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
super(DenseLayer, self).__init__()
self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
self.nonlinear = get_nonlinear(config_str, out_channels)
def forward(self, x):
if len(x.shape) == 2:
x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
else:
x = self.linear(x)
x = self.nonlinear(x)
return x
+146
View File
@@ -0,0 +1,146 @@
#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
# MIT License (https://opensource.org/licenses/MIT)
# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
import time
import torch
import numpy as np
from collections import OrderedDict
from contextlib import contextmanager
from distutils.version import LooseVersion
from funasr.register import tables
from funasr.models.campplus.utils import extract_feature
from funasr.utils.load_utils import load_audio_text_image_video
from funasr.models.campplus.components import (
DenseLayer,
StatsPool,
TDNNLayer,
CAMDenseTDNNBlock,
TransitLayer,
get_nonlinear,
FCM,
)
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
from torch.cuda.amp import autocast
else:
# Nothing to do if torch<1.6.0
@contextmanager
def autocast(enabled=True):
yield
@tables.register("model_classes", "CAMPPlus")
class CAMPPlus(torch.nn.Module):
def __init__(
self,
feat_dim=80,
embedding_size=192,
growth_rate=32,
bn_size=4,
init_channels=128,
config_str="batchnorm-relu",
memory_efficient=True,
output_level="segment",
**kwargs,
):
super().__init__()
self.head = FCM(feat_dim=feat_dim)
channels = self.head.out_channels
self.output_level = output_level
self.xvector = torch.nn.Sequential(
OrderedDict(
[
(
"tdnn",
TDNNLayer(
channels,
init_channels,
5,
stride=2,
dilation=1,
padding=-1,
config_str=config_str,
),
),
]
)
)
channels = init_channels
for i, (num_layers, kernel_size, dilation) in enumerate(
zip((12, 24, 16), (3, 3, 3), (1, 2, 2))
):
block = CAMDenseTDNNBlock(
num_layers=num_layers,
in_channels=channels,
out_channels=growth_rate,
bn_channels=bn_size * growth_rate,
kernel_size=kernel_size,
dilation=dilation,
config_str=config_str,
memory_efficient=memory_efficient,
)
self.xvector.add_module("block%d" % (i + 1), block)
channels = channels + num_layers * growth_rate
self.xvector.add_module(
"transit%d" % (i + 1),
TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
)
channels //= 2
self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
if self.output_level == "segment":
self.xvector.add_module("stats", StatsPool())
self.xvector.add_module(
"dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")
)
else:
assert (
self.output_level == "frame"
), "`output_level` should be set to 'segment' or 'frame'. "
for m in self.modules():
if isinstance(m, (torch.nn.Conv1d, torch.nn.Linear)):
torch.nn.init.kaiming_normal_(m.weight.data)
if m.bias is not None:
torch.nn.init.zeros_(m.bias)
def forward(self, x):
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
x = self.head(x)
x = self.xvector(x)
if self.output_level == "frame":
x = x.transpose(1, 2)
return x
def inference(
self,
data_in,
data_lengths=None,
key: list = None,
tokenizer=None,
frontend=None,
**kwargs,
):
# extract fbank feats
meta_data = {}
time1 = time.perf_counter()
audio_sample_list = load_audio_text_image_video(
data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound"
)
time2 = time.perf_counter()
meta_data["load_data"] = f"{time2 - time1:0.3f}"
speech, speech_lengths, speech_times = extract_feature(audio_sample_list)
speech = speech.to(device=kwargs["device"])
time3 = time.perf_counter()
meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
meta_data["batch_data_time"] = np.array(speech_times).sum().item() / 16000.0
results = [{"spk_embedding": self.forward(speech.to(torch.float32))}]
return results, meta_data
+23
View File
@@ -0,0 +1,23 @@
# This is an example that demonstrates how to configure a model file.
# You can modify the configuration according to your own requirements.
# to print the register_table:
# from funasr.register import tables
# tables.print()
# network architecture
model: CAMPPlus
model_conf:
feat_dim: 80
embedding_size: 192
growth_rate: 32
bn_size: 4
init_channels: 128
config_str: 'batchnorm-relu'
memory_efficient: True
output_level: 'segment'
# frontend related
frontend: WavFrontend
frontend_conf:
fs: 16000
+504
View File
@@ -0,0 +1,504 @@
#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
# MIT License (https://opensource.org/licenses/MIT)
# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
import io
import os
import torch
import requests
import tempfile
import contextlib
import numpy as np
import librosa as sf
from typing import Union
from pathlib import Path
from typing import Generator, Union
from abc import ABCMeta, abstractmethod
import torchaudio.compliance.kaldi as Kaldi
from funasr.models.transformer.utils.nets_utils import pad_list
def check_audio_list(audio: list):
audio_dur = 0
for i in range(len(audio)):
seg = audio[i]
assert seg[1] >= seg[0], "modelscope error: Wrong time stamps."
assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type."
assert (
int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0]
), "modelscope error: audio data in list is inconsistent with time length."
if i > 0:
assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps."
audio_dur += seg[1] - seg[0]
return audio_dur
# assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.'
def sv_preprocess(inputs: Union[np.ndarray, list]):
output = []
for i in range(len(inputs)):
if isinstance(inputs[i], str):
file_bytes = File.read(inputs[i])
data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32")
if len(data.shape) == 2:
data = data[:, 0]
data = torch.from_numpy(data).unsqueeze(0)
data = data.squeeze(0)
elif isinstance(inputs[i], np.ndarray):
assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]"
data = inputs[i]
if data.dtype in ["int16", "int32", "int64"]:
data = (data / (1 << 15)).astype("float32")
else:
data = data.astype("float32")
data = torch.from_numpy(data)
else:
raise ValueError(
"modelscope error: The input type is restricted to audio address and nump array."
)
output.append(data)
return output
def sv_chunk(vad_segments: list, fs=16000) -> list:
config = {
"seg_dur": 1.5,
"seg_shift": 0.75,
}
def seg_chunk(seg_data):
seg_st = seg_data[0]
data = seg_data[2]
chunk_len = int(config["seg_dur"] * fs)
chunk_shift = int(config["seg_shift"] * fs)
last_chunk_ed = 0
seg_res = []
for chunk_st in range(0, data.shape[0], chunk_shift):
chunk_ed = min(chunk_st + chunk_len, data.shape[0])
if chunk_ed <= last_chunk_ed:
break
last_chunk_ed = chunk_ed
chunk_st = max(0, chunk_ed - chunk_len)
chunk_data = data[chunk_st:chunk_ed]
if chunk_data.shape[0] < chunk_len:
chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant")
seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data])
return seg_res
segs = []
for i, s in enumerate(vad_segments):
segs.extend(seg_chunk(s))
return segs
def extract_feature(audio):
features = []
feature_times = []
feature_lengths = []
for au in audio:
feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80)
feature = feature - feature.mean(dim=0, keepdim=True)
features.append(feature)
feature_times.append(au.shape[0])
feature_lengths.append(feature.shape[0])
# padding for batch inference
features_padded = pad_list(features, pad_value=0)
# features = torch.cat(features)
return features_padded, feature_lengths, feature_times
def postprocess(
segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray
) -> list:
assert len(segments) == len(labels)
labels = correct_labels(labels)
distribute_res = []
for i in range(len(segments)):
distribute_res.append([segments[i][0], segments[i][1], labels[i]])
# merge the same speakers chronologically
distribute_res = merge_seque(distribute_res)
# accquire speaker center
spk_embs = []
for i in range(labels.max() + 1):
spk_emb = embeddings[labels == i].mean(0)
spk_embs.append(spk_emb)
spk_embs = np.stack(spk_embs)
def is_overlapped(t1, t2):
if t1 > t2 + 1e-4:
return True
return False
# distribute the overlap region
for i in range(1, len(distribute_res)):
if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]):
p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2
distribute_res[i][0] = p
distribute_res[i - 1][1] = p
# smooth the result
distribute_res = smooth(distribute_res)
return distribute_res
def correct_labels(labels):
labels_id = 0
id2id = {}
new_labels = []
for i in labels:
if i not in id2id:
id2id[i] = labels_id
labels_id += 1
new_labels.append(id2id[i])
return np.array(new_labels)
def merge_seque(distribute_res):
res = [distribute_res[0]]
for i in range(1, len(distribute_res)):
if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]:
res.append(distribute_res[i])
else:
res[-1][1] = distribute_res[i][1]
return res
def smooth(res, mindur=1):
# if only one segment, return directly
if len(res) < 2:
return res
# short segments are assigned to nearest speakers.
for i in range(len(res)):
res[i][0] = round(res[i][0], 2)
res[i][1] = round(res[i][1], 2)
if res[i][1] - res[i][0] < mindur:
if i == 0:
res[i][2] = res[i + 1][2]
elif i == len(res) - 1:
res[i][2] = res[i - 1][2]
elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]:
res[i][2] = res[i - 1][2]
else:
res[i][2] = res[i + 1][2]
# merge the speakers
res = merge_seque(res)
return res
def distribute_spk(sentence_list, sd_time_list):
sd_sentence_list = []
for d in sentence_list:
sentence_start = d["start"]
sentence_end = d["end"]
sentence_spk = 0
max_overlap = 0
for sd_time in sd_time_list:
spk_st, spk_ed, spk = sd_time
spk_st = spk_st * 1000
spk_ed = spk_ed * 1000
overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0)
if overlap > max_overlap:
max_overlap = overlap
sentence_spk = spk
d["spk"] = int(sentence_spk)
sd_sentence_list.append(d)
return sd_sentence_list
class Storage(metaclass=ABCMeta):
"""Abstract class of storage.
All backends need to implement two apis: ``read()`` and ``read_text()``.
``read()`` reads the file as a byte stream and ``read_text()`` reads
the file as texts.
"""
@abstractmethod
def read(self, filepath: str):
pass
@abstractmethod
def read_text(self, filepath: str):
pass
@abstractmethod
def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
pass
@abstractmethod
def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
pass
class LocalStorage(Storage):
"""Local hard disk storage"""
def read(self, filepath: Union[str, Path]) -> bytes:
"""Read data from a given ``filepath`` with 'rb' mode.
Args:
filepath (str or Path): Path to read data.
Returns:
bytes: Expected bytes object.
"""
with open(filepath, "rb") as f:
content = f.read()
return content
def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str:
"""Read data from a given ``filepath`` with 'r' mode.
Args:
filepath (str or Path): Path to read data.
encoding (str): The encoding format used to open the ``filepath``.
Default: 'utf-8'.
Returns:
str: Expected text reading from ``filepath``.
"""
with open(filepath, "r", encoding=encoding) as f:
value_buf = f.read()
return value_buf
def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
"""Write data to a given ``filepath`` with 'wb' mode.
Note:
``write`` will create a directory if the directory of ``filepath``
does not exist.
Args:
obj (bytes): Data to be written.
filepath (str or Path): Path to write data.
"""
dirname = os.path.dirname(filepath)
if dirname and not os.path.exists(dirname):
os.makedirs(dirname, exist_ok=True)
with open(filepath, "wb") as f:
f.write(obj)
def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
"""Write data to a given ``filepath`` with 'w' mode.
Note:
``write_text`` will create a directory if the directory of
``filepath`` does not exist.
Args:
obj (str): Data to be written.
filepath (str or Path): Path to write data.
encoding (str): The encoding format used to open the ``filepath``.
Default: 'utf-8'.
"""
dirname = os.path.dirname(filepath)
if dirname and not os.path.exists(dirname):
os.makedirs(dirname, exist_ok=True)
with open(filepath, "w", encoding=encoding) as f:
f.write(obj)
@contextlib.contextmanager
def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]:
"""Only for unified API and do nothing."""
yield filepath
class HTTPStorage(Storage):
"""HTTP and HTTPS storage."""
def read(self, url):
# TODO @wenmeng.zwm add progress bar if file is too large
r = requests.get(url)
r.raise_for_status()
return r.content
def read_text(self, url):
r = requests.get(url)
r.raise_for_status()
return r.text
@contextlib.contextmanager
def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
"""Download a file from ``filepath``.
``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
can be called with ``with`` statement, and when exists from the
``with`` statement, the temporary path will be released.
Args:
filepath (str): Download a file from ``filepath``.
Examples:
>>> storage = HTTPStorage()
>>> # After existing from the ``with`` clause,
>>> # the path will be removed
>>> with storage.get_local_path('http://path/to/file') as path:
... # do something here
"""
try:
f = tempfile.NamedTemporaryFile(delete=False)
f.write(self.read(filepath))
f.close()
yield f.name
finally:
os.remove(f.name)
def write(self, obj: bytes, url: Union[str, Path]) -> None:
raise NotImplementedError("write is not supported by HTTP Storage")
def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None:
raise NotImplementedError("write_text is not supported by HTTP Storage")
class OSSStorage(Storage):
"""OSS storage."""
def __init__(self, oss_config_file=None):
# read from config file or env var
raise NotImplementedError("OSSStorage.__init__ to be implemented in the future")
def read(self, filepath):
raise NotImplementedError("OSSStorage.read to be implemented in the future")
def read_text(self, filepath, encoding="utf-8"):
raise NotImplementedError("OSSStorage.read_text to be implemented in the future")
@contextlib.contextmanager
def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
"""Download a file from ``filepath``.
``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
can be called with ``with`` statement, and when exists from the
``with`` statement, the temporary path will be released.
Args:
filepath (str): Download a file from ``filepath``.
Examples:
>>> storage = OSSStorage()
>>> # After existing from the ``with`` clause,
>>> # the path will be removed
>>> with storage.get_local_path('http://path/to/file') as path:
... # do something here
"""
try:
f = tempfile.NamedTemporaryFile(delete=False)
f.write(self.read(filepath))
f.close()
yield f.name
finally:
os.remove(f.name)
def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
raise NotImplementedError("OSSStorage.write to be implemented in the future")
def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
raise NotImplementedError("OSSStorage.write_text to be implemented in the future")
G_STORAGES = {}
class File(object):
_prefix_to_storage: dict = {
"oss": OSSStorage,
"http": HTTPStorage,
"https": HTTPStorage,
"local": LocalStorage,
}
@staticmethod
def _get_storage(uri):
assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}"
if "://" not in uri:
# local path
storage_type = "local"
else:
prefix, _ = uri.split("://")
storage_type = prefix
assert storage_type in File._prefix_to_storage, (
f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}"
)
if storage_type not in G_STORAGES:
G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]()
return G_STORAGES[storage_type]
@staticmethod
def read(uri: str) -> bytes:
"""Read data from a given ``filepath`` with 'rb' mode.
Args:
filepath (str or Path): Path to read data.
Returns:
bytes: Expected bytes object.
"""
storage = File._get_storage(uri)
return storage.read(uri)
@staticmethod
def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str:
"""Read data from a given ``filepath`` with 'r' mode.
Args:
filepath (str or Path): Path to read data.
encoding (str): The encoding format used to open the ``filepath``.
Default: 'utf-8'.
Returns:
str: Expected text reading from ``filepath``.
"""
storage = File._get_storage(uri)
return storage.read_text(uri)
@staticmethod
def write(obj: bytes, uri: Union[str, Path]) -> None:
"""Write data to a given ``filepath`` with 'wb' mode.
Note:
``write`` will create a directory if the directory of ``filepath``
does not exist.
Args:
obj (bytes): Data to be written.
filepath (str or Path): Path to write data.
"""
storage = File._get_storage(uri)
return storage.write(obj, uri)
@staticmethod
def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None:
"""Write data to a given ``filepath`` with 'w' mode.
Note:
``write_text`` will create a directory if the directory of
``filepath`` does not exist.
Args:
obj (str): Data to be written.
filepath (str or Path): Path to write data.
encoding (str): The encoding format used to open the ``filepath``.
Default: 'utf-8'.
"""
storage = File._get_storage(uri)
return storage.write_text(obj, uri)
@contextlib.contextmanager
def as_local_path(uri: str) -> Generator[Union[str, Path], None, None]:
"""Only for unified API and do nothing."""
storage = File._get_storage(uri)
with storage.as_local_path(uri) as local_path:
yield local_path