[preprocessors] TEED preprocessor

README.md

@@ -60,4 +60,5 @@ This library uses inference and models from various authors. If you think you mu
 
 * [LineArt](https://github.com/carolineec/informative-drawings)
 * [Depth Anything V2](https://depth-anything-v2.github.io)
+* [TEED](https://github.com/xavysp/TEED)
 * [ZoeDepth](https://github.com/isl-org/ZoeDepth)

src/image_gen_aux/__init__.py

@@ -49,7 +49,7 @@
 else:
     _import_structure["upscalers"].extend(["UpscaleWithModel"])
 
-    _import_structure["preprocessors"].extend(["LineArtPreprocessor", "DepthPreprocessor"])
+    _import_structure["preprocessors"].extend(["LineArtPreprocessor", "DepthPreprocessor", "TeedPreprocessor"])
 
 if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
     try:
@@ -58,7 +58,7 @@
     except OptionalDependencyNotAvailable:
         ...
     else:
-        from .preprocessors import LineArtPreprocessor
+        from .preprocessors import LineArtPreprocessor, TeedPreprocessor
         from .upscalers import UpscaleWithModel
 
     try:

src/image_gen_aux/image_processor.py

@@ -34,7 +34,7 @@ class ImageMixin:
     """
 
     def convert_image_to_tensor(
-        self, image: Union[PIL.Image.Image, np.ndarray, List[PIL.Image.Image]]
+        self, image: Union[PIL.Image.Image, np.ndarray, List[PIL.Image.Image]], normalize: bool = True
     ) -> torch.Tensor:
         """
         Convert a PIL image or a NumPy array to a PyTorch tensor.
@@ -53,7 +53,7 @@ def convert_image_to_tensor(
                     if not isinstance(single_image, PIL.Image.Image):
                         raise ValueError("All images in the list must be Pillow images.")
 
-            image = self.pil_to_numpy(image)
+            image = self.pil_to_numpy(image, normalize)
 
         return self.numpy_to_pt(image)
 
@@ -79,7 +79,7 @@ def post_process_image(self, image: torch.Tensor, return_type: str):
         return image
 
     @staticmethod
-    def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
+    def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image], normalize: bool = True) -> np.ndarray:
         """
         Convert a PIL image or a list of PIL images to NumPy arrays.
 
@@ -91,7 +91,12 @@ def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.nd
         """
         if not isinstance(images, list):
             images = [images]
-        images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+
+        if normalize:
+            images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+        else:
+            images = [np.array(image).astype(np.float32) for image in images]
+
         images = np.stack(images, axis=0)
 
         return images
@@ -109,8 +114,8 @@ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
         """
         if images.ndim == 3:
             images = images[..., None]
+        images = torch.from_numpy(images.transpose(0, 3, 1, 2)).float()
 
-        images = torch.from_numpy(images.transpose(0, 3, 1, 2))
         return images
 
     @staticmethod

src/image_gen_aux/preprocessors/README.md

@@ -7,6 +7,8 @@ Preprocessors in this context refer to the application of machine learning model
 This is a list of the currently supported preprocessors.
 
 * [LineArtPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/lineart/README.md)
+* [DepthPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/depth/README.md)
+* [TeedPreprocessor](https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/preprocessors/teed/README.md)
 
 ## General preprocessor usage
 
@@ -39,6 +41,13 @@ This is the current list of safetensor checkpoints available on the Hub.
 
 |Preprocessor|Repository|Author|
 |---|---|---|
+|Depth Anything V2 Small|depth-anything/Depth-Anything-V2-Small-hf|<https://depth-anything-v2.github.io/>|
+|Depth Anything V2 Base|depth-anything/Depth-Anything-V2-Base-hf|<https://depth-anything-v2.github.io/>|
+|Depth Anything V2 Large|depth-anything/Depth-Anything-V2-Large-hf|<https://depth-anything-v2.github.io/>|
 |LineArt|OzzyGT/lineart|[Caroline Chan](https://github.com/carolineec)|
+|Teed|OzzyGT/teed|<https://github.com/xavysp/TEED>|
+|ZoeDepth NYU|Intel/zoedepth-nyu|<https://github.com/isl-org/ZoeDepth>|
+|ZoeDepth KITTI|Intel/zoedepth-kitti|<https://github.com/isl-org/ZoeDepth>|
+|ZoeDepth NYU and KITTI|Intel/zoedepth-nyu-kitti|<https://github.com/isl-org/ZoeDepth>|
 
 If you own the model and want us to change the repository to your name/organization please open an issue.

src/image_gen_aux/preprocessors/__init__.py

@@ -12,6 +12,7 @@
 _import_structure = {
     "lineart": [],
     "depth": [],
+    "teed": [],
 }
 
 
@@ -21,9 +22,8 @@
 except OptionalDependencyNotAvailable:
     ...
 else:
-    _import_structure["lineart"] = [
-        "LineArtPreprocessor",
-    ]
+    _import_structure["lineart"] = ["LineArtPreprocessor"]
+    _import_structure["teed"] = ["TeedPreprocessor"]
 
 try:
     if not (is_torch_available() and is_transformers_available()):
@@ -43,6 +43,7 @@
         ...
     else:
         from .lineart import LineArtPreprocessor
+        from .teed import TeedPreprocessor
 
     try:
         if not (is_torch_available() and is_transformers_available()):

src/image_gen_aux/preprocessors/teed/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Xavier Soria Poma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/image_gen_aux/preprocessors/teed/README.md

@@ -0,0 +1,26 @@
+# TEED: Tiny and Efficient Model for the Edge Detection Generalization
+
+Tiny and Efficient Edge Detector (TEED) is a light convolutional neural network with only 58K parameters, less than 0.2% of the state-of-the-art models. Training on the [BIPED](https://www.kaggle.com/datasets/xavysp/biped) dataset takes less than 30 minutes,
+with each epoch requiring less than 5 minutes. Our proposed model is easy to train and it quickly converges within very first few
+epochs, while the predicted edge-maps are crisp and of high quality, see image above.
+[This paper has been accepted by ICCV 2023-Workshop RCV](https://arxiv.org/abs/2308.06468).
+
+## Usage
+
+```python
+from image_gen_aux import TeedPreprocessor
+from image_gen_aux.utils import load_image
+
+input_image = load_image(
+    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/teed/20240922043215.png"
+)
+
+teed_preprocessor = TeedPreprocessor.from_pretrained("OzzyGT/teed").to("cuda")
+image = teed_preprocessor(input_image)[0]
+image.save("teed.png")
+```
+
+## Additional resources
+
+* [Project page](https://github.com/xavysp/TEED)
+* [Paper](https://arxiv.org/abs/2308.06468)

src/image_gen_aux/preprocessors/teed/__init__.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from ...utils import IMAGE_AUX_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {}
+
+try:
+    if not (is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    ...
+else:
+    _import_structure["teed_preprocessor"] = [
+        "TeedPreprocessor",
+    ]
+
+if TYPE_CHECKING or IMAGE_AUX_SLOW_IMPORT:
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        ...
+    else:
+        from .teed_preprocessor import (
+            TeedPreprocessor,
+        )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )