diff --git a/depth_anything_v2/dpt.py b/depth_anything_v2/dpt.py index c221f7d..1e8040c 100644 --- a/depth_anything_v2/dpt.py +++ b/depth_anything_v2/dpt.py @@ -1,13 +1,9 @@ -#import cv2 import torch import torch.nn as nn import torch.nn.functional as F -from torchvision.transforms import Compose from .dinov2 import DINOv2 from .util.blocks import FeatureFusionBlock, _make_scratch -from .util.transform import Resize, NormalizeImage, PrepareForNet - def _make_fusion_block(features, use_bn, size=None): return FeatureFusionBlock( @@ -20,7 +16,6 @@ def _make_fusion_block(features, use_bn, size=None): size=size, ) - class ConvBlock(nn.Module): def __init__(self, in_feature, out_feature): super().__init__() @@ -34,7 +29,6 @@ class ConvBlock(nn.Module): def forward(self, x): return self.conv_block(x) - class DPTHead(nn.Module): def __init__( self, @@ -199,41 +193,4 @@ class DepthAnythingV2(nn.Module): depth = self.depth_head(features, patch_h, patch_w) depth = F.relu(depth) - return depth.squeeze(1) - - @torch.no_grad() - def infer_image(self, raw_image, input_size=518): - #image, (h, w) = self.image2tensor(raw_image, input_size) - - depth = self.forward(raw_image) - - #depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0] - - return depth - - # def image2tensor(self, raw_image, input_size=518): - # transform = Compose([ - # Resize( - # width=input_size, - # height=input_size, - # 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(), - # ]) - - # h, w = raw_image.shape[:2] - - # image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0 - - # image = transform({'image': image})['image'] - # image = torch.from_numpy(image).unsqueeze(0) - - # DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu' - # image = image.to(DEVICE) - - # return image, (h, w) + return depth.squeeze(1) \ No newline at end of file diff --git a/depth_anything_v2/util/transform.py b/depth_anything_v2/util/transform.py deleted file mode 100644 index b14aacd..0000000 --- a/depth_anything_v2/util/transform.py +++ /dev/null @@ -1,158 +0,0 @@ -import numpy as np -import cv2 - - -class Resize(object): - """Resize sample to given size (width, height). - """ - - def __init__( - self, - width, - height, - resize_target=True, - keep_aspect_ratio=False, - ensure_multiple_of=1, - resize_method="lower_bound", - image_interpolation_method=cv2.INTER_AREA, - ): - """Init. - - Args: - width (int): desired output width - height (int): desired output height - resize_target (bool, optional): - True: Resize the full sample (image, mask, target). - False: Resize image only. - Defaults to True. - keep_aspect_ratio (bool, optional): - True: Keep the aspect ratio of the input sample. - Output sample might not have the given width and height, and - resize behaviour depends on the parameter 'resize_method'. - Defaults to False. - ensure_multiple_of (int, optional): - Output width and height is constrained to be multiple of this parameter. - Defaults to 1. - resize_method (str, optional): - "lower_bound": Output will be at least as large as the given size. - "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.) - "minimal": Scale as least as possible. (Output size might be smaller than given size.) - Defaults to "lower_bound". - """ - self.__width = width - self.__height = height - - self.__resize_target = resize_target - self.__keep_aspect_ratio = keep_aspect_ratio - self.__multiple_of = ensure_multiple_of - self.__resize_method = resize_method - self.__image_interpolation_method = image_interpolation_method - - def constrain_to_multiple_of(self, x, min_val=0, max_val=None): - y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) - - if max_val is not None and y > max_val: - y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int) - - if y < min_val: - y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int) - - return y - - def get_size(self, width, height): - # determine new height and width - scale_height = self.__height / height - scale_width = self.__width / width - - if self.__keep_aspect_ratio: - if self.__resize_method == "lower_bound": - # scale such that output size is lower bound - if scale_width > scale_height: - # fit width - scale_height = scale_width - else: - # fit height - scale_width = scale_height - elif self.__resize_method == "upper_bound": - # scale such that output size is upper bound - if scale_width < scale_height: - # fit width - scale_height = scale_width - else: - # fit height - scale_width = scale_height - elif self.__resize_method == "minimal": - # scale as least as possbile - if abs(1 - scale_width) < abs(1 - scale_height): - # fit width - scale_height = scale_width - else: - # fit height - scale_width = scale_height - else: - raise ValueError(f"resize_method {self.__resize_method} not implemented") - - if self.__resize_method == "lower_bound": - new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height) - new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width) - elif self.__resize_method == "upper_bound": - new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height) - new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width) - elif self.__resize_method == "minimal": - new_height = self.constrain_to_multiple_of(scale_height * height) - new_width = self.constrain_to_multiple_of(scale_width * width) - else: - raise ValueError(f"resize_method {self.__resize_method} not implemented") - - return (new_width, new_height) - - def __call__(self, sample): - width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0]) - - # resize sample - sample["image"] = cv2.resize(sample["image"], (width, height), interpolation=self.__image_interpolation_method) - - if self.__resize_target: - if "depth" in sample: - sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST) - - if "mask" in sample: - sample["mask"] = cv2.resize(sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST) - - return sample - - -class NormalizeImage(object): - """Normlize image by given mean and std. - """ - - def __init__(self, mean, std): - self.__mean = mean - self.__std = std - - def __call__(self, sample): - sample["image"] = (sample["image"] - self.__mean) / self.__std - - return sample - - -class PrepareForNet(object): - """Prepare sample for usage as network input. - """ - - def __init__(self): - pass - - def __call__(self, sample): - image = np.transpose(sample["image"], (2, 0, 1)) - sample["image"] = np.ascontiguousarray(image).astype(np.float32) - - if "depth" in sample: - depth = sample["depth"].astype(np.float32) - sample["depth"] = np.ascontiguousarray(depth) - - if "mask" in sample: - sample["mask"] = sample["mask"].astype(np.float32) - sample["mask"] = np.ascontiguousarray(sample["mask"]) - - return sample \ No newline at end of file diff --git a/nodes.py b/nodes.py index 1695fce..e3cf7b9 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,6 @@ import torch import torch.nn.functional as F from torchvision import transforms -from huggingface_hub import hf_hub_download import os from contextlib import nullcontext @@ -144,7 +143,7 @@ https://depth-anything-v2.github.io autocast_condition = (dtype != torch.float32) and not mm.is_device_mps(device) with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): for img in normalized_images: - depth = model.infer_image(img.unsqueeze(0).to(device)) + depth = model(img.unsqueeze(0).to(device)) depth = (depth - depth.min()) / (depth.max() - depth.min()) out.append(depth.cpu()) pbar.update(1)