add first files

app.py

@@ -0,0 +1,53 @@
+import torch
+import gradio as gr
+import numpy as np
+
+import requests
+from PIL import Image
+from io import BytesIO
+from torchvision import transforms
+
+from transformers import AutoConfig, AutoModel
+from transformers import AutoModel
+
+from focusondepth.model_config import FocusOnDepthConfig
+from focusondepth.model_definition import FocusOnDepth
+
+AutoConfig.register("focusondepth", FocusOnDepthConfig)
+AutoModel.register(FocusOnDepthConfig, FocusOnDepth)
+
+original_image_cache = {}
+transform = transforms.Compose([
+    transforms.Resize((384, 384)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+])
+model = AutoModel.from_pretrained('ybelkada/focusondepth', trust_remote_code=True)
+model.load_state_dict(torch.load('./focusondepth/FocusOnDepth_vit_base_patch16_384.p', map_location=torch.device('cpu'))['model_state_dict'])
+
+@torch.no_grad()
+def inference(input_image):
+    global model, transform
+    
+    model.eval()
+    input_image = Image.fromarray(input_image)
+    original_size = input_image.size
+    tensor_image = transform(input_image)
+    
+    depth, segmentation = model(tensor_image.unsqueeze(0))
+    depth = 1-depth
+
+    depth = transforms.ToPILImage()(depth[0, :])
+    segmentation = transforms.ToPILImage()(segmentation.argmax(dim=1).float())
+
+    return [depth.resize(original_size, resample=Image.BICUBIC), segmentation.resize(original_size, resample=Image.NEAREST)]
+
+iface = gr.Interface(
+    fn=inference, 
+    inputs=gr.inputs.Image(label="Input Image"), 
+    outputs = [
+        gr.outputs.Image(label="Depth Map:"),
+        gr.outputs.Image(label="Segmentation Map:"),
+    ],
+)
+iface.launch()

focusondepth/fusion.py

@@ -0,0 +1,41 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+class ResidualConvUnit(nn.Module):
+    def __init__(self, features):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.conv2 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """Forward pass.
+        Args:
+            x (tensor): input
+        Returns:
+            tensor: output
+        """
+        out = self.relu(x)
+        out = self.conv1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        return out + x
+
+class Fusion(nn.Module):
+    def __init__(self, resample_dim):
+        super(Fusion, self).__init__()
+        self.res_conv1 = ResidualConvUnit(resample_dim)
+        self.res_conv2 = ResidualConvUnit(resample_dim)
+        
+    def forward(self, x, previous_stage=None):
+        if previous_stage == None:
+            previous_stage = torch.zeros_like(x)
+        output_stage1 = self.res_conv1(x)
+        output_stage1 += previous_stage
+        output_stage2 = self.res_conv2(output_stage1)
+        output_stage2 = nn.functional.interpolate(output_stage2, scale_factor=2, mode="bilinear", align_corners=True)
+        return output_stage2

focusondepth/head.py

@@ -0,0 +1,50 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+class Interpolate(nn.Module):
+    def __init__(self, scale_factor, mode, align_corners=False):
+        super(Interpolate, self).__init__()
+        self.interp = nn.functional.interpolate
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        x = self.interp(
+            x,
+            scale_factor=self.scale_factor,
+            mode=self.mode,
+            align_corners=self.align_corners)
+        return x
+
+class HeadDepth(nn.Module):
+    def __init__(self, features):
+        super(HeadDepth, self).__init__()
+        self.head = nn.Sequential(
+            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
+            # nn.ReLU()
+            nn.Sigmoid()
+        )
+    def forward(self, x):
+        x = self.head(x)
+        # x = (x - x.min())/(x.max()-x.min() + 1e-15)
+        return x
+
+class HeadSeg(nn.Module):
+    def __init__(self, features, nclasses=2):
+        super(HeadSeg, self).__init__()
+        self.head = nn.Sequential(
+            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, nclasses, kernel_size=1, stride=1, padding=0)
+        )
+    def forward(self, x):
+        x = self.head(x)
+        return x

focusondepth/model_config.py

@@ -0,0 +1,45 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+class FocusOnDepthConfig(PretrainedConfig):
+    model_type = "focusondepth"
+
+    def __init__(
+        self,
+        image_size         = (3, 384, 384),
+        patch_size         = 16,
+        emb_dim            = 768,
+        resample_dim       = 256,
+        read               = 'projection',
+        num_layers_encoder = 24,
+        hooks              = [2, 5, 8, 11],
+        reassemble_s       = [4, 8, 16, 32],
+        transformer_dropout= 0,
+        nclasses           = 2,
+        type_               = "full",
+        model_timm         = "vit_base_patch16_384",
+        **kwargs,
+    ):
+        if type_ not in ["full", "depth", "segmentation"]:
+            raise ValueError(f"`type_` must be 'full' or depth' or 'segmentation, got {type_}.")
+        if read not in ["ignore", "add", "projection"]:
+            raise ValueError(f"`read` must be '', 'ignore' or 'add' or 'projection, got {read}.")
+        if image_size[0] != 3 and image_size[1] != 384 and image_size[2] != 384:
+            raise ValueError(f"`image_size` must be 3, 384, 384, got {image_size}.")
+        if patch_size != 16:
+            raise ValueError(f"`patch_size` must be 16, got {patch_size}.")
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.emb_dim = emb_dim
+        self.resample_dim = resample_dim
+        self.read = read
+        self.num_layers_encoder = num_layers_encoder
+        self.hooks = hooks
+        self.reassemble_s = reassemble_s
+        self.transformer_dropout = transformer_dropout
+        self.nclasses = nclasses
+        self.type_ = type_
+        self.model_timm = model_timm
+        super().__init__(**kwargs)

focusondepth/model_definition.py

@@ -0,0 +1,68 @@
+from transformers import PreTrainedModel
+import timm
+import torch.nn as nn
+import numpy as np
+
+from .model_config import FocusOnDepthConfig
+from .reassemble import Reassemble
+from .fusion import Fusion
+from .head import HeadDepth, HeadSeg
+
+
+class FocusOnDepth(PreTrainedModel):
+    config_class = FocusOnDepthConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer_encoders = timm.create_model(config.model_timm, pretrained=True)
+        self.type_ = config.type_
+
+        #Register hooks
+        self.activation = {}
+        self.hooks = config.hooks
+        self._get_layers_from_hooks(self.hooks)
+
+        #Reassembles Fusion
+        self.reassembles = []
+        self.fusions = []
+        for s in config.reassemble_s:
+            self.reassembles.append(Reassemble(config.image_size, config.read, config.patch_size, s, config.emb_dim, config.resample_dim))
+            self.fusions.append(Fusion(config.resample_dim))
+        self.reassembles = nn.ModuleList(self.reassembles)
+        self.fusions = nn.ModuleList(self.fusions)
+
+        #Head
+        if self.type_ == "full":
+            self.head_depth = HeadDepth(config.resample_dim)
+            self.head_segmentation = HeadSeg(config.resample_dim, nclasses=config.nclasses)
+        elif self.type_ == "depth":
+            self.head_depth = HeadDepth(config.resample_dim)
+            self.head_segmentation = None
+        else:
+            self.head_depth = None
+            self.head_segmentation = HeadSeg(config.resample_dim, nclasses=config.nclasses)
+
+    def forward(self, img):
+        _ = self.transformer_encoders(img)
+        previous_stage = None
+        for i in np.arange(len(self.fusions)-1, -1, -1):
+            hook_to_take = 't'+str(self.hooks[i])
+            activation_result = self.activation[hook_to_take]
+            reassemble_result = self.reassembles[i](activation_result)
+            fusion_result = self.fusions[i](reassemble_result, previous_stage)
+            previous_stage = fusion_result
+        out_depth = None
+        out_segmentation = None
+        if self.head_depth != None:
+            out_depth = self.head_depth(previous_stage)
+        if self.head_segmentation != None:
+            out_segmentation = self.head_segmentation(previous_stage)
+        return out_depth, out_segmentation
+
+    def _get_layers_from_hooks(self, hooks):
+        def get_activation(name):
+            def hook(model, input, output):
+                self.activation[name] = output
+            return hook
+        for h in hooks:
+            self.transformer_encoders.blocks[h].register_forward_hook(get_activation('t'+str(h)))

