Spaces:
Runtime error
Runtime error
| # based on https://github.com/isl-org/MiDaS | |
| import cv2 | |
| import torch | |
| import torch.nn as nn | |
| import os | |
| models_path = 'pretrained/controlnet/preprocess' | |
| from torchvision.transforms import Compose | |
| from .midas.dpt_depth import DPTDepthModel | |
| from .midas.midas_net import MidasNet | |
| from .midas.midas_net_custom import MidasNet_small | |
| from .midas.transforms import Resize, NormalizeImage, PrepareForNet | |
| base_model_path = os.path.join(models_path, "midas") | |
| old_modeldir = os.path.dirname(os.path.realpath(__file__)) | |
| remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt" | |
| ISL_PATHS = { | |
| "dpt_large": os.path.join(base_model_path, "dpt_large-midas-2f21e586.pt"), | |
| "dpt_hybrid": os.path.join(base_model_path, "dpt_hybrid-midas-501f0c75.pt"), | |
| "midas_v21": "", | |
| "midas_v21_small": "", | |
| } | |
| OLD_ISL_PATHS = { | |
| "dpt_large": os.path.join(old_modeldir, "dpt_large-midas-2f21e586.pt"), | |
| "dpt_hybrid": os.path.join(old_modeldir, "dpt_hybrid-midas-501f0c75.pt"), | |
| "midas_v21": "", | |
| "midas_v21_small": "", | |
| } | |
| def disabled_train(self, mode=True): | |
| """Overwrite model.train with this function to make sure train/eval mode | |
| does not change anymore.""" | |
| return self | |
| def load_midas_transform(model_type): | |
| # https://github.com/isl-org/MiDaS/blob/master/run.py | |
| # load transform only | |
| if model_type == "dpt_large": # DPT-Large | |
| net_w, net_h = 384, 384 | |
| resize_mode = "minimal" | |
| normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) | |
| elif model_type == "dpt_hybrid": # DPT-Hybrid | |
| net_w, net_h = 384, 384 | |
| resize_mode = "minimal" | |
| normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) | |
| elif model_type == "midas_v21": | |
| net_w, net_h = 384, 384 | |
| resize_mode = "upper_bound" | |
| normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| elif model_type == "midas_v21_small": | |
| net_w, net_h = 256, 256 | |
| resize_mode = "upper_bound" | |
| normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| else: | |
| assert False, f"model_type '{model_type}' not implemented, use: --model_type large" | |
| transform = Compose( | |
| [ | |
| Resize( | |
| net_w, | |
| net_h, | |
| resize_target=None, | |
| keep_aspect_ratio=True, | |
| ensure_multiple_of=32, | |
| resize_method=resize_mode, | |
| image_interpolation_method=cv2.INTER_CUBIC, | |
| ), | |
| normalization, | |
| PrepareForNet(), | |
| ] | |
| ) | |
| return transform | |
| def load_file_from_url(url, model_dir=None, progress=True, file_name=None): | |
| """Load file form http url, will download models if necessary. | |
| Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py | |
| Args: | |
| url (str): URL to be downloaded. | |
| model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir. | |
| Default: None. | |
| progress (bool): Whether to show the download progress. Default: True. | |
| file_name (str): The downloaded file name. If None, use the file name in the url. Default: None. | |
| Returns: | |
| str: The path to the downloaded file. | |
| """ | |
| from torch.hub import download_url_to_file, get_dir | |
| from urllib.parse import urlparse | |
| if model_dir is None: # use the pytorch hub_dir | |
| hub_dir = get_dir() | |
| model_dir = os.path.join(hub_dir, 'checkpoints') | |
| os.makedirs(model_dir, exist_ok=True) | |
| parts = urlparse(url) | |
| filename = os.path.basename(parts.path) | |
| if file_name is not None: | |
| filename = file_name | |
| cached_file = os.path.abspath(os.path.join(model_dir, filename)) | |
| if not os.path.exists(cached_file): | |
| print(f'Downloading: "{url}" to {cached_file}\n') | |
| download_url_to_file(url, cached_file, hash_prefix=None, progress=progress) | |
| return cached_file | |
| def load_model(model_type): | |
| # https://github.com/isl-org/MiDaS/blob/master/run.py | |
| # load network | |
| model_path = ISL_PATHS[model_type] | |
| old_model_path = OLD_ISL_PATHS[model_type] | |
| if model_type == "dpt_large": # DPT-Large | |
| model = DPTDepthModel( | |
| path=model_path, | |
| backbone="vitl16_384", | |
| non_negative=True, | |
| ) | |
| net_w, net_h = 384, 384 | |
| resize_mode = "minimal" | |
| normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) | |
| elif model_type == "dpt_hybrid": # DPT-Hybrid | |
| if os.path.exists(old_model_path): | |
| model_path = old_model_path | |
| elif not os.path.exists(model_path): | |
| load_file_from_url(remote_model_path, model_dir=base_model_path) | |
| model = DPTDepthModel( | |
| path=model_path, | |
| backbone="vitb_rn50_384", | |
| non_negative=True, | |
| ) | |
| net_w, net_h = 384, 384 | |
| resize_mode = "minimal" | |
| normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) | |
| elif model_type == "midas_v21": | |
| model = MidasNet(model_path, non_negative=True) | |
| net_w, net_h = 384, 384 | |
| resize_mode = "upper_bound" | |
| normalization = NormalizeImage( | |
| mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] | |
| ) | |
| elif model_type == "midas_v21_small": | |
| model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True, | |
| non_negative=True, blocks={'expand': True}) | |
| net_w, net_h = 256, 256 | |
| resize_mode = "upper_bound" | |
| normalization = NormalizeImage( | |
| mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] | |
| ) | |
| else: | |
| print(f"model_type '{model_type}' not implemented, use: --model_type large") | |
| assert False | |
| transform = Compose( | |
| [ | |
| Resize( | |
| net_w, | |
| net_h, | |
| resize_target=None, | |
| keep_aspect_ratio=True, | |
| ensure_multiple_of=32, | |
| resize_method=resize_mode, | |
| image_interpolation_method=cv2.INTER_CUBIC, | |
| ), | |
| normalization, | |
| PrepareForNet(), | |
| ] | |
| ) | |
| return model.eval(), transform | |
| class MiDaSInference(nn.Module): | |
| MODEL_TYPES_TORCH_HUB = [ | |
| "DPT_Large", | |
| "DPT_Hybrid", | |
| "MiDaS_small" | |
| ] | |
| MODEL_TYPES_ISL = [ | |
| "dpt_large", | |
| "dpt_hybrid", | |
| "midas_v21", | |
| "midas_v21_small", | |
| ] | |
| def __init__(self, model_type): | |
| super().__init__() | |
| assert (model_type in self.MODEL_TYPES_ISL) | |
| model, _ = load_model(model_type) | |
| self.model = model | |
| self.model.train = disabled_train | |
| def forward(self, x): | |
| with torch.no_grad(): | |
| prediction = self.model(x) | |
| return prediction | |