Spaces:
Running
Running
File size: 1,697 Bytes
c608946 c74a070 c608946 c74a070 c608946 c74a070 c608946 c74a070 c608946 c74a070 c608946 c74a070 c608946 c74a070 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
weight_urls = {
"DKMv3": {
"outdoor": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_outdoor.pth",
"indoor": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_indoor.pth",
},
}
import torch
from .DKMv3 import DKMv3
def DKMv3_outdoor(path_to_weights=None, device=None):
"""
Loads DKMv3 outdoor weights, uses internal resolution of (540, 720) by default
resolution can be changed by setting model.h_resized, model.w_resized later.
Additionally upsamples preds to fixed resolution of (864, 1152),
can be turned off by model.upsample_preds = False
"""
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if path_to_weights is not None:
weights = torch.load(path_to_weights, map_location="cpu")
else:
weights = torch.hub.load_state_dict_from_url(
weight_urls["DKMv3"]["outdoor"], map_location="cpu"
)
return DKMv3(weights, 540, 720, upsample_preds=True, device=device)
def DKMv3_indoor(path_to_weights=None, device=None):
"""
Loads DKMv3 indoor weights, uses internal resolution of (480, 640) by default
Resolution can be changed by setting model.h_resized, model.w_resized later.
"""
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if path_to_weights is not None:
weights = torch.load(path_to_weights, map_location=device)
else:
weights = torch.hub.load_state_dict_from_url(
weight_urls["DKMv3"]["indoor"], map_location=device
)
return DKMv3(weights, 480, 640, upsample_preds=False, device=device)
|