src/depthmap_generation.py

import gc
import os.path
from operator import getitem

import cv2
import numpy as np
import skimage.measure
from PIL import Image
import torch
from torchvision.transforms import Compose, transforms

# midas imports
from dmidas.dpt_depth import DPTDepthModel
from dmidas.midas_net import MidasNet
from dmidas.midas_net_custom import MidasNet_small
from dmidas.transforms import Resize, NormalizeImage, PrepareForNet
# zoedepth
from dzoedepth.models.builder import build_model
from dzoedepth.utils.config import get_config
# AdelaiDepth/LeReS imports
from lib.multi_depth_model_woauxi import RelDepthModel
from lib.net_tools import strip_prefix_if_present
from pix2pix.models.pix2pix4depth_model import Pix2Pix4DepthModel
# Marigold
from dmarigold.marigold import MarigoldPipeline
# pix2pix/merge net imports
from pix2pix.options.test_options import TestOptions
# depthanyting v2
try:
    from ddepth_anything_v2 import DepthAnythingV2
except:
    print('depth_anything_v2 import failed... somehow')

# Our code
from src.misc import *
from src import backbone

global depthmap_device

class ModelHolder:
    def __init__(self):
        self.depth_model = None
        self.pix2pix_model = None
        self.depth_model_type = None
        self.device = None  # Target device, the model may be swapped from VRAM into RAM.
        self.offloaded = False  # True means current device is not the target device

        # Extra stuff
        self.resize_mode = None
        self.normalization = None
        self.tiling_mode = False


    def update_settings(self, **kvargs):
        # Opens the pandora box
        for k, v in kvargs.items():
            setattr(self, k, v)


    def ensure_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False):
        # TODO: could make it more granular
        if model_type == -1 or model_type is None:
            self.unload_models()
            return
        # Certain optimisations are irreversible and not device-agnostic, thus changing device requires reloading
        if (
                model_type != self.depth_model_type or
                boost != (self.pix2pix_model is not None) or
                device != self.device or
                tiling_mode != self.tiling_mode
        ):
            self.unload_models()
            self.load_models(model_type, device, boost, tiling_mode)
        self.reload()

    def load_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False):
        """Ensure that the depth model is loaded"""

        # TODO: we need to at least try to find models downloaded by other plugins (e.g. controlnet)

        # model path and name
        # ZoeDepth and Marigold do not use this
        model_dir = "./models/midas"
        if model_type == 0:
            model_dir = "./models/leres"
        if model_type == 11:
            model_dir = "./models/depth_anything"
        if model_type in [12, 13, 14]:
            model_dir = "./models/depth_anything_v2"

        # create paths to model if not present
        os.makedirs(model_dir, exist_ok=True)
        os.makedirs('./models/pix2pix', exist_ok=True)

        print("Loading model weights from ", end=" ")

        resize_mode = "minimal"
        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

        model = None
        if model_type == 0:  # "res101"
            model_path = f"{model_dir}/res101.pth"
            print(model_path)
            ensure_file_downloaded(
                model_path,
                ["https://cloudstor.aarnet.edu.au/plus/s/lTIJF4vrvHCAI31/download",
                 "https://huggingface.co/lllyasviel/Annotators/resolve/5bc80eec2b4fddbb/res101.pth",
                 ],
                "1d696b2ef3e8336b057d0c15bc82d2fecef821bfebe5ef9d7671a5ec5dde520b")
            if device != torch.device('cpu'):
                checkpoint = torch.load(model_path)
            else:
                checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
            model = RelDepthModel(backbone='resnext101')
            model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), strict=True)
            del checkpoint
            backbone.torch_gc()

        if model_type == 1:  # "dpt_beit_large_512" midas 3.1
            model_path = f"{model_dir}/dpt_beit_large_512.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt")
            model = DPTDepthModel(
                path=model_path,
                backbone="beitl16_512",
                non_negative=True,
            )
            resize_mode = "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

        if model_type == 2:  # "dpt_beit_large_384" midas 3.1
            model_path = f"{model_dir}/dpt_beit_large_384.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_384.pt")
            model = DPTDepthModel(
                path=model_path,
                backbone="beitl16_384",
                non_negative=True,
            )
            resize_mode = "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

        if model_type == 3:  # "dpt_large_384" midas 3.0
            model_path = f"{model_dir}/dpt_large-midas-2f21e586.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt")
            model = DPTDepthModel(
                path=model_path,
                backbone="vitl16_384",
                non_negative=True,
            )
            resize_mode = "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

        elif model_type == 4:  # "dpt_hybrid_384" midas 3.0
            model_path = f"{model_dir}/dpt_hybrid-midas-501f0c75.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt")
            model = DPTDepthModel(
                path=model_path,
                backbone="vitb_rn50_384",
                non_negative=True,
            )
            resize_mode = "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

        elif model_type == 5:  # "midas_v21"
            model_path = f"{model_dir}/midas_v21-f6b98070.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt")
            model = MidasNet(model_path, non_negative=True)
            resize_mode = "upper_bound"
            normalization = NormalizeImage(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            )

        elif model_type == 6:  # "midas_v21_small"
            model_path = f"{model_dir}/midas_v21_small-70d6b9c8.pt"
            print(model_path)
            ensure_file_downloaded(model_path,
                                   "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt")
            model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
                                   non_negative=True, blocks={'expand': True})
            resize_mode = "upper_bound"
            normalization = NormalizeImage(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            )

        # When loading, zoedepth models will report the default net size.
        # It will be overridden by the generation settings.
        elif model_type == 7:  # zoedepth_n
            print("zoedepth_n\n")
            conf = get_config("zoedepth", "infer")
            model = build_model(conf)

        elif model_type == 8:  # zoedepth_k
            print("zoedepth_k\n")
            conf = get_config("zoedepth", "infer", config_version="kitti")
            model = build_model(conf)

        elif model_type == 9:  # zoedepth_nk
            print("zoedepth_nk\n")
            conf = get_config("zoedepth_nk", "infer")
            model = build_model(conf)

        elif model_type == 10:  # Marigold v1
            model_path = "Bingxin/Marigold"
            print(model_path)
            dtype = torch.float32 if self.no_half else torch.float16
            model = MarigoldPipeline.from_pretrained(model_path, torch_dtype=dtype)
            try:
                import xformers
                model.enable_xformers_memory_efficient_attention()
            except:
                pass  # run without xformers
        elif model_type == 11:  # depth_anything
            from depth_anything.dpt import DPT_DINOv2
            # This will download the model... to some place
            model = (
                DPT_DINOv2(
                    encoder="vitl",
                    features=256,
                    out_channels=[256, 512, 1024, 1024],
                    localhub=False,
                ).to(device).eval()
            )
            model_path = f"{model_dir}/depth_anything_vitl14.pth"
            ensure_file_downloaded(model_path,
                                   "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitl14.pth")

            model.load_state_dict(torch.load(model_path))
        elif model_type in [12, 13, 14]:  # depth_anything_v2 small, base, large
            letter = {12: 's', 13: 'b', 14: 'l'}[model_type]
            word = {12: 'Small', 13: 'Base', 14: 'Large'}[model_type]
            model_path = f"{model_dir}/depth_anything_v2_vit{letter}.pth"
            ensure_file_downloaded(model_path,
                                   f"https://huggingface.co/depth-anything/Depth-Anything-V2-{word}/resolve/main/depth_anything_v2_vit{letter}.pth")
            model_configs = {'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
                             'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
                             'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
                             'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}}
            model = DepthAnythingV2(**model_configs[f'vit{letter}'])
            model.load_state_dict(torch.load(model_path, map_location='cpu'))
        # 15 is reserved for Depth Anything V2 Giant

        if tiling_mode:
            def flatten(el):
                flattened = [flatten(children) for children in el.children()]
                res = [el]
                for c in flattened:
                    res += c
                return res
            layers = flatten(model)  # Hijacking the model
            for layer in [layer for layer in layers if type(layer) == torch.nn.Conv2d or type(layer) == torch.nn.Conv1d]:
                layer.padding_mode = 'circular'

        if model_type in range(0, 10):
            model.eval()  # prepare for evaluation
        # optimize
        if device == torch.device("cuda"):
            if model_type in [0, 1, 2, 3, 4, 5, 6]:
                model = model.to(memory_format=torch.channels_last)  # TODO: weird
            if not self.no_half:
                # Marigold can be done
                # TODO: Fix for zoedepth_n - it completely trips and generates black images
                if model_type in [1, 2, 3, 4, 5, 6, 8, 9, 11] and not boost:
                    model = model.half()
                if model_type in [12, 13, 14]:
                    model.depth_head.half()
                    model.pretrained.half()
        model.to(device)  # to correct device

        self.depth_model = model
        self.depth_model_type = model_type
        self.resize_mode = resize_mode
        self.normalization = normalization
        self.tiling_mode = tiling_mode

        self.device = device

        if boost:
            # sfu.ca unfortunately is not very reliable, we use a mirror just in case
            ensure_file_downloaded(
                './models/pix2pix/latest_net_G.pth',
                ["https://huggingface.co/lllyasviel/Annotators/resolve/9a7d84251d487d11/latest_net_G.pth",
                 "https://sfu.ca/~yagiz/CVPR21/latest_net_G.pth"],
                '50ec735d74ed6499562d898f41b49343e521808b8dae589aa3c2f5c9ac9f7462')
            opt = TestOptions().parse()
            if device == torch.device('cpu'):
                opt.gpu_ids = []
            self.pix2pix_model = Pix2Pix4DepthModel(opt)
            self.pix2pix_model.save_dir = './models/pix2pix'
            self.pix2pix_model.load_networks('latest')
            self.pix2pix_model.eval()

        backbone.torch_gc()

    @staticmethod
    def get_default_net_size(model_type):
        # Have you ever wondered why so many things in so many code repositories are not optimal?
        # For example, this here is a set of int:tuple. Why wouldn't it be a set of enum:tuple?
        # Or even better, why won't every model be defined separately with all it's necessary values and constants in one place? And why one like of this comment is much longer than the other ones?!
        # Why won't the models indexed by enum elements, not integers?
        # The answer is as definite as it is horrifying: tech depth.
        # This here is a prime example of how tech debt piles up: one slightly iffy decision a long time ago,
        # then nothing is done with it for quite some time, stuff starts depending on it, more stuff is added.
        # The old code are like blocks are like jenga blocks that are experiencing ever-increasing pressure,
        # in tower that (just as code) grows to infinity. And noone wants to knock out the jenga.
        # Noone wants to spend hours of their life fixing it - because adding new features is more exciting.
        # Once merely a suboptimal thing, that worked perfectly at a time, turns into this monster that slowly
        # takes your sanity away. It's not that it ambushes you directly - like a hungry moskquito it knows that
        # being too annoying will warrant immediate action and smashing. Instead, it bothers you just a
        # couple of sound decibels and droplets of blood less than necessary for you to actually go and deal with it.
        # And mind you, this is one buffed maskito: well, actually it got beefed up with time.
        # Now it is just a giant mockyto monster. Noone wants to fight it because it is scary,
        # and thus this threshold of pain is much higher. Don't repeat our mistakes: fight the giant mojito monsters and
        # don't let them spread!
        sizes = {
            0: [448, 448],
            1: [512, 512],
            2: [384, 384],
            3: [384, 384],
            4: [384, 384],
            5: [384, 384],
            6: [256, 256],
            7: [384, 512],
            8: [384, 768],
            9: [384, 512],
            10: [768, 768],
            11: [518, 518],
            12: [518, 518],
            13: [518, 518],
            14: [518, 518]
        }
        if model_type in sizes:
            return sizes[model_type]
        return [512, 512]

    def offload(self):
        """Move to RAM to conserve VRAM"""
        if self.device != torch.device('cpu') and not self.offloaded:
            self.move_models_to(torch.device('cpu'))
            self.offloaded = True

    def reload(self):
        """Undoes offload"""
        if self.offloaded:
            self.move_models_to(self.device)
            self.offloaded = False

    def move_models_to(self, device):
        if self.depth_model is not None:
            self.depth_model.to(device)
        if self.pix2pix_model is not None:
            pass
            # TODO: pix2pix offloading not implemented

    def unload_models(self):
        if self.depth_model is not None or self.pix2pix_model is not None:
            del self.depth_model
            self.depth_model = None
            del self.pix2pix_model
            self.pix2pix_model = None
            gc.collect()
            backbone.torch_gc()

        self.depth_model_type = None
        self.device = None

    def get_raw_prediction(self, input, net_width, net_height):
        """Get prediction from the model currently loaded by the ModelHolder object.
        If boost is enabled, net_width and net_height will be ignored."""
        global depthmap_device
        depthmap_device = self.device
        # input image
        img = cv2.cvtColor(np.asarray(input), cv2.COLOR_BGR2RGB) / 255.0
        # compute depthmap
        if self.pix2pix_model is None:
            if self.depth_model_type == 0:
                raw_prediction = estimateleres(img, self.depth_model, net_width, net_height)
            elif self.depth_model_type in [7, 8, 9]:
                raw_prediction = estimatezoedepth(input, self.depth_model, net_width, net_height)
            elif self.depth_model_type in [1, 2, 3, 4, 5, 6]:
                raw_prediction = estimatemidas(img, self.depth_model, net_width, net_height,
                                               self.resize_mode, self.normalization, self.no_half,
                                               self.precision == "autocast")
            elif self.depth_model_type == 10:
                raw_prediction = estimatemarigold(img, self.depth_model, net_width, net_height,
                                                  self.marigold_ensembles, self.marigold_steps)
            elif self.depth_model_type == 11:
                raw_prediction = estimatedepthanything(img, self.depth_model, net_width, net_height)
            elif self.depth_model_type in [12, 13, 14]:
                raw_prediction = estimatedepthanything_v2(img, self.depth_model, net_width, net_height)
        else:
            raw_prediction = estimateboost(img, self.depth_model, self.depth_model_type, self.pix2pix_model,
                                           self.boost_rmax)
        raw_prediction_invert = self.depth_model_type in [0, 7, 8, 9, 10]
        return raw_prediction, raw_prediction_invert


def estimateleres(img, model, w, h):
    # leres transform input
    rgb_c = img[:, :, ::-1].copy()
    A_resize = cv2.resize(rgb_c, (w, h))
    img_torch = scale_torch(A_resize)[None, :, :, :]

    # compute
    with torch.no_grad():
        if depthmap_device == torch.device("cuda"):
            img_torch = img_torch.cuda()
        prediction = model.depth_model(img_torch)

    prediction = prediction.squeeze().cpu().numpy()
    prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)

    return prediction


