Additional fix for dataset download / reporting dataset size

votchallenge · May 16, 2023 · 42d8cac · 42d8cac
1 parent ce19ee2
commit 42d8cac
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/vot/dataset/vot.py b/vot/dataset/vot.py
@@ -165,10 +165,17 @@ def list(self):
 
 def download_dataset_meta(url, path):
     from vot.utilities.net import download_uncompress, download_json, get_base_url, join_url, NetworkException
+    from vot.utilities import format_size
 
     meta = download_json(url)
 
-    logger.info('Downloading sequence dataset "%s" with %s sequences.', meta["name"], len(meta["sequences"]))
+    total_size = 0
+    for sequence in meta["sequences"]:
+        total_size += sequence["annotations"]["uncompressed"]
+        for channel in sequence["channels"].values():
+            total_size += channel["uncompressed"]
+
+    logger.info('Downloading sequence dataset "%s" with %s sequences (total %s).', meta["name"], len(meta["sequences"]), format_size(total_size))
 
     base_url = get_base_url(url) + "/"
 
@@ -203,6 +210,8 @@ def download_dataset_meta(url, path):
                 failed.append(sequence["name"])
                 continue
 
+            failure = False
+
             for cname, channel in sequence["channels"].items():
                 channel_directory = os.path.join(sequence_directory, cname)
                 os.makedirs(channel_directory, exist_ok=True)
@@ -214,17 +223,20 @@ def download_dataset_meta(url, path):
                 except NetworkException as e:
                     logger.exception(e)
                     failed.append(sequence["name"])
-                    continue
+                    failure = False
                 except IOError as e:
                     logger.exception(e)
                     failed.append(sequence["name"])
-                    continue
+                    failure = False
 
                 if "pattern" in channel:
                     data["channels." + cname] = cname + os.path.sep + channel["pattern"]
                 else:
                     data["channels." + cname] = cname + os.path.sep
 
+                if failure:
+                    continue
+
             write_properties(os.path.join(sequence_directory, 'sequence'), data)
             progress.relative(1)
 

diff --git a/vot/utilities/__init__.py b/vot/utilities/__init__.py
@@ -339,6 +339,13 @@ def to_logical(val):
     except ValueError:
         raise RuntimeError("Logical value conversion error")
 
+def format_size(num, suffix="B"):
+    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
+        if abs(num) < 1024.0:
+            return f"{num:3.1f}{unit}{suffix}"
+        num /= 1024.0
+    return f"{num:.1f}Yi{suffix}"
+
 def singleton(class_):
     instances = {}
     def getinstance(*args, **kwargs):

diff --git a/vot/utilities/net.py b/vot/utilities/net.py
@@ -115,6 +115,9 @@ def download(url, output, callback=None, chunk_size=1024*32, retry=10):
                         if callback:
                             callback(position, total)
 
+                    if position < total:
+                        raise requests.exceptions.RequestException("Connection closed")
+
                     if tmp_file:
                         filehandle.close()
                         shutil.copy(tmp_file, output)