# Path: third_party/midas_blocks.py
# class FeatureFusionBlock(nn.Module):
#     """Feature fusion block.
#     """
# 
#     def __init__(self, features):
#         """Init.
#         Args:
#             features (int): number of features
#         """
#         super(FeatureFusionBlock, self).__init__()
# 
#         self.resConfUnit1 = ResidualConvUnit(features)
#         self.resConfUnit2 = ResidualConvUnit(features)
# 
#     def forward(self, *xs):
#         """Forward pass.
#         Returns:
#             tensor: output
#         """
#         output = xs[0]
# 
#         if len(xs) == 2:
#             output += self.resConfUnit1(xs[1])
# 
#         output = self.resConfUnit2(output)
# 
#         output = nn.functional.interpolate(
#             output, scale_factor=2, mode="bilinear", align_corners=True
#         )
# 
#         return output
# 
# class Interpolate(nn.Module):
#     """Interpolation module.
#     """
# 
#     def __init__(self, scale_factor, mode):
#         """Init.
#         Args:
#             scale_factor (float): scaling
#             mode (str): interpolation mode
#         """
#         super(Interpolate, self).__init__()
# 
#         self.interp = nn.functional.interpolate
#         self.scale_factor = scale_factor
#         self.mode = mode
# 
#     def forward(self, x):
#         """Forward pass.
#         Args:
#             x (tensor): input
#         Returns:
#             tensor: interpolated data
#         """
# 
#         x = self.interp(
#             x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False
#         )
# 
#         return x
# 
# def _make_encoder(features, use_pretrained):
#     pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
#     scratch = _make_scratch([256, 512, 1024, 2048], features)
# 
#     return pretrained, scratch


import torch
import torch.nn as nn
from .midas_blocks import FeatureFusionBlock, Interpolate, _make_encoder

        if "optimizer" in parameters:
            parameters = parameters["model"]

        self.load_state_dict(parameters)


class MidasNet_mod(BaseModel):
    def __init__(self, path=None, features=256, non_negative=True, normalize_input=False, resize=None, freeze_backbone=False, mask_branch=False):
        """Init.
        Args:
            path (str, optional): Path to saved model. Defaults to None.
            features (int, optional): Number of features. Defaults to 256.
            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
        """
        print("Loading weights: ", path)

        super(MidasNet_mod, self).__init__()

        use_pretrained = False if path is None else True

        self.pretrained, self.scratch = _make_encoder(features, use_pretrained)

        self.scratch.refinenet4 = FeatureFusionBlock(features)
        self.scratch.refinenet3 = FeatureFusionBlock(features)
        self.scratch.refinenet2 = FeatureFusionBlock(features)
        self.scratch.refinenet1 = FeatureFusionBlock(features)

        self.scratch.output_conv = nn.Sequential(
            nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),

# Path: third_party/midas_blocks.py
# class FeatureFusionBlock(nn.Module):
#     """Feature fusion block.
#     """
# 
#     def __init__(self, features):
#         """Init.
#         Args:
#             features (int): number of features
#         """
#         super(FeatureFusionBlock, self).__init__()
# 
#         self.resConfUnit1 = ResidualConvUnit(features)
#         self.resConfUnit2 = ResidualConvUnit(features)
# 
#     def forward(self, *xs):
#         """Forward pass.
#         Returns:
#             tensor: output
#         """
#         output = xs[0]
# 
#         if len(xs) == 2:
#             output += self.resConfUnit1(xs[1])
# 
#         output = self.resConfUnit2(output)
# 
#         output = nn.functional.interpolate(
#             output, scale_factor=2, mode="bilinear", align_corners=True
#         )
# 
#         return output
# 
# class Interpolate(nn.Module):
#     """Interpolation module.
#     """
# 
#     def __init__(self, scale_factor, mode):
#         """Init.
#         Args:
#             scale_factor (float): scaling
#             mode (str): interpolation mode
#         """
#         super(Interpolate, self).__init__()
# 
#         self.interp = nn.functional.interpolate
#         self.scale_factor = scale_factor
#         self.mode = mode
# 
#     def forward(self, x):
#         """Forward pass.
#         Args:
#             x (tensor): input
#         Returns:
#             tensor: interpolated data
#         """
# 
#         x = self.interp(
#             x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False
#         )
# 
#         return x
# 
# def _make_encoder(features, use_pretrained):
#     pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
#     scratch = _make_scratch([256, 512, 1024, 2048], features)
# 
#     return pretrained, scratch


# Path: third_party/MiDaS.py
import torch
import torch.nn as nn
from .midas_blocks import FeatureFusionBlock, Interpolate, _make_encoder


        if "optimizer" in parameters:
            parameters = parameters["model"]

        self.load_state_dict(parameters)


class MidasNet_mod(BaseModel):
    def __init__(self, path=None, features=256, non_negative=True, normalize_input=False, resize=None, freeze_backbone=False, mask_branch=False):
        """Init.
        Args:
            path (str, optional): Path to saved model. Defaults to None.
            features (int, optional): Number of features. Defaults to 256.
            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
        """
        print("Loading weights: ", path)

        super(MidasNet_mod, self).__init__()

        use_pretrained = False if path is None else True

        self.pretrained, self.scratch = _make_encoder(features, use_pretrained)

        self.scratch.refinenet4 = FeatureFusionBlock(features)
        self.scratch.refinenet3 = FeatureFusionBlock(features)
        self.scratch.refinenet2 = FeatureFusionBlock(features)
        self.scratch.refinenet1 = FeatureFusionBlock(features)

        self.scratch.output_conv = nn.Sequential(
            nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),

            Interpolate(scale_factor=2, mode="bilinear"),