sad

nnet/sad_rnn.py

@@ -6,6 +6,7 @@ import torch
 import torch.nn as nn
 from torch import optim
 from torch.utils.data import Dataset
+from sidekit.frontend.io import read_hdf5_segment
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -52,17 +53,18 @@ class SAD_Dataset(Dataset):
             features, _ = features_server.load(show)
             labels = numpy.zeros((len(features), 1), dtype=numpy.int)
 
-            for seg in train_list[show]:
-                labels[seg['start']:seg['stop']] = 1
-            self.vad[show] = labels
-
-            for seg in uem_list[show]:
-                if seg['start'] is not None:
-                    start, stop = seg['start'], seg['stop']
-                else:
-                    start, stop = 0, len(features)
-                for i in range(start, min(stop, len(features)) - self.duration, self.step):
-                    self.segments.append((show, i, i + self.duration))
+            if show in train_list and show in uem_list:
+                for seg in train_list[show]:
+                    labels[seg['start']:seg['stop']] = 1
+                self.vad[show] = labels
+
+                for seg in uem_list[show]:
+                    if seg['start'] is not None:
+                        start, stop = seg['start'], seg['stop']
+                    else:
+                        start, stop = 0, len(features)
+                    for i in range(start, min(stop, len(features)) - self.duration, self.step):
+                        self.segments.append((show, i, i + self.duration))
         self.input_size = features.shape[1]
 
         if shuffle:
@@ -84,6 +86,98 @@ class SAD_Dataset(Dataset):
         return self.len
 
 
+
+class SAD_Dataset(Dataset):
+    """
+    Object that takes a list of files from a file and initialize a Dataset
+    """
+    def __init__(self, mdtm_file, features_file, batch_size=512, duration=3.2, step=0.8, uem_file=None, shuffle=False):
+
+        self.batch_size = batch_size
+        self.duration = int(duration * 100)
+        self.step = int(step * 100)
+        self.feature_file = open(feature_file, 'r')
+
+
+        train_list = {}
+        with open(mdtm_file, 'r') as f:
+            for line in f:
+                show, _, start, dur, _, _, _, _ = line.rstrip().split()
+                if show not in train_list:
+                    train_list[show] = []
+                train_list[show].append({
+                    "start": int(float(start) * 100),
+                    "stop": int((float(start) + float(dur)) * 100)})
+
+        uem_list = {}
+        if uem_file is not None:
+            with open(uem_file, 'r') as f:
+                for line in f:
+                    show, _, start, stop = line.rstrip().split()
+                    if show not in uem_list:
+                        uem_list[show] = []
+                    uem_list[show].append({
+                        "start": int(float(start) * 100),
+                        "stop": int(float(stop) * 100)})
+        else:
+            for show in train_list.keys():
+                uem_list[show].append({"start": None, "stop": None})
+
+        self.vad = {}
+        self.segments = []
+
+        for show in sorted(train_list.keys()):
+            #features, _ = features_server.load(show)
+            features = read_hdf5_segment(self.feature_file,
+                                         show,
+                                         ['energy', 'cep'],
+                                         label=None,
+                                         start=None, stop=None,
+                                         global_cmvn=False)[0]
+
+            labels = numpy.zeros((len(features), 1), dtype=numpy.int)
+
+            if show in train_list and show in uem_list:
+                for seg in train_list[show]:
+                    labels[seg['start']:seg['stop']] = 1
+                self.vad[show] = labels
+
+                for seg in uem_list[show]:
+                    if seg['start'] is not None:
+                        start, stop = seg['start'], seg['stop']
+                    else:
+                        start, stop = 0, len(features)
+                    for i in range(start, min(stop, len(features)) - self.duration, self.step):
+                        self.segments.append((show, i, i + self.duration))
+        self.input_size = features.shape[1]
+
+        if shuffle:
+            random.shuffle(self.segments)
+
+        self.len = len(self.segments) // self.batch_size
+
+    def __getitem__(self, index):
+        batch_X = numpy.zeros((self.batch_size, self.duration, self.input_size))
+        batch_Y = numpy.zeros((self.batch_size, self.duration, 1))
+        for i in range(self.batch_size):
+            show, start, stop = self.segments[index * self.batch_size + i]
+            #features, _ = self.features_server.load(show)
+            features = read_hdf5_segment(self.feature_file,
+                                         show,
+                                         ['energy', 'cep'],
+                                         label=None,
+                                         start=start, stop=stop,
+                                         global_cmvn=False)[0]
+            batch_X[i] = features
+            batch_Y[i] = self.vad[show][start:stop]
+            #batch_X[i] = features[start:stop]
+            #batch_Y[i] = self.vad[show][start:stop]
+        return torch.Tensor(batch_X), torch.Tensor(batch_Y)
+
+    def __len__(self):
+        return self.len
+
+
 class SAD_RNN():
     """
     A SAD_RNN is meant to use a PyTorch RNN model for Speech Activity Detection