Spaces:
Running
Running
| 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) | |