otmanheddouch/house_design

1# based on https://github.com/isl-org/MiDaS2 3import cv24import os5import torch6import torch.nn as nn7from torchvision.transforms import Compose8 9from .midas.dpt_depth import DPTDepthModel10from .midas.midas_net import MidasNet11from .midas.midas_net_custom import MidasNet_small12from .midas.transforms import Resize, NormalizeImage, PrepareForNet13from annotator.util import annotator_ckpts_path14 15 16ISL_PATHS = {17    "dpt_large": os.path.join(annotator_ckpts_path, "dpt_large-midas-2f21e586.pt"),18    "dpt_hybrid": os.path.join(annotator_ckpts_path, "dpt_hybrid-midas-501f0c75.pt"),19    "midas_v21": "",20    "midas_v21_small": "",21}22 23remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"24 25 26def disabled_train(self, mode=True):27    """Overwrite model.train with this function to make sure train/eval mode28    does not change anymore."""29    return self30 31 32def load_midas_transform(model_type):33    # https://github.com/isl-org/MiDaS/blob/master/run.py34    # load transform only35    if model_type == "dpt_large":  # DPT-Large36        net_w, net_h = 384, 38437        resize_mode = "minimal"38        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])39 40    elif model_type == "dpt_hybrid":  # DPT-Hybrid41        net_w, net_h = 384, 38442        resize_mode = "minimal"43        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])44 45    elif model_type == "midas_v21":46        net_w, net_h = 384, 38447        resize_mode = "upper_bound"48        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])49 50    elif model_type == "midas_v21_small":51        net_w, net_h = 256, 25652        resize_mode = "upper_bound"53        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])54 55    else:56        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"57 58    transform = Compose(59        [60            Resize(61                net_w,62                net_h,63                resize_target=None,64                keep_aspect_ratio=True,65                ensure_multiple_of=32,66                resize_method=resize_mode,67                image_interpolation_method=cv2.INTER_CUBIC,68            ),69            normalization,70            PrepareForNet(),71        ]72    )73 74    return transform75 76 77def load_model(model_type):78    # https://github.com/isl-org/MiDaS/blob/master/run.py79    # load network80    model_path = ISL_PATHS[model_type]81    if model_type == "dpt_large":  # DPT-Large82        model = DPTDepthModel(83            path=model_path,84            backbone="vitl16_384",85            non_negative=True,86        )87        net_w, net_h = 384, 38488        resize_mode = "minimal"89        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])90 91    elif model_type == "dpt_hybrid":  # DPT-Hybrid92        if not os.path.exists(model_path):93            from basicsr.utils.download_util import load_file_from_url94            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)95 96        model = DPTDepthModel(97            path=model_path,98            backbone="vitb_rn50_384",99            non_negative=True,100        )101        net_w, net_h = 384, 384102        resize_mode = "minimal"103        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])104 105    elif model_type == "midas_v21":106        model = MidasNet(model_path, non_negative=True)107        net_w, net_h = 384, 384108        resize_mode = "upper_bound"109        normalization = NormalizeImage(110            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]111        )112 113    elif model_type == "midas_v21_small":114        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,115                               non_negative=True, blocks={'expand': True})116        net_w, net_h = 256, 256117        resize_mode = "upper_bound"118        normalization = NormalizeImage(119            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]120        )121 122    else:123        print(f"model_type '{model_type}' not implemented, use: --model_type large")124        assert False125 126    transform = Compose(127        [128            Resize(129                net_w,130                net_h,131                resize_target=None,132                keep_aspect_ratio=True,133                ensure_multiple_of=32,134                resize_method=resize_mode,135                image_interpolation_method=cv2.INTER_CUBIC,136            ),137            normalization,138            PrepareForNet(),139        ]140    )141 142    return model.eval(), transform143 144 145class MiDaSInference(nn.Module):146    MODEL_TYPES_TORCH_HUB = [147        "DPT_Large",148        "DPT_Hybrid",149        "MiDaS_small"150    ]151    MODEL_TYPES_ISL = [152        "dpt_large",153        "dpt_hybrid",154        "midas_v21",155        "midas_v21_small",156    ]157 158    def __init__(self, model_type):159        super().__init__()160        assert (model_type in self.MODEL_TYPES_ISL)161        model, _ = load_model(model_type)162        self.model = model163        self.model.train = disabled_train164 165    def forward(self, x):166        with torch.no_grad():167            prediction = self.model(x)168        return prediction169 170