Duplicate from ybelkada/FocusOnDepth

Co-authored-by: Younes Belkada <ybelkada@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,28 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,10 @@
+---
+title: FocusOnDepth
+emoji: 🐨
+colorFrom: pink
+colorTo: pink
+sdk: gradio
+app_file: app.py
+pinned: false
+duplicated_from: ybelkada/FocusOnDepth
+---

app.py

@@ -0,0 +1,75 @@
+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)
+
+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)
+
+@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)]
+
+description = """
+<center>
+Can a single model predict both segmentation and depth estimation? At least, if the segmentation is constrained for a single class, the answer is yes! <br>
+In this project, we use a DPT model to predict the depth and the segmentation mask of the class human, of an image. This model could be potentially used for an autofocus application where you would need the segmentation mask of the humans on the picture, as well as the depth estimation of the scene<br>
+Credits also to <div style='text-align: center;'><a href='https://github.com/antocad' target='_blank'>@antocad</a> !
+</center>
+"""
+title="""
+FocusOnDepth - A single DPT encoder for Dense Prediction Tasks
+"""
+css = """
+"""
+article = """
+<center>
+Example image taken from <a href="https://www.flickr.com/photos/17423713@N03/29129350066">here</a>. The image is free to share and use. <br>
+</center>
+<div style='text-align: center;'><a href='https://github.com/isl-org/DPT' target='_blank'>Original Paper</a> | <a href='https://github.com/antocad/FocusOnDepth' target='_blank'>Extended Version</a></div>
+"""
+
+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:"),
+    ],
+    examples=['example_image.jpg'],
+    description=description,
+    title=title,
+    css=css,
+    article=article
+)
+iface.launch(enable_queue=True, cache_examples=True)

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

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+timm
+numpy
+pillow
+transformers
+einops