def scale_torch(img):
    """
    Scale the image and output it in torch.tensor.
    :param img: input rgb is in shape [H, W, C], input depth/disp is in shape [H, W]
    :param scale: the scale factor. float
    :return: img. [C, H, W]
    """
    if len(img.shape) == 2:
        img = img[np.newaxis, :, :]
    if img.shape[2] == 3:
        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
        img = transform(img.astype(np.float32))
    else:
        img = img.astype(np.float32)
        img = torch.from_numpy(img)
    return img


def estimatezoedepth(img, model, w, h):
    # x = transforms.ToTensor()(img).unsqueeze(0)
    # x = x.type(torch.float32)
    # x.to(depthmap_device)
    # prediction = model.infer(x)
    model.core.prep.resizer._Resize__width = w
    model.core.prep.resizer._Resize__height = h
    prediction = model.infer_pil(img)

    return prediction


def estimatemidas(img, model, w, h, resize_mode, normalization, no_half, precision_is_autocast):
    import contextlib
    # init transform
    transform = Compose(
        [
            Resize(
                w,
                h,
                resize_target=None,
                keep_aspect_ratio=True,
                ensure_multiple_of=32,
                resize_method=resize_mode,
                image_interpolation_method=cv2.INTER_CUBIC,
            ),
            normalization,
            PrepareForNet(),
        ]
    )

    # transform input
    img_input = transform({"image": img})["image"]

    # compute
    precision_scope = torch.autocast if precision_is_autocast and depthmap_device == torch.device(
        "cuda") else contextlib.nullcontext
    with torch.no_grad(), precision_scope("cuda"):
        sample = torch.from_numpy(img_input).to(depthmap_device).unsqueeze(0)
        if depthmap_device == torch.device("cuda"):
            sample = sample.to(memory_format=torch.channels_last)
            if not no_half:
                sample = sample.half()
        prediction = model.forward(sample)
        prediction = (
            torch.nn.functional.interpolate(
                prediction.unsqueeze(1),
                size=img.shape[:2],
                mode="bicubic",
                align_corners=False,
            )
            .squeeze()
            .cpu()
            .numpy()
        )

    return prediction


