Fix bug with empty RTTMs (#81)

* Update README.md * Update README.md * Fix bug with empty RTTMs * Add uri to empty annotation for consistency * Check non-empty annotation list in a more pythonic way Co-authored-by: Juan Coria <[email protected]>
juanmc2005 · Aug 17, 2022 · 6bbc226 · 6bbc226
1 parent a7befd7
commit 6bbc226
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/src/diart/inference.py b/src/diart/inference.py
@@ -68,7 +68,11 @@ def __call__(self, pipeline: OnlineSpeakerDiarization, source: src.AudioSource)
         # Stream audio through the pipeline
         source.read()
 
-        return load_rttm(rttm_path)[source.uri]
+        annotations = load_rttm(rttm_path)
+        if source.uri in annotations:
+            return annotations[source.uri]
+        else:
+            return Annotation(uri=source.uri)
 
 
 class Benchmark:

diff --git a/src/diart/sinks.py b/src/diart/sinks.py
@@ -24,9 +24,12 @@ def __init__(self, path: Union[Path, Text], patch_collar: float = 0.05):
 
     def patch_rttm(self):
         """Stitch same-speaker turns that are close to each other"""
-        annotation = list(load_rttm(self.path).values())[0]
-        with open(self.path, 'w') as file:
-            annotation.support(self.patch_collar).write_rttm(file)
+
+        annotations = list(load_rttm(self.path).values())
+        if annotations:
+            annotation = annotations[0]
+            with open(self.path, 'w') as file:
+                annotation.support(self.patch_collar).write_rttm(file)
 
     def on_next(self, value: Tuple[Annotation, SlidingWindowFeature]):
         with open(self.path, 'a') as file: