From a2ca25f26dd52b5eaf315d66261c6beaea101f83 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian-Robert=20St=C3=B6ter?= <mail@faroit.com>
Date: Tue, 27 Jun 2023 22:31:21 +0200
Subject: [PATCH] change preview location, more modern downloader

---
 musdb/__init__.py      | 56 ++++++++++++++++++++--------
 musdb/configs/mus.yaml | 85 +++++++++++++++++++++---------------------
 2 files changed, 83 insertions(+), 58 deletions(-)

diff --git a/musdb/__init__.py b/musdb/__init__.py
index 0ad41c2..93776ea 100644
--- a/musdb/__init__.py
+++ b/musdb/__init__.py
@@ -1,16 +1,18 @@
 from .audio_classes import MultiTrack, Source, Target
 from os import path as op
+from pathlib import Path
 import stempeg
-import urllib.request
+from urllib.request import urlopen, Request
 import collections
 import numpy as np
-import functools
+from tqdm import tqdm
 import zipfile
 import yaml
-import musdb
 import errno
+import musdb
+import shutil
 import os
-
+import tempfile
 
 class DB(object):
     """
@@ -351,7 +353,7 @@ def save_estimates(
     def _check_exists(self):
         return os.path.exists(os.path.join(self.root, "train"))
 
-    def download(self):
+    def download(self, progress: bool = True, suffix: str = ".zip"):
         """Download the MUSDB Sample data"""
         if self._check_exists():
             return
@@ -364,16 +366,40 @@ def download(self):
                 pass
             else:
                 raise
-
         print('Downloading MUSDB 7s Sample Dataset to %s...' % self.root)
-        data = urllib.request.urlopen(self.url)
-        filename = 'MUSDB18-7-STEMS.zip'
-        file_path = os.path.join(self.root, filename)
-        with open(file_path, 'wb') as f:
-            f.write(data.read())
-        zip_ref = zipfile.ZipFile(file_path, 'r')
-        zip_ref.extractall(os.path.join(self.root))
-        zip_ref.close()
-        os.unlink(file_path)
+
+        file_size = None
+        req = Request(self.url, headers={"User-Agent": "musdb_downloader"})
+        u = urlopen(req)
+        meta = u.info()
+        if hasattr(meta, 'getheaders'):
+            content_length = meta.getheaders("Content-Length")
+        else:
+            content_length = meta.get_all("Content-Length")
+        if content_length is not None and len(content_length) > 0:
+            file_size = int(content_length[0])
+
+        # We deliberately save it in a temp file and move it after
+        # download is complete. This prevents a local working checkpoint
+        # being overridden by a broken download.
+        f = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+
+        try:
+            with tqdm(total=file_size, disable=not progress, unit='B', unit_scale=True, unit_divisor=1024) as pbar:
+                while True:
+                    buffer = u.read(8192)
+                    if len(buffer) == 0:
+                        break
+                    f.write(buffer)
+                    pbar.update(len(buffer))
+
+            f.close()
+            zip_ref = zipfile.ZipFile(f.name, 'r')
+            zip_ref.extractall(os.path.join(self.root))
+            zip_ref.close()
+        finally:
+            f.close()
+            if os.path.exists(f.name):
+                os.remove(f.name)
 
         print('Done!')
diff --git a/musdb/configs/mus.yaml b/musdb/configs/mus.yaml
index ed180c5..07e5d1a 100644
--- a/musdb/configs/mus.yaml
+++ b/musdb/configs/mus.yaml
@@ -1,58 +1,57 @@
-sample-url: https://github.com/sigsep/sigsep-mus-db/releases/download/v0.3.0/mus-sample.zip
+sample-url: https://github.com/sigsep/sigsep-mus-db/releases/download/v0.4.0/MUSDB18-7-STEMS.zip
 
 # global dataset samplerate
 sample_rate: 44100.0
 
 # Define the Input sources for each track
 sources:
-    vocals: vocals.wav
-    drums: drums.wav
-    bass: bass.wav
-    other: other.wav
+  vocals: vocals.wav
+  drums: drums.wav
+  bass: bass.wav
+  other: other.wav
 
 # Additionally provide the artistic mix
-mixture:
-    mixture.wav
+mixture: mixture.wav
 
 stem_ids:
-    mixture: 0
-    drums: 1
-    bass: 2
-    other: 3
-    vocals: 4
+  mixture: 0
+  drums: 1
+  bass: 2
+  other: 3
+  vocals: 4
 
 # Define output targets to be tested using linear mixtures
 targets:
-    vocals:
-        vocals: 1
-    drums:
-        drums: 1
-    bass:
-        bass: 1
-    other:
-        other: 1
-    accompaniment:
-        bass: 1
-        drums: 1
-        other: 1
-    linear_mixture:
-        vocals: 1
-        bass: 1
-        drums: 1
-        other: 1
+  vocals:
+    vocals: 1
+  drums:
+    drums: 1
+  bass:
+    bass: 1
+  other:
+    other: 1
+  accompaniment:
+    bass: 1
+    drums: 1
+    other: 1
+  linear_mixture:
+    vocals: 1
+    bass: 1
+    drums: 1
+    other: 1
 
 validation_tracks:
-    - 'Actions - One Minute Smile'
-    - 'Clara Berry And Wooldog - Waltz For My Victims'
-    - 'Johnny Lokke - Promises & Lies'
-    - 'Patrick Talbot - A Reason To Leave'
-    - 'Triviul - Angelsaint'
-    - 'Alexander Ross - Goodbye Bolero'
-    - 'Fergessen - Nos Palpitants'
-    - 'Leaf - Summerghost'
-    - 'Skelpolu - Human Mistakes'
-    - 'Young Griffo - Pennies'
-    - 'ANiMAL - Rockshow'
-    - 'James May - On The Line'
-    - 'Meaxic - Take A Step'
-    - 'Traffic Experiment - Sirens'
+- 'Actions - One Minute Smile'
+- 'Clara Berry And Wooldog - Waltz For My Victims'
+- 'Johnny Lokke - Promises & Lies'
+- 'Patrick Talbot - A Reason To Leave'
+- 'Triviul - Angelsaint'
+- 'Alexander Ross - Goodbye Bolero'
+- 'Fergessen - Nos Palpitants'
+- 'Leaf - Summerghost'
+- 'Skelpolu - Human Mistakes'
+- 'Young Griffo - Pennies'
+- 'ANiMAL - Rockshow'
+- 'James May - On The Line'
+- 'Meaxic - Take A Step'
+- 'Traffic Experiment - Sirens'