# TODO: correct values for BOOST
# TODO: "h" is not used
def estimatemarigold(image, model, w, h, marigold_ensembles=5, marigold_steps=12):
    # This hideous thing should be re-implemented once there is support from the upstream.
    # TODO: re-implement this hideous thing by using features from the upstream
    img = cv2.cvtColor((image * 255.0001).astype('uint8'), cv2.COLOR_BGR2RGB)
    img = Image.fromarray(img)
    with torch.no_grad():
        pipe_out = model(img, processing_res=w, show_progress_bar=False,
                         ensemble_size=marigold_ensembles, denoising_steps=marigold_steps,
                         match_input_res=False)
        return cv2.resize(pipe_out.depth_np, (image.shape[:2][::-1]), interpolation=cv2.INTER_CUBIC)


def estimatedepthanything(image, model, w, h):
    from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
    transform = Compose(
        [
            Resize(
                width=w // 14 * 14,
                height=h // 14 * 14,
                resize_target=False,
                keep_aspect_ratio=True,
                ensure_multiple_of=14,
                resize_method="lower_bound",
                image_interpolation_method=cv2.INTER_CUBIC,
            ),
            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            PrepareForNet(),
        ]
    )

    timage = transform({"image": image})["image"]
    timage = torch.from_numpy(timage).unsqueeze(0).to(device=next(model.parameters()).device,
                                                      dtype=next(model.parameters()).dtype)

    with torch.no_grad():
        depth = model(timage)
    import torch.nn.functional as F
    depth = F.interpolate(
        depth[None], (image.shape[0], image.shape[1]), mode="bilinear", align_corners=False
    )[0, 0]

    return depth.cpu().numpy()


