Added download file functions

Author: giswqs

requirements.txt

@@ -6,4 +6,5 @@ onnxruntime; python_version < '3.11'
 onnx
 geopandas
 rasterio
-tqdm
+tqdm
+gdown

samgeo/common.py

@@ -66,6 +66,136 @@ def temp_file_path(extension):
     return file_path
 
 
+def github_raw_url(url):
+    """Get the raw URL for a GitHub file.
+
+    Args:
+        url (str): The GitHub URL.
+    Returns:
+        str: The raw URL.
+    """
+    if isinstance(url, str) and url.startswith("https://github.com/") and "blob" in url:
+        url = url.replace("github.com", "raw.githubusercontent.com").replace(
+            "blob/", ""
+        )
+    return url
+
+
+def download_file(
+    url=None,
+    output=None,
+    quiet=False,
+    proxy=None,
+    speed=None,
+    use_cookies=True,
+    verify=True,
+    id=None,
+    fuzzy=False,
+    resume=False,
+    unzip=True,
+    overwrite=False,
+    subfolder=False,
+):
+    """Download a file from URL, including Google Drive shared URL.
+
+    Args:
+        url (str, optional): Google Drive URL is also supported. Defaults to None.
+        output (str, optional): Output filename. Default is basename of URL.
+        quiet (bool, optional): Suppress terminal output. Default is False.
+        proxy (str, optional): Proxy. Defaults to None.
+        speed (float, optional): Download byte size per second (e.g., 256KB/s = 256 * 1024). Defaults to None.
+        use_cookies (bool, optional): Flag to use cookies. Defaults to True.
+        verify (bool | str, optional): Either a bool, in which case it controls whether the server's TLS certificate is verified, or a string,
+            in which case it must be a path to a CA bundle to use. Default is True.. Defaults to True.
+        id (str, optional): Google Drive's file ID. Defaults to None.
+        fuzzy (bool, optional): Fuzzy extraction of Google Drive's file Id. Defaults to False.
+        resume (bool, optional): Resume the download from existing tmp file if possible. Defaults to False.
+        unzip (bool, optional): Unzip the file. Defaults to True.
+        overwrite (bool, optional): Overwrite the file if it already exists. Defaults to False.
+        subfolder (bool, optional): Create a subfolder with the same name as the file. Defaults to False.
+
+    Returns:
+        str: The output file path.
+    """
+    import zipfile
+    try:
+        import gdown
+    except ImportError:
+        print(
+            "The gdown package is required for this function. Use `pip install gdown` to install it."
+        )
+        return
+
+    if output is None:
+        if isinstance(url, str) and url.startswith("http"):
+            output = os.path.basename(url)
+
+    out_dir = os.path.abspath(os.path.dirname(output))
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+
+    if isinstance(url, str):
+        if os.path.exists(os.path.abspath(output)) and (not overwrite):
+            print(
+                f"{output} already exists. Skip downloading. Set overwrite=True to overwrite."
+            )
+            return os.path.abspath(output)
+        else:
+            url = github_raw_url(url)
+
+    if "https://drive.google.com/file/d/" in url:
+        fuzzy = True
+
+    output = gdown.download(
+        url, output, quiet, proxy, speed, use_cookies, verify, id, fuzzy, resume
+    )
+
+    if unzip and output.endswith(".zip"):
+        with zipfile.ZipFile(output, "r") as zip_ref:
+            if not quiet:
+                print("Extracting files...")
+            if subfolder:
+                basename = os.path.splitext(os.path.basename(output))[0]
+
+                output = os.path.join(out_dir, basename)
+                if not os.path.exists(output):
+                    os.makedirs(output)
+                zip_ref.extractall(output)
+            else:
+                zip_ref.extractall(os.path.dirname(output))
+
+    return os.path.abspath(output)
+
+
+def download_checkpoint(url=None, output=None, overwrite=False, **kwargs):
+    """Download a checkpoint from URL. It can be one of the following: sam_vit_h_4b8939.pth, sam_vit_l_0b3195.pth, sam_vit_b_01ec64.pth.
+
+    Args:
+        url (str, optional): The checkpoint URL. Defaults to None. 
+        output (str, optional): The output file path. Defaults to None.
+        overwrite (bool, optional): Overwrite the file if it already exists. Defaults to False.
+
+    Returns:
+        str: The output file path.
+    """
+    checkpoints = {
+        'sam_vit_h_4b8939.pth': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+        'sam_vit_l_0b3195.pth': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+        'sam_vit_b_01ec64.pth': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+    }
+
+    if isinstance(url, str) and url in checkpoints:
+        url = checkpoints[url]
+
+    if url is None:
+        url = checkpoints['sam_vit_h_4b8939.pth']
+
+    if output is None:
+        output = os.path.basename(url)
+
+    return download_file(url, output,overwrite=overwrite, **kwargs)
+
+
 def image_to_cog(source, dst_path=None, profile="deflate", **kwargs):
     """Converts an image to a COG file.
 

samgeo/samgeo.py

@@ -2,7 +2,7 @@
 The source code is adapted from https://github.com/aliaksandr960/segment-anything-eo. Credit to the author Aliaksandr Hancharenka.
 """
 
-
+import os
 import numpy as np
 import cv2
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
@@ -41,6 +41,12 @@ def __init__(
         mask_multiplier=255,
         sam_kwargs=None,
     ):
+        
+        if not os.path.exists(checkpoint):
+            print(f'Checkpoint {checkpoint} does not exist.')
+            download_checkpoint(output=checkpoint)
+
+
         self.checkpoint = checkpoint
         self.model_type = model_type
         self.device = device