focusondepth/reassemble.py

@@ -0,0 +1,115 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+class Read_ignore(nn.Module):
+    def __init__(self, start_index=1):
+        super(Read_ignore, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        return x[:, self.start_index:]
+
+
+class Read_add(nn.Module):
+    def __init__(self, start_index=1):
+        super(Read_add, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        if self.start_index == 2:
+            readout = (x[:, 0] + x[:, 1]) / 2
+        else:
+            readout = x[:, 0]
+        return x[:, self.start_index :] + readout.unsqueeze(1)
+
+
+class Read_projection(nn.Module):
+    def __init__(self, in_features, start_index=1):
+        super(Read_projection, self).__init__()
+        self.start_index = start_index
+        self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU())
+
+    def forward(self, x):
+        readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :])
+        features = torch.cat((x[:, self.start_index :], readout), -1)
+        return self.project(features)
+
+class MyConvTranspose2d(nn.Module):
+    def __init__(self, conv, output_size):
+        super(MyConvTranspose2d, self).__init__()
+        self.output_size = output_size
+        self.conv = conv
+
+    def forward(self, x):
+        x = self.conv(x, output_size=self.output_size)
+        return x
+
+class Resample(nn.Module):
+    def __init__(self, p, s, h, emb_dim, resample_dim):
+        super(Resample, self).__init__()
+        assert (s in [4, 8, 16, 32]), "s must be in [0.5, 4, 8, 16, 32]"
+        self.conv1 = nn.Conv2d(emb_dim, resample_dim, kernel_size=1, stride=1, padding=0)
+        if s == 4:
+            self.conv2 = nn.ConvTranspose2d(resample_dim,
+                                resample_dim,
+                                kernel_size=4,
+                                stride=4,
+                                padding=0,
+                                bias=True,
+                                dilation=1,
+                                groups=1)
+        elif s == 8:
+            self.conv2 = nn.ConvTranspose2d(resample_dim,
+                                resample_dim,
+                                kernel_size=2,
+                                stride=2,
+                                padding=0,
+                                bias=True,
+                                dilation=1,
+                                groups=1)
+        elif s == 16:
+            self.conv2 = nn.Identity()
+        else:
+            self.conv2 = nn.Conv2d(resample_dim, resample_dim, kernel_size=2,stride=2, padding=0, bias=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+class Reassemble(nn.Module):
+    def __init__(self, image_size, read, p, s, emb_dim, resample_dim):
+        """
+        p = patch size
+        s = coefficient resample
+        emb_dim <=> D (in the paper)
+        resample_dim <=> ^D (in the paper)
+        read : {"ignore", "add", "projection"}
+        """
+        super(Reassemble, self).__init__()
+        channels, image_height, image_width = image_size
+
+        #Read
+        self.read = Read_ignore()
+        if read == 'add':
+            self.read = Read_add()
+        elif read == 'projection':
+            self.read = Read_projection(emb_dim)
+
+        #Concat after read
+        self.concat = Rearrange('b (h w) c -> b c h w',
+                                c=emb_dim,
+                                h=(image_height // p),
+                                w=(image_width // p))
+
+        #Projection + Resample
+        self.resample = Resample(p, s, image_height, emb_dim, resample_dim)
+
+    def forward(self, x):
+        x = self.read(x)
+        x = self.concat(x)
+        x = self.resample(x)
+        return x