def estimatedepthanything_v2(image, model, w, h):
    # This is an awkward re-conversion, but I believe it should not impact quality
    img = cv2.cvtColor((image * 255.1).astype('uint8'), cv2.COLOR_BGR2RGB)
    with torch.no_grad():
        # Compare to: model.infer_image(img, w)
        image, (h, w) = model.image2tensor(img, w)
        # Casting to correct type, it is the same as type of some model tensor (the one here is arbitrary)
        image_casted = image.type_as(model.pretrained.blocks[0].norm1.weight.data)
        depth = model.forward(image_casted).type_as(image)
        import torch.nn.functional as F
        depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
        return depth.cpu().numpy()


class ImageandPatchs:
    def __init__(self, root_dir, name, patchsinfo, rgb_image, scale=1):
        self.root_dir = root_dir
        self.patchsinfo = patchsinfo
        self.name = name
        self.patchs = patchsinfo
        self.scale = scale

        self.rgb_image = cv2.resize(rgb_image, (round(rgb_image.shape[1] * scale), round(rgb_image.shape[0] * scale)),
                                    interpolation=cv2.INTER_CUBIC)

        self.do_have_estimate = False
        self.estimation_updated_image = None
        self.estimation_base_image = None

    def __len__(self):
        return len(self.patchs)

    def set_base_estimate(self, est):
        self.estimation_base_image = est
        if self.estimation_updated_image is not None:
            self.do_have_estimate = True

    def set_updated_estimate(self, est):
        self.estimation_updated_image = est
        if self.estimation_base_image is not None:
            self.do_have_estimate = True

    def __getitem__(self, index):
        patch_id = int(self.patchs[index][0])
        rect = np.array(self.patchs[index][1]['rect'])
        msize = self.patchs[index][1]['size']

        ## applying scale to rect:
        rect = np.round(rect * self.scale)
        rect = rect.astype('int')
        msize = round(msize * self.scale)

        patch_rgb = impatch(self.rgb_image, rect)
        if self.do_have_estimate:
            patch_whole_estimate_base = impatch(self.estimation_base_image, rect)
            patch_whole_estimate_updated = impatch(self.estimation_updated_image, rect)
            return {'patch_rgb': patch_rgb, 'patch_whole_estimate_base': patch_whole_estimate_base,
                    'patch_whole_estimate_updated': patch_whole_estimate_updated, 'rect': rect,
                    'size': msize, 'id': patch_id}
        else:
            return {'patch_rgb': patch_rgb, 'rect': rect, 'size': msize, 'id': patch_id}

    def print_options(self, opt):
        """Print and save options

        It will print both current options and default values(if different).
        It will save options into a text file / [checkpoints_dir] / opt.txt
        """
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        print(message)

        # save to the disk
        """
        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
        util.mkdirs(expr_dir)
        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
        with open(file_name, 'wt') as opt_file:
            opt_file.write(message)
            opt_file.write('\n')
        """

    def parse(self):
        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
        opt = self.gather_options()
        opt.isTrain = self.isTrain  # train or test

        # process opt.suffix
        if opt.suffix:
            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
            opt.name = opt.name + suffix

        # self.print_options(opt)

        # set gpu ids
        str_ids = opt.gpu_ids.split(',')
        opt.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                opt.gpu_ids.append(id)
        # if len(opt.gpu_ids) > 0:
        #    torch.cuda.set_device(opt.gpu_ids[0])

        self.opt = opt
        return self.opt


