Skip to content

Commit 116e515

Browse files
authored
Change checkpoints download URLs (#259)
1 parent d0e7162 commit 116e515

File tree

4 files changed

+148
-10
lines changed

4 files changed

+148
-10
lines changed

docs/examples/maxar_open_data.ipynb

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,7 @@
6363
},
6464
"outputs": [],
6565
"source": [
66-
"url = (\n",
67-
" \"https://drive.google.com/file/d/1jIIC5hvSPeJEC0fbDhtxVWk2XV9AxsQD/view?usp=sharing\"\n",
68-
")"
66+
"url = \"https://github.com/opengeos/datasets/releases/download/raster/Derna_sample.tif\""
6967
]
7068
},
7169
{

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ leafmap
55
localtileserver
66
matplotlib
77
opencv-python
8+
patool
89
pycocotools
910
pyproj
1011
rasterio

samgeo/common.py

Lines changed: 144 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -209,19 +209,22 @@ def download_checkpoint(model_type="vit_h", checkpoint_dir=None, hq=False):
209209
model_types = {
210210
"vit_h": {
211211
"name": "sam_hq_vit_h.pth",
212-
"url": "https://drive.google.com/file/d/1qobFYrI4eyIANfBSmYcGuWRaSIXfMOQ8/view?usp=sharing",
212+
"url": [
213+
"https://github.com/opengeos/datasets/releases/download/models/sam_hq_vit_h.zip",
214+
"https://github.com/opengeos/datasets/releases/download/models/sam_hq_vit_h.z01",
215+
],
213216
},
214217
"vit_l": {
215218
"name": "sam_hq_vit_l.pth",
216-
"url": "https://drive.google.com/file/d/1Uk17tDKX1YAKas5knI4y9ZJCo0lRVL0G/view?usp=sharing",
219+
"url": "https://github.com/opengeos/datasets/releases/download/models/sam_hq_vit_l.pth",
217220
},
218221
"vit_b": {
219222
"name": "sam_hq_vit_b.pth",
220-
"url": "https://drive.google.com/file/d/11yExZLOve38kRZPfRx_MRxfIAKmfMY47/view?usp=sharing",
223+
"url": "https://github.com/opengeos/datasets/releases/download/models/sam_hq_vit_b.pth",
221224
},
222225
"vit_tiny": {
223226
"name": "sam_hq_vit_tiny.pth",
224-
"url": "https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_tiny.pth",
227+
"url": "https://github.com/opengeos/datasets/releases/download/models/sam_hq_vit_tiny.pth",
225228
},
226229
}
227230

@@ -239,7 +242,10 @@ def download_checkpoint(model_type="vit_h", checkpoint_dir=None, hq=False):
239242
if not os.path.exists(checkpoint):
240243
print(f"Model checkpoint for {model_type} not found.")
241244
url = model_types[model_type]["url"]
242-
download_file(url, checkpoint)
245+
if isinstance(url, str):
246+
download_file(url, checkpoint)
247+
elif isinstance(url, list):
248+
download_files(url, checkpoint_dir, multi_part=True)
243249
return checkpoint
244250

245251

@@ -2987,3 +2993,136 @@ def merge_rasters(
29872993
dstNodata=output_nodata,
29882994
options=output_options,
29892995
)
2996+
2997+
2998+
def extract_archive(archive, outdir=None, **kwargs):
2999+
"""
3000+
Extracts a multipart archive.
3001+
3002+
This function uses the patoolib library to extract a multipart archive.
3003+
If the patoolib library is not installed, it attempts to install it.
3004+
If the archive does not end with ".zip", it appends ".zip" to the archive name.
3005+
If the extraction fails (for example, if the files already exist), it skips the extraction.
3006+
3007+
Args:
3008+
archive (str): The path to the archive file.
3009+
outdir (str): The directory where the archive should be extracted.
3010+
**kwargs: Arbitrary keyword arguments for the patoolib.extract_archive function.
3011+
3012+
Returns:
3013+
None
3014+
3015+
Raises:
3016+
Exception: An exception is raised if the extraction fails for reasons other than the files already existing.
3017+
3018+
Example:
3019+
3020+
files = ["sam_hq_vit_tiny.zip", "sam_hq_vit_tiny.z01", "sam_hq_vit_tiny.z02", "sam_hq_vit_tiny.z03"]
3021+
base_url = "https://github.com/opengeos/datasets/releases/download/models/"
3022+
urls = [base_url + f for f in files]
3023+
leafmap.download_files(urls, out_dir="models", multi_part=True)
3024+
3025+
"""
3026+
try:
3027+
import patoolib
3028+
except ImportError:
3029+
install_package("patool")
3030+
import patoolib
3031+
3032+
if not archive.endswith(".zip"):
3033+
archive = archive + ".zip"
3034+
3035+
if outdir is None:
3036+
outdir = os.path.dirname(archive)
3037+
3038+
try:
3039+
patoolib.extract_archive(archive, outdir=outdir, **kwargs)
3040+
except Exception as e:
3041+
print("The unzipped files might already exist. Skipping extraction.")
3042+
return
3043+
3044+
3045+
def download_files(
3046+
urls,
3047+
out_dir=None,
3048+
filenames=None,
3049+
quiet=False,
3050+
proxy=None,
3051+
speed=None,
3052+
use_cookies=True,
3053+
verify=True,
3054+
id=None,
3055+
fuzzy=False,
3056+
resume=False,
3057+
unzip=True,
3058+
overwrite=False,
3059+
subfolder=False,
3060+
multi_part=False,
3061+
):
3062+
"""Download files from URLs, including Google Drive shared URL.
3063+
3064+
Args:
3065+
urls (list): The list of urls to download. Google Drive URL is also supported.
3066+
out_dir (str, optional): The output directory. Defaults to None.
3067+
filenames (list, optional): Output filename. Default is basename of URL.
3068+
quiet (bool, optional): Suppress terminal output. Default is False.
3069+
proxy (str, optional): Proxy. Defaults to None.
3070+
speed (float, optional): Download byte size per second (e.g., 256KB/s = 256 * 1024). Defaults to None.
3071+
use_cookies (bool, optional): Flag to use cookies. Defaults to True.
3072+
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.
3073+
id (str, optional): Google Drive's file ID. Defaults to None.
3074+
fuzzy (bool, optional): Fuzzy extraction of Google Drive's file Id. Defaults to False.
3075+
resume (bool, optional): Resume the download from existing tmp file if possible. Defaults to False.
3076+
unzip (bool, optional): Unzip the file. Defaults to True.
3077+
overwrite (bool, optional): Overwrite the file if it already exists. Defaults to False.
3078+
subfolder (bool, optional): Create a subfolder with the same name as the file. Defaults to False.
3079+
multi_part (bool, optional): If the file is a multi-part file. Defaults to False.
3080+
3081+
Examples:
3082+
3083+
files = ["sam_hq_vit_tiny.zip", "sam_hq_vit_tiny.z01", "sam_hq_vit_tiny.z02", "sam_hq_vit_tiny.z03"]
3084+
base_url = "https://github.com/opengeos/datasets/releases/download/models/"
3085+
urls = [base_url + f for f in files]
3086+
leafmap.download_files(urls, out_dir="models", multi_part=True)
3087+
"""
3088+
3089+
if out_dir is None:
3090+
out_dir = os.getcwd()
3091+
3092+
if filenames is None:
3093+
filenames = [None] * len(urls)
3094+
3095+
filepaths = []
3096+
for url, output in zip(urls, filenames):
3097+
if output is None:
3098+
filename = os.path.join(out_dir, os.path.basename(url))
3099+
else:
3100+
filename = os.path.join(out_dir, output)
3101+
3102+
filepaths.append(filename)
3103+
if multi_part:
3104+
unzip = False
3105+
3106+
download_file(
3107+
url,
3108+
filename,
3109+
quiet,
3110+
proxy,
3111+
speed,
3112+
use_cookies,
3113+
verify,
3114+
id,
3115+
fuzzy,
3116+
resume,
3117+
unzip,
3118+
overwrite,
3119+
subfolder,
3120+
)
3121+
3122+
if multi_part:
3123+
archive = os.path.splitext(filename)[0] + ".zip"
3124+
out_dir = os.path.dirname(filename)
3125+
extract_archive(archive, out_dir)
3126+
3127+
for file in filepaths:
3128+
os.remove(file)

samgeo/fast_sam.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@ def __init__(self, model="FastSAM-x.pt", **kwargs):
3030
)
3131

3232
models = {
33-
"FastSAM-x.pt": "https://drive.google.com/file/d/1m1sjY4ihXBU1fZXdQ-Xdj-mDltW-2Rqv/view?usp=sharing",
34-
"FastSAM-s.pt": "https://drive.google.com/file/d/10XmSj6mmpmRb8NhXbtiuO9cTTBwR_9SV/view?usp=sharing",
33+
"FastSAM-x.pt": "https://github.com/opengeos/datasets/releases/download/models/FastSAM-x.pt",
34+
"FastSAM-s.pt": "https://github.com/opengeos/datasets/releases/download/models/FastSAM-s.pt",
3535
}
3636

3737
if model not in models:

0 commit comments

Comments
 (0)