def impatch(image, rect):
    # Extract the given patch pixels from a given image.
    w1 = rect[0]
    h1 = rect[1]
    w2 = w1 + rect[2]
    h2 = h1 + rect[3]
    image_patch = image[h1:h2, w1:w2]
    return image_patch


class ImageandPatchs:
    def __init__(self, root_dir, name, patchsinfo, rgb_image, scale=1):
        self.root_dir = root_dir
        self.patchsinfo = patchsinfo
        self.name = name
        self.patchs = patchsinfo
        self.scale = scale

        self.rgb_image = cv2.resize(rgb_image, (round(rgb_image.shape[1] * scale), round(rgb_image.shape[0] * scale)),
                                    interpolation=cv2.INTER_CUBIC)

        self.do_have_estimate = False
        self.estimation_updated_image = None
        self.estimation_base_image = None

    def __len__(self):
        return len(self.patchs)

    def set_base_estimate(self, est):
        self.estimation_base_image = est
        if self.estimation_updated_image is not None:
            self.do_have_estimate = True

    def set_updated_estimate(self, est):
        self.estimation_updated_image = est
        if self.estimation_base_image is not None:
            self.do_have_estimate = True

    def __getitem__(self, index):
        patch_id = int(self.patchs[index][0])
        rect = np.array(self.patchs[index][1]['rect'])
        msize = self.patchs[index][1]['size']

        ## applying scale to rect:
        rect = np.round(rect * self.scale)
        rect = rect.astype('int')
        msize = round(msize * self.scale)

        patch_rgb = impatch(self.rgb_image, rect)
        if self.do_have_estimate:
            patch_whole_estimate_base = impatch(self.estimation_base_image, rect)
            patch_whole_estimate_updated = impatch(self.estimation_updated_image, rect)
            return {'patch_rgb': patch_rgb, 'patch_whole_estimate_base': patch_whole_estimate_base,
                    'patch_whole_estimate_updated': patch_whole_estimate_updated, 'rect': rect,
                    'size': msize, 'id': patch_id}
        else:
            return {'patch_rgb': patch_rgb, 'rect': rect, 'size': msize, 'id': patch_id}

    def print_options(self, opt):
        """Print and save options

        It will print both current options and default values(if different).
        It will save options into a text file / [checkpoints_dir] / opt.txt
        """
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        print(message)

        # save to the disk
        """
        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
        util.mkdirs(expr_dir)
        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
        with open(file_name, 'wt') as opt_file:
            opt_file.write(message)
            opt_file.write('\n')
        """

    def parse(self):
        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
        opt = self.gather_options()
        opt.isTrain = self.isTrain  # train or test

        # process opt.suffix
        if opt.suffix:
            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
            opt.name = opt.name + suffix

        # self.print_options(opt)

        # set gpu ids
        str_ids = opt.gpu_ids.split(',')
        opt.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                opt.gpu_ids.append(id)
        # if len(opt.gpu_ids) > 0:
        #    torch.cuda.set_device(opt.gpu_ids[0])

        self.opt = opt
        return self.opt


def estimateboost(img, model, model_type, pix2pixmodel, whole_size_threshold):
    pix2pixsize = 1024  # TODO: pix2pixsize and whole_size_threshold to setting?

    if model_type == 0:  # leres
        net_receptive_field_size = 448
    elif model_type == 1:  # dpt_beit_large_512
        net_receptive_field_size = 512
    elif model_type == 11:  # depth_anything
        net_receptive_field_size = 518
    elif model_type in [12, 13, 14]:  # depth_anything_v2
        net_receptive_field_size = 518
    else:  # other midas  # TODO Marigold support
        net_receptive_field_size = 384
    patch_netsize = 2 * net_receptive_field_size
    # Good luck trying to use zoedepth

    gc.collect()
    backbone.torch_gc()

    # Generate mask used to smoothly blend the local pathc estimations to the base estimate.
    # It is arbitrarily large to avoid artifacts during rescaling for each crop.
    mask_org = generatemask((3000, 3000))
    mask = mask_org.copy()

    # Value x of R_x defined in the section 5 of the main paper.
    r_threshold_value = 0.2
    # if R0:
    #    r_threshold_value = 0

    input_resolution = img.shape
    scale_threshold = 3  # Allows up-scaling with a scale up to 3

    # Find the best input resolution R-x. The resolution search described in section 5-double estimation of the main paper and section B of the
    # supplementary material.
    whole_image_optimal_size, patch_scale = calculateprocessingres(img, net_receptive_field_size, r_threshold_value,
                                                                   scale_threshold, whole_size_threshold)

    print('wholeImage being processed in :', whole_image_optimal_size)

    # Generate the base estimate using the double estimation.
    whole_estimate = doubleestimate(img, net_receptive_field_size, whole_image_optimal_size, pix2pixsize, model,
                                    model_type, pix2pixmodel)

    # Compute the multiplier described in section 6 of the main paper to make sure our initial patch can select
    # small high-density regions of the image.
    factor = max(min(1, 4 * patch_scale * whole_image_optimal_size / whole_size_threshold), 0.2)
    print('Adjust factor is:', 1 / factor)

    # Compute the default target resolution.
    if img.shape[0] > img.shape[1]:
        a = 2 * whole_image_optimal_size
        b = round(2 * whole_image_optimal_size * img.shape[1] / img.shape[0])
    else:
        a = round(2 * whole_image_optimal_size * img.shape[0] / img.shape[1])
        b = 2 * whole_image_optimal_size
    b = int(round(b / factor))
    a = int(round(a / factor))

    """
    # recompute a, b and saturate to max res.
    if max(a,b) > max_res:
        print('Default Res is higher than max-res: Reducing final resolution')
        if img.shape[0] > img.shape[1]:
            a = max_res
            b = round(option.max_res * img.shape[1] / img.shape[0])
        else:
            a = round(option.max_res * img.shape[0] / img.shape[1])
            b = max_res
        b = int(b)
        a = int(a)
    """

    img = cv2.resize(img, (b, a), interpolation=cv2.INTER_CUBIC)

    # Extract selected patches for local refinement
    base_size = net_receptive_field_size * 2
    patchset = generatepatchs(img, base_size, factor)

    print('Target resolution: ', img.shape)

    # Computing a scale in case user prompted to generate the results as the same resolution of the input.
    # Notice that our method output resolution is independent of the input resolution and this parameter will only
    # enable a scaling operation during the local patch merge implementation to generate results with the same resolution
    # as the input.
    """
    if output_resolution == 1:
        mergein_scale = input_resolution[0] / img.shape[0]
        print('Dynamicly change merged-in resolution; scale:', mergein_scale)
    else:
        mergein_scale = 1
    """
    # always rescale to input res for now
    mergein_scale = input_resolution[0] / img.shape[0]

    imageandpatchs = ImageandPatchs('', '', patchset, img, mergein_scale)
    whole_estimate_resized = cv2.resize(whole_estimate, (round(img.shape[1] * mergein_scale),
                                                         round(img.shape[0] * mergein_scale)),
                                        interpolation=cv2.INTER_CUBIC)
    imageandpatchs.set_base_estimate(whole_estimate_resized.copy())
    imageandpatchs.set_updated_estimate(whole_estimate_resized.copy())

    print('Resulting depthmap resolution will be :', whole_estimate_resized.shape[:2])
    print('patches to process: ' + str(len(imageandpatchs)))

    # Enumerate through all patches, generate their estimations and refining the base estimate.
    for patch_ind in range(len(imageandpatchs)):

        # Get patch information
        patch = imageandpatchs[patch_ind]  # patch object
        patch_rgb = patch['patch_rgb']  # rgb patch
        patch_whole_estimate_base = patch['patch_whole_estimate_base']  # corresponding patch from base
        rect = patch['rect']  # patch size and location
        patch_id = patch['id']  # patch ID
        org_size = patch_whole_estimate_base.shape  # the original size from the unscaled input
        print('\t processing patch', patch_ind, '/', len(imageandpatchs) - 1, '|', rect)

        # We apply double estimation for patches. The high resolution value is fixed to twice the receptive
        # field size of the network for patches to accelerate the process.
        patch_estimation = doubleestimate(patch_rgb, net_receptive_field_size, patch_netsize, pix2pixsize, model,
                                          model_type, pix2pixmodel)
        patch_estimation = cv2.resize(patch_estimation, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
        patch_whole_estimate_base = cv2.resize(patch_whole_estimate_base, (pix2pixsize, pix2pixsize),
                                               interpolation=cv2.INTER_CUBIC)

        # Merging the patch estimation into the base estimate using our merge network:
        # We feed the patch estimation and the same region from the updated base estimate to the merge network
        # to generate the target estimate for the corresponding region.
        pix2pixmodel.set_input(patch_whole_estimate_base, patch_estimation)

        # Run merging network
        pix2pixmodel.test()
        visuals = pix2pixmodel.get_current_visuals()

        prediction_mapped = visuals['fake_B']
        prediction_mapped = (prediction_mapped + 1) / 2
        prediction_mapped = prediction_mapped.squeeze().cpu().numpy()

        mapped = prediction_mapped

        # We use a simple linear polynomial to make sure the result of the merge network would match the values of
        # base estimate
        p_coef = np.polyfit(mapped.reshape(-1), patch_whole_estimate_base.reshape(-1), deg=1)
        merged = np.polyval(p_coef, mapped.reshape(-1)).reshape(mapped.shape)

        merged = cv2.resize(merged, (org_size[1], org_size[0]), interpolation=cv2.INTER_CUBIC)

        # Get patch size and location
        w1 = rect[0]
        h1 = rect[1]
        w2 = w1 + rect[2]
        h2 = h1 + rect[3]

        # To speed up the implementation, we only generate the Gaussian mask once with a sufficiently large size
        # and resize it to our needed size while merging the patches.
        if mask.shape != org_size:
            mask = cv2.resize(mask_org, (org_size[1], org_size[0]), interpolation=cv2.INTER_LINEAR)

        tobemergedto = imageandpatchs.estimation_updated_image

        # Update the whole estimation:
        # We use a simple Gaussian mask to blend the merged patch region with the base estimate to ensure seamless
        # blending at the boundaries of the patch region.
        tobemergedto[h1:h2, w1:w2] = np.multiply(tobemergedto[h1:h2, w1:w2], 1 - mask) + np.multiply(merged, mask)
        imageandpatchs.set_updated_estimate(tobemergedto)

    # output
    return cv2.resize(imageandpatchs.estimation_updated_image, (input_resolution[1], input_resolution[0]),
                      interpolation=cv2.INTER_CUBIC)


def generatemask(size):
    # Generates a Guassian mask
    mask = np.zeros(size, dtype=np.float32)
    sigma = int(size[0] / 16)
    k_size = int(2 * np.ceil(2 * int(size[0] / 16)) + 1)
    mask[int(0.15 * size[0]):size[0] - int(0.15 * size[0]), int(0.15 * size[1]): size[1] - int(0.15 * size[1])] = 1
    mask = cv2.GaussianBlur(mask, (int(k_size), int(k_size)), sigma)
    mask = (mask - mask.min()) / (mask.max() - mask.min())
    mask = mask.astype(np.float32)
    return mask


def rgb2gray(rgb):
    # Converts rgb to gray
    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])


def resizewithpool(img, size):
    i_size = img.shape[0]
    n = int(np.floor(i_size / size))

    out = skimage.measure.block_reduce(img, (n, n), np.max)
    return out


def calculateprocessingres(img, basesize, confidence=0.1, scale_threshold=3, whole_size_threshold=3000):
    # Returns the R_x resolution described in section 5 of the main paper.

    # Parameters:
    #    img :input rgb image
    #    basesize : size the dilation kernel which is equal to receptive field of the network.
    #    confidence: value of x in R_x; allowed percentage of pixels that are not getting any contextual cue.
    #    scale_threshold: maximum allowed upscaling on the input image ; it has been set to 3.
    #    whole_size_threshold: maximum allowed resolution. (R_max from section 6 of the main paper)

    # Returns:
    #    outputsize_scale*speed_scale :The computed R_x resolution
    #    patch_scale: K parameter from section 6 of the paper

    # speed scale parameter is to process every image in a smaller size to accelerate the R_x resolution search
    speed_scale = 32
    image_dim = int(min(img.shape[0:2]))

    gray = rgb2gray(img)
    grad = np.abs(cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)) + np.abs(cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3))
    grad = cv2.resize(grad, (image_dim, image_dim), cv2.INTER_AREA)

    # thresholding the gradient map to generate the edge-map as a proxy of the contextual cues
    m = grad.min()
    M = grad.max()
    middle = m + (0.4 * (M - m))
    grad[grad < middle] = 0
    grad[grad >= middle] = 1

    # dilation kernel with size of the receptive field
    kernel = np.ones((int(basesize / speed_scale), int(basesize / speed_scale)), float)
    # dilation kernel with size of the a quarter of receptive field used to compute k
    # as described in section 6 of main paper
    kernel2 = np.ones((int(basesize / (4 * speed_scale)), int(basesize / (4 * speed_scale))), float)

    # Output resolution limit set by the whole_size_threshold and scale_threshold.
    threshold = min(whole_size_threshold, scale_threshold * max(img.shape[:2]))

    outputsize_scale = basesize / speed_scale
    for p_size in range(int(basesize / speed_scale), int(threshold / speed_scale), int(basesize / (2 * speed_scale))):
        grad_resized = resizewithpool(grad, p_size)
        grad_resized = cv2.resize(grad_resized, (p_size, p_size), cv2.INTER_NEAREST)
        grad_resized[grad_resized >= 0.5] = 1
        grad_resized[grad_resized < 0.5] = 0

        dilated = cv2.dilate(grad_resized, kernel, iterations=1)
        meanvalue = (1 - dilated).mean()
        if meanvalue > confidence:
            break
        else:
            outputsize_scale = p_size

    grad_region = cv2.dilate(grad_resized, kernel2, iterations=1)
    patch_scale = grad_region.mean()

    return int(outputsize_scale * speed_scale), patch_scale


# Generate a double-input depth estimation
def doubleestimate(img, size1, size2, pix2pixsize, model, net_type, pix2pixmodel):
    # Generate the low resolution estimation
    estimate1 = singleestimate(img, size1, model, net_type)
    # Resize to the inference size of merge network.
    estimate1 = cv2.resize(estimate1, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)

    # Generate the high resolution estimation
    estimate2 = singleestimate(img, size2, model, net_type)
    # Resize to the inference size of merge network.
    estimate2 = cv2.resize(estimate2, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)

    # Inference on the merge model
    pix2pixmodel.set_input(estimate1, estimate2)
    pix2pixmodel.test()
    visuals = pix2pixmodel.get_current_visuals()
    prediction_mapped = visuals['fake_B']
    prediction_mapped = (prediction_mapped + 1) / 2
    prediction_mapped = (prediction_mapped - torch.min(prediction_mapped)) / (
            torch.max(prediction_mapped) - torch.min(prediction_mapped))
    prediction_mapped = prediction_mapped.squeeze().cpu().numpy()

    return prediction_mapped


# Generate a single-input depth estimation
def singleestimate(img, msize, model, net_type):
    if net_type == 0:
        return estimateleres(img, model, msize, msize)
    elif net_type == 10:
        return estimatemarigold(img, model, msize, msize)
    elif net_type == 11:
        return estimatedepthanything(img, model, msize, msize)
    elif net_type in [12, 13, 14]:
        return estimatedepthanything_v2(img, model, msize, msize)
    elif net_type >= 7:
        # np to PIL
        return estimatezoedepth(Image.fromarray(np.uint8(img * 255)).convert('RGB'), model, msize, msize)
    else:
        return estimatemidasBoost(img, model, msize, msize)


# Generating local patches to perform the local refinement described in section 6 of the main paper.
def generatepatchs(img, base_size, factor):
    # Compute the gradients as a proxy of the contextual cues.
    img_gray = rgb2gray(img)
    whole_grad = np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 0, 1, ksize=3)) + \
                 np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 1, 0, ksize=3))

    threshold = whole_grad[whole_grad > 0].mean()
    whole_grad[whole_grad < threshold] = 0

    # We use the integral image to speed-up the evaluation of the amount of gradients for each patch.
    gf = whole_grad.sum() / len(whole_grad.reshape(-1))
    grad_integral_image = cv2.integral(whole_grad)

    # Variables are selected such that the initial patch size would be the receptive field size
    # and the stride is set to 1/3 of the receptive field size.
    blsize = int(round(base_size / 2))
    stride = int(round(blsize * 0.75))

    # Get initial Grid
    patch_bound_list = applyGridpatch(blsize, stride, img, [0, 0, 0, 0])

    # Refine initial Grid of patches by discarding the flat (in terms of gradients of the rgb image) ones. Refine
    # each patch size to ensure that there will be enough depth cues for the network to generate a consistent depth map.
    print("Selecting patches ...")
    patch_bound_list = adaptiveselection(grad_integral_image, patch_bound_list, gf, factor)

    # Sort the patch list to make sure the merging operation will be done with the correct order: starting from biggest
    # patch
    patchset = sorted(patch_bound_list.items(), key=lambda x: getitem(x[1], 'size'), reverse=True)
    return patchset


def applyGridpatch(blsize, stride, img, box):
    # Extract a simple grid patch.
    counter1 = 0
    patch_bound_list = {}
    for k in range(blsize, img.shape[1] - blsize, stride):
        for j in range(blsize, img.shape[0] - blsize, stride):
            patch_bound_list[str(counter1)] = {}
            patchbounds = [j - blsize, k - blsize, j - blsize + 2 * blsize, k - blsize + 2 * blsize]
            patch_bound = [box[0] + patchbounds[1], box[1] + patchbounds[0], patchbounds[3] - patchbounds[1],
                           patchbounds[2] - patchbounds[0]]
            patch_bound_list[str(counter1)]['rect'] = patch_bound
            patch_bound_list[str(counter1)]['size'] = patch_bound[2]
            counter1 = counter1 + 1
    return patch_bound_list


# Adaptively select patches
def adaptiveselection(integral_grad, patch_bound_list, gf, factor):
    patchlist = {}
    count = 0
    height, width = integral_grad.shape

    search_step = int(32 / factor)

    # Go through all patches
    for c in range(len(patch_bound_list)):
        # Get patch
        bbox = patch_bound_list[str(c)]['rect']

        # Compute the amount of gradients present in the patch from the integral image.
        cgf = getGF_fromintegral(integral_grad, bbox) / (bbox[2] * bbox[3])

        # Check if patching is beneficial by comparing the gradient density of the patch to
        # the gradient density of the whole image
        if cgf >= gf:
            bbox_test = bbox.copy()
            patchlist[str(count)] = {}

            # Enlarge each patch until the gradient density of the patch is equal
            # to the whole image gradient density
            while True:

                bbox_test[0] = bbox_test[0] - int(search_step / 2)
                bbox_test[1] = bbox_test[1] - int(search_step / 2)

                bbox_test[2] = bbox_test[2] + search_step
                bbox_test[3] = bbox_test[3] + search_step

                # Check if we are still within the image
                if bbox_test[0] < 0 or bbox_test[1] < 0 or bbox_test[1] + bbox_test[3] >= height \
                        or bbox_test[0] + bbox_test[2] >= width:
                    break

                # Compare gradient density
                cgf = getGF_fromintegral(integral_grad, bbox_test) / (bbox_test[2] * bbox_test[3])
                if cgf < gf:
                    break
                bbox = bbox_test.copy()

            # Add patch to selected patches
            patchlist[str(count)]['rect'] = bbox
            patchlist[str(count)]['size'] = bbox[2]
            count = count + 1

    # Return selected patches
    return patchlist


def getGF_fromintegral(integralimage, rect):
    # Computes the gradient density of a given patch from the gradient integral image.
    x1 = rect[1]
    x2 = rect[1] + rect[3]
    y1 = rect[0]
    y2 = rect[0] + rect[2]
    value = integralimage[x2, y2] - integralimage[x1, y2] - integralimage[x2, y1] + integralimage[x1, y1]
    return value


def estimatemidasBoost(img, model, w, h):
    # init transform
    transform = Compose(
        [
            Resize(
                w,
                h,
                resize_target=None,
                keep_aspect_ratio=True,
                ensure_multiple_of=32,
                resize_method="upper_bound",
                image_interpolation_method=cv2.INTER_CUBIC,
            ),
            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            PrepareForNet(),
        ]
    )

    # transform input
    img_input = transform({"image": img})["image"]

    # compute
    with torch.no_grad():
        sample = torch.from_numpy(img_input).to(depthmap_device).unsqueeze(0)
        if depthmap_device == torch.device("cuda"):
            sample = sample.to(memory_format=torch.channels_last)
        prediction = model.forward(sample)

    prediction = prediction.squeeze().cpu().numpy()
    prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)

    # normalization
    depth_min = prediction.min()
    depth_max = prediction.max()

    if depth_max - depth_min > np.finfo("float").eps:
        prediction = (prediction - depth_min) / (depth_max - depth_min)
    else:
        prediction = 0

    return prediction