diff --git a/.gitignore b/.gitignore index 26fd6c7..6cf3119 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .idea __pycache__ .DS_Store -logs -debug.py \ No newline at end of file +debug.py +build +dist \ No newline at end of file diff --git a/README.md b/README.md index 4ee4acb..e8961e6 100644 --- a/README.md +++ b/README.md @@ -1 +1,216 @@ ![torchlm-logo](docs/res/logo.png) + +
+ + + + + + +
+ + +## 🤗 Introduction +**torchlm** is a PyTorch landmarks-only library with **100+ data augmentations**, **training** and **inference**. **torchlm** is only focus on any landmarks detection, such as face landmarks, hand keypoints and body keypoints, etc. It provides **30+** native data augmentations and compatible with **80+** torchvision and albumations's transforms, no matter the input is a np.ndarray or a torch Tensor, **torchlm** will **automatically** be compatible with different data types through a **autodtype** wrapper. Further, in the future **torchlm** will add modules for **training** and **inference**. + +# 🆕 What's New + +* [2022/02/13]: Add **30+** native data augmentations and **bind** 80+ torchvision and albumations's transforms. + +## 🛠️ Usage + +### Requirements +* opencv-python-headless>=4.5.2 +* numpy>=1.14.4 +* torch>=1.6.0 +* torchvision>=0.9.0 +* albumentations>=1.1.0 + +### Installation +you can install **torchlm** directly from pip. +```shell +pip3 install torchlm +# install from specific pypi mirrors use '-i' +pip3 install torchlm -i https://pypi.org/simple/ +``` + +### Data Augmentation +**torchlm** provides 30+ native data augmentations for landmarks and is compatible with 80+ transforms from torchvision and albumations, no matter the input is a np.ndarray or a torch Tensor, torchlm will automatically be compatible with different data types through a autodtype wrapper. +* use native torchlm's transforms +```python +import torchlm +transform = torchlm.LandmarksCompose([ + # use native torchlm transforms + torchlm.LandmarksRandomScale(prob=0.5), + torchlm.LandmarksRandomTranslate(prob=0.5), + torchlm.LandmarksRandomShear(prob=0.5), + torchlm.LandmarksRandomMask(prob=0.5), + torchlm.LandmarksRandomBlur(kernel_range=(5, 25), prob=0.5), + torchlm.LandmarksRandomBrightness(prob=0.), + torchlm.LandmarksRandomRotate(40, prob=0.5, bins=8), + torchlm.LandmarksRandomCenterCrop((0.5, 1.0), (0.5, 1.0), prob=0.5), + torchlm.LandmarksResize((256, 256)), + torchlm.LandmarksNormalize(), + torchlm.LandmarksToTensor(), + torchlm.LandmarksToNumpy(), + torchlm.LandmarksUnNormalize() + ]) +``` +* **bind** torchvision and albumations's transform, using **torchlm.bind** +```python +import torchvision +import albumentations +import torchlm +transform = torchlm.LandmarksCompose([ + # use native torchlm transforms + torchlm.LandmarksRandomScale(prob=0.5), + # ... + # bind torchvision image only transforms + torchlm.bind(torchvision.transforms.GaussianBlur(kernel_size=(5, 25))), + torchlm.bind(torchvision.transforms.RandomAutocontrast(p=0.5)), + torchlm.bind(torchvision.transforms.RandomAdjustSharpness(sharpness_factor=3, p=0.5)), + # bind albumentations image only transforms + torchlm.bind(albumentations.ColorJitter(p=0.5)), + torchlm.bind(albumentations.GlassBlur(p=0.5)), + torchlm.bind(albumentations.RandomShadow(p=0.5)), + # bind albumentations dual transforms + torchlm.bind(albumentations.RandomCrop(height=200, width=200, p=0.5)), + torchlm.bind(albumentations.RandomScale(p=0.5)), + torchlm.bind(albumentations.Rotate(p=0.5)), + torchlm.LandmarksResize((256, 256)), + torchlm.LandmarksNormalize(), + torchlm.LandmarksToTensor(), + torchlm.LandmarksToNumpy(), + torchlm.LandmarksUnNormalize() + ]) +``` +* **bind** custom callable array or Tensor functions, using **torchlm.bind** + +```python +# First, defined your custom functions +def callable_array_noop( + img: np.ndarray, + landmarks: np.ndarray +) -> Tuple[np.ndarray, np.ndarray]: + # Do some transform here ... + return img.astype(np.uint32), landmarks.astype(np.float32) + + +def callable_tensor_noop( + img: Tensor, + landmarks: Tensor +) -> Tuple[Tensor, Tensor]: + # Do some transform here ... + return img, landmarks +``` + +```python +# Then, bind your functions and put it into transforms pipeline. +transform = torchlm.LandmarksCompose([ + # use native torchlm transforms + torchlm.LandmarksRandomScale(prob=0.5), + # ... + # bind torchvision image only transforms + torchlm.bind(torchvision.transforms.GaussianBlur(kernel_size=(5, 25))), + torchlm.bind(torchvision.transforms.RandomAutocontrast(p=0.5)), + torchlm.bind(torchvision.transforms.RandomAdjustSharpness(sharpness_factor=3, p=0.5)), + # bind albumentations image only transforms + torchlm.bind(albumentations.ColorJitter(p=0.5)), + torchlm.bind(albumentations.GlassBlur(p=0.5)), + torchlm.bind(albumentations.RandomShadow(p=0.5)), + # bind albumentations dual transforms + torchlm.bind(albumentations.RandomCrop(height=200, width=200, p=0.5)), + torchlm.bind(albumentations.RandomScale(p=0.5)), + torchlm.bind(albumentations.Rotate(p=0.5)), + # bind custom callable array functions + torchlm.bind(callable_array_noop, bind_type=torchlm.BindEnum.Callable_Array), + # bind custom callable Tensor functions + torchlm.bind(callable_tensor_noop, bind_type=torchlm.BindEnum.Callable_Tensor), + torchlm.LandmarksResize((256, 256)), + torchlm.LandmarksNormalize(), + torchlm.LandmarksToTensor(), + torchlm.LandmarksToNumpy(), + torchlm.LandmarksUnNormalize() + ]) +``` +
+ + + + + + + +
+ + +* setup logging mode as `True` globally might help you figure out the runtime details +```python +import torchlm +# some global setting +torchlm.set_transforms_debug(True) +torchlm.set_transforms_logging(True) +torchlm.set_autodtype_logging(True) +``` +Some details logs will show you at each runtime, just like the follows +```shell +LandmarksRandomHorizontalFlip() AutoDtype Info: AutoDtypeEnum.Array_InOut +LandmarksRandomHorizontalFlip() Execution Flag: True +LandmarksRandomScale() AutoDtype Info: AutoDtypeEnum.Array_InOut +LandmarksRandomScale() Execution Flag: False +... +BindTorchVisionTransform(GaussianBlur())() AutoDtype Info: AutoDtypeEnum.Tensor_InOut +BindTorchVisionTransform(GaussianBlur())() Execution Flag: True +... +BindAlbumentationsTransform(ColorJitter())() AutoDtype Info: AutoDtypeEnum.Array_InOut +BindAlbumentationsTransform(ColorJitter())() Execution Flag: True +... +BindArrayCallable(callable_array_noop())() AutoDtype Info: AutoDtypeEnum.Array_InOut +BindArrayCallable(callable_array_noop())() Execution Flag: True +BindTensorCallable(callable_tensor_noop())() AutoDtype Info: AutoDtypeEnum.Tensor_InOut +BindTensorCallable(callable_tensor_noop())() Execution Flag: True +... +LandmarksUnNormalize() AutoDtype Info: AutoDtypeEnum.Array_InOut +LandmarksUnNormalize() Execution Flag: True +``` +* Execution Flag: True means current transform was executed successful, False means it was not executed because of the random probability or some Runtime Exceptions(torchlm will should the error infos if debug mode is True). +* AutoDtype Info: + * Array_InOut means current transform need a np.ndnarray as input and then output a np.ndarray. + * Tensor_InOut means current transform need a torch Tensor as input and then output torch Tensor. + * Array_In means current transform needs a np.ndarray input and then output a torch Tensor. + * Tensor_In means current transform needs a torch Tensor input and then output a np.ndarray. + + But, is ok if your pass a Tensor to a np.ndarray like transform, **torchlm** will automatically be compatible with different data types and then wrap back to the original type through a autodtype wrapper. + + +* Supported Transforms Sets, see [transforms.md](docs/api/transfroms.md). A detail example can be found at [test/transforms.py](test/transforms.py). + +### Training(TODO) +* [ ] YOLOX +* [ ] YOLOv5 +* [ ] NanoDet +* [ ] PIPNet +* [ ] ResNet +* [ ] MobileNet +* [ ] ShuffleNet +* [ ] ... + +### Inference(TODO) +* [ ] ONNXRuntime +* [ ] MNN +* [ ] NCNN +* [ ] TNN +* [ ] ... + +## 📖 Documentations +* [ ] Native Data Augmentation's API (TODO) +* [ ] ... + +## 🎓 License +The code of torchlm is released under the MIT License. + +## 👋 Contributing +If you like this project please consider ⭐ this repo, as it is the simplest way to support me. + +## 🎓 Acknowledgement +The implementation of torchlm's transforms borrow the code from [Paperspace](ttps://github.com/Paperspace/DataAugmentationForObjectDetection/blob/master/data_aug/bbox_util.py). diff --git a/docs/api/.gitignore b/docs/api/.gitignore index a9a9176..c08061c 100644 --- a/docs/api/.gitignore +++ b/docs/api/.gitignore @@ -1,4 +1,3 @@ .idea __pycache__ .DS_Store -logs diff --git a/docs/api/transfroms.md b/docs/api/transfroms.md new file mode 100644 index 0000000..af411d8 --- /dev/null +++ b/docs/api/transfroms.md @@ -0,0 +1,139 @@ +## Supported Transforms Set + +### native torchlm's transforms + +```python +__all__ = [ + "LandmarksCompose", + "LandmarksNormalize", + "LandmarksUnNormalize", + "LandmarksToTensor", + "LandmarksToNumpy", + "LandmarksResize", + "LandmarksClip", + "LandmarksAlign", + "LandmarksRandomAlign", + "LandmarksRandomCenterCrop", + "LandmarksRandomHorizontalFlip", + "LandmarksHorizontalFlip", + "LandmarksRandomScale", + "LandmarksRandomTranslate", + "LandmarksRandomRotate", + "LandmarksRandomShear", + "LandmarksRandomHSV", + "LandmarksRandomMask", + "LandmarksRandomBlur", + "LandmarksRandomBrightness", + "LandmarksRandomPatches", + "LandmarksRandomBackground", + "LandmarksRandomPatchesWithAlpha", + "LandmarksRandomBackgroundWithAlpha", + "LandmarksRandomMaskWithAlpha", + "BindAlbumentationsTransform", + "BindTorchVisionTransform", + "BindArrayCallable", + "BindTensorCallable", + "BindEnum", + "bind", + "set_transforms_logging", + "set_transforms_debug" +] +``` + +### transforms from torchvision + +```python +# torchvision >= 0.9.0 +_Supported_Image_Only_Transform_Set: Tuple = ( + torchvision.transforms.Normalize, + torchvision.transforms.ColorJitter, + torchvision.transforms.Grayscale, + torchvision.transforms.RandomGrayscale, + torchvision.transforms.RandomErasing, + torchvision.transforms.GaussianBlur, + torchvision.transforms.RandomInvert, + torchvision.transforms.RandomPosterize, + torchvision.transforms.RandomSolarize, + torchvision.transforms.RandomAdjustSharpness, + torchvision.transforms.RandomAutocontrast, + torchvision.transforms.RandomEqualize +) +``` + +### transforms from albumentations + +```python +# albumentations >= v 1.1.0 +_Supported_Image_Only_Transform_Set: Tuple = ( + albumentations.Blur, + albumentations.CLAHE, + albumentations.ChannelDropout, + albumentations.ChannelShuffle, + albumentations.ColorJitter, + albumentations.Downscale, + albumentations.Emboss, + albumentations.Equalize, + albumentations.FDA, + albumentations.FancyPCA, + albumentations.FromFloat, + albumentations.GaussNoise, + albumentations.GaussianBlur, + albumentations.GlassBlur, + albumentations.HistogramMatching, + albumentations.HueSaturationValue, + albumentations.ISONoise, + albumentations.ImageCompression, + albumentations.InvertImg, + albumentations.MedianBlur, + albumentations.MotionBlur, + albumentations.Normalize, + albumentations.PixelDistributionAdaptation, + albumentations.Posterize, + albumentations.RGBShift, + albumentations.RandomBrightnessContrast, + albumentations.RandomFog, + albumentations.RandomGamma, + albumentations.RandomRain, + albumentations.RandomShadow, + albumentations.RandomSnow, + albumentations.RandomSunFlare, + albumentations.RandomToneCurve, + albumentations.Sharpen, + albumentations.Solarize, + albumentations.Superpixels, + albumentations.TemplateTransform, + albumentations.ToFloat, + albumentations.ToGray +) + +_Supported_Dual_Transform_Set: Tuple = ( + albumentations.Affine, + albumentations.CenterCrop, + albumentations.CoarseDropout, + albumentations.Crop, + albumentations.CropAndPad, + albumentations.CropNonEmptyMaskIfExists, + albumentations.Flip, + albumentations.HorizontalFlip, + albumentations.Lambda, + albumentations.LongestMaxSize, + albumentations.NoOp, + albumentations.PadIfNeeded, + albumentations.Perspective, + albumentations.PiecewiseAffine, + albumentations.RandomCrop, + albumentations.RandomCropNearBBox, + albumentations.RandomGridShuffle, + albumentations.RandomResizedCrop, + albumentations.RandomRotate90, + albumentations.RandomScale, + albumentations.RandomSizedCrop, + albumentations.Resize, + albumentations.Rotate, + albumentations.SafeRotate, + albumentations.ShiftScaleRotate, + albumentations.SmallestMaxSize, + albumentations.Transpose, + albumentations.VerticalFlip +) +``` \ No newline at end of file diff --git a/docs/res/.gitignore b/docs/res/.gitignore index a9a9176..dd02cf0 100644 --- a/docs/res/.gitignore +++ b/docs/res/.gitignore @@ -1,4 +1,4 @@ .idea __pycache__ .DS_Store -logs + diff --git a/docs/res/124.jpg b/docs/res/124.jpg new file mode 100644 index 0000000..afe3075 Binary files /dev/null and b/docs/res/124.jpg differ diff --git a/docs/res/158.jpg b/docs/res/158.jpg new file mode 100644 index 0000000..27dee15 Binary files /dev/null and b/docs/res/158.jpg differ diff --git a/docs/res/386.jpg b/docs/res/386.jpg new file mode 100644 index 0000000..1cdd35b Binary files /dev/null and b/docs/res/386.jpg differ diff --git a/docs/res/478.jpg b/docs/res/478.jpg new file mode 100644 index 0000000..a69e4d1 Binary files /dev/null and b/docs/res/478.jpg differ diff --git a/docs/res/537.jpg b/docs/res/537.jpg new file mode 100644 index 0000000..5f3ef42 Binary files /dev/null and b/docs/res/537.jpg differ diff --git a/docs/res/605.jpg b/docs/res/605.jpg new file mode 100644 index 0000000..67d8561 Binary files /dev/null and b/docs/res/605.jpg differ diff --git a/docs/res/802.jpg b/docs/res/802.jpg new file mode 100644 index 0000000..2f04589 Binary files /dev/null and b/docs/res/802.jpg differ diff --git a/docs/res/836.jpg b/docs/res/836.jpg new file mode 100644 index 0000000..366ec1e Binary files /dev/null and b/docs/res/836.jpg differ diff --git a/docs/res/logo.png b/docs/res/logo.png index 14496b4..b0faaf9 100644 Binary files a/docs/res/logo.png and b/docs/res/logo.png differ diff --git a/requirements.txt b/requirements.txt index e69de29..a504238 100644 --- a/requirements.txt +++ b/requirements.txt @@ -0,0 +1,9 @@ +# torchlm +opencv-python-headless>=4.5.2 +numpy>=1.14.4 +torch>=1.6.0 +torchvision>=0.9.0 +albumentations>=1.1.0 + +# test +matplotlib>=3.4.1 \ No newline at end of file diff --git a/setup.py b/setup.py index 0af0fbc..d36bddd 100644 --- a/setup.py +++ b/setup.py @@ -6,21 +6,27 @@ setuptools.setup( name="torchlm", version="0.1.0", - author="Yan Jun", + author="DefTruth", author_email="qyjdef@163.com", - description="A PyTorch library for landmarks detection, include data augmentation, training and inference.", + description="A PyTorch landmarks-only library with 100+ data augmentations, training and inference.", long_description=long_description, long_description_content_type="text/markdown", - url="https://github.com/DefTruth/torchaug", - packages=["torchlm"], + url="https://github.com/DefTruth/torchlm", + packages=setuptools.find_packages(), install_requires=[ - "opencv-python>=4.2.1", + "opencv-python-headless>=4.5.2", "numpy>=1.14.4", - "torch>=1.6.0" + "torch>=1.6.0", + "torchvision>=0.9.0", + "albumentations>=1.1.0" ], - classifiers=( + classifiers=[ "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", "License :: OSI Approved :: MIT License", - "Operating System :: OS Independent", - ), + "Operating System :: OS Independent" + ], ) diff --git a/test/logs/.gitignore b/test/logs/.gitignore new file mode 100644 index 0000000..4bc4044 --- /dev/null +++ b/test/logs/.gitignore @@ -0,0 +1 @@ +*.jpg \ No newline at end of file diff --git a/test/transforms.py b/test/transforms.py new file mode 100644 index 0000000..14e4e19 --- /dev/null +++ b/test/transforms.py @@ -0,0 +1,93 @@ +import cv2 +import numpy as np +import torchvision +import albumentations +from torch import Tensor +from typing import Tuple +import matplotlib.pyplot as plt + +import torchlm + + +def callable_array_noop( + img: np.ndarray, + landmarks: np.ndarray +) -> Tuple[np.ndarray, np.ndarray]: + # Do some transform here ... + return img.astype(np.uint32), landmarks.astype(np.float32) + + +def callable_tensor_noop( + img: Tensor, + landmarks: Tensor +) -> Tuple[Tensor, Tensor]: + # Do some transform here ... + return img, landmarks + + +def test_torchlm_transforms(): + + print(f"torchlm version: {torchlm.__version__}") + seed = np.random.randint(0, 1000) + np.random.seed(seed) + + img_path = "./1.jpg" + anno_path = "./1.txt" + save_path = f"./logs/{seed}.jpg" + img = cv2.imread(img_path)[:, :, ::-1].copy() # RGB + with open(anno_path, 'r') as fr: + lm_info = fr.readlines()[0].strip('\n').split(' ')[1:] + + landmarks = [float(x) for x in lm_info[4:]] + landmarks = np.array(landmarks).reshape(5, 2) # (5,2) + + # some global setting + torchlm.set_transforms_debug(True) + torchlm.set_transforms_logging(True) + torchlm.set_autodtype_logging(True) + + transform = torchlm.LandmarksCompose([ + # use native torchlm transforms + torchlm.LandmarksRandomHorizontalFlip(prob=0.5), + torchlm.LandmarksRandomScale(prob=0.5), + torchlm.LandmarksRandomTranslate(prob=0.5), + torchlm.LandmarksRandomShear(prob=0.5), + torchlm.LandmarksRandomMask(prob=0.5), + torchlm.LandmarksRandomBlur(kernel_range=(5, 25), prob=0.5), + torchlm.LandmarksRandomBrightness(prob=0.), + torchlm.LandmarksRandomRotate(40, prob=0.5, bins=8), + torchlm.LandmarksRandomCenterCrop((0.5, 1.0), (0.5, 1.0), prob=0.5), + # bind torchvision image only transforms + torchlm.bind(torchvision.transforms.GaussianBlur(kernel_size=(5, 25))), + torchlm.bind(torchvision.transforms.RandomAutocontrast(p=0.5)), + torchlm.bind(torchvision.transforms.RandomAdjustSharpness(sharpness_factor=3, p=0.5)), + # bind albumentations image only transforms + torchlm.bind(albumentations.ColorJitter(p=0.5)), + torchlm.bind(albumentations.GlassBlur(p=0.5)), + torchlm.bind(albumentations.RandomShadow(p=0.5)), + # bind albumentations dual transforms + torchlm.bind(albumentations.RandomCrop(height=200, width=200, p=0.5)), + torchlm.bind(albumentations.RandomScale(p=0.5)), + torchlm.bind(albumentations.Rotate(p=0.5)), + # bind custom callable array functions + torchlm.bind(callable_array_noop, bind_type=torchlm.BindEnum.Callable_Array), + # bind custom callable Tensor functions + torchlm.bind(callable_tensor_noop, bind_type=torchlm.BindEnum.Callable_Tensor), + torchlm.LandmarksResize((256, 256)), + torchlm.LandmarksNormalize(), + torchlm.LandmarksToTensor(), + torchlm.LandmarksToNumpy(), + torchlm.LandmarksUnNormalize() + ]) + + trans_img, trans_landmarks = transform(img, landmarks) + new_img = torchlm.draw_landmarks(trans_img, trans_landmarks) + plt.imshow(new_img) + plt.show() + cv2.imwrite(save_path, new_img[:, :, ::-1]) + print("transformed landmarks: ", trans_landmarks.flatten().tolist()) + print("original landmarks: ", landmarks.flatten().tolist()) + + +if __name__ == "__main__": + test_torchlm_transforms() diff --git a/torchlm/__init__.py b/torchlm/__init__.py index e69de29..d5bed03 100644 --- a/torchlm/__init__.py +++ b/torchlm/__init__.py @@ -0,0 +1,8 @@ +# Versions +__version__ = '0.1.0' +# Transforms Module: 100+ transforms available, can bind torchvision and +# albumentations into torchlm pipeline with autodtype wrapper. +from .transfroms import * +# Utils Module: some utils methods. +from .utils import * +# Other Modules: TODO \ No newline at end of file diff --git a/test/test.py b/torchlm/data/converter.py similarity index 100% rename from test/test.py rename to torchlm/data/converter.py diff --git a/torchlm/smooth/__init__.py b/torchlm/smooth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/torchlm/trackers/__init__.py b/torchlm/trackers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/torchlm/transfroms/__init__.py b/torchlm/transfroms/__init__.py index e69de29..2468b22 100644 --- a/torchlm/transfroms/__init__.py +++ b/torchlm/transfroms/__init__.py @@ -0,0 +1,2 @@ +from .transforms import * +from .autodtypes import set_autodtype_logging \ No newline at end of file diff --git a/torchlm/transfroms/autodtypes.py b/torchlm/transfroms/autodtypes.py index 33b0ca7..e1ac4c5 100644 --- a/torchlm/transfroms/autodtypes.py +++ b/torchlm/transfroms/autodtypes.py @@ -1,6 +1,6 @@ import numpy as np from torch import Tensor -from typing import Tuple, Union, Callable +from typing import Tuple, Union, Callable, Any, Dict from . import functional as F # base element_type @@ -16,6 +16,23 @@ class AutoDtypeEnum: Tensor_In: int = 2 Tensor_InOut: int = 3 +AutoDtypeLoggingMode: bool = False + +def set_autodtype_logging(logging: bool = False): + global AutoDtypeLoggingMode + AutoDtypeLoggingMode = logging + +def _autodtype_api_logging(self: Any, mode: int): + global AutoDtypeLoggingMode + if AutoDtypeLoggingMode: + mode_info_map: Dict[int, str] = { + AutoDtypeEnum.Array_In: "AutoDtypeEnum.Array_In", + AutoDtypeEnum.Array_InOut: "AutoDtypeEnum.Array_InOut", + AutoDtypeEnum.Tensor_In: "AutoDtypeEnum.Tensor_In", + AutoDtypeEnum.Tensor_InOut: "AutoDtypeEnum.Tensor_InOut" + } + print(f"{self}() AutoDtype Info: {mode_info_map[mode]}") + def autodtype(mode: int) -> Callable: # A Pythonic style to auto convert input dtype and let the output dtype unchanged @@ -39,6 +56,8 @@ def apply( ), "Error dtype, must be np.ndarray or Tensor!" # force array before transform and then wrap back. if mode == AutoDtypeEnum.Array_InOut: + _autodtype_api_logging(self, mode) + if any(( isinstance(img, Tensor), isinstance(landmarks, Tensor) @@ -53,8 +72,17 @@ def apply( ) img = F.to_tensor(img) landmarks = F.to_tensor(landmarks) + else: + img, landmarks = callable_array_or_tensor_func( + self, + img, + landmarks, + **kwargs + ) # force array before transform and don't wrap back. elif mode == AutoDtypeEnum.Array_In: + _autodtype_api_logging(self, mode) + if any(( isinstance(img, Tensor), isinstance(landmarks, Tensor) @@ -67,8 +95,17 @@ def apply( landmarks, **kwargs ) + else: + img, landmarks = callable_array_or_tensor_func( + self, + img, + landmarks, + **kwargs + ) # force tensor before transform and then wrap back. elif mode == AutoDtypeEnum.Tensor_InOut: + _autodtype_api_logging(self, mode) + if any(( isinstance(img, np.ndarray), isinstance(landmarks, np.ndarray) @@ -83,8 +120,17 @@ def apply( ) img = F.to_numpy(img) landmarks = F.to_numpy(landmarks) + else: + img, landmarks = callable_array_or_tensor_func( + self, + img, + landmarks, + **kwargs + ) # force tensor before transform and don't wrap back. elif mode == AutoDtypeEnum.Tensor_In: + _autodtype_api_logging(self, mode) + if any(( isinstance(img, np.ndarray), isinstance(landmarks, np.ndarray) @@ -97,7 +143,15 @@ def apply( landmarks, **kwargs ) + else: + img, landmarks = callable_array_or_tensor_func( + self, + img, + landmarks, + **kwargs + ) else: + _autodtype_api_logging(self, mode) img, landmarks = callable_array_or_tensor_func( self, img, diff --git a/torchlm/transfroms/functional.py b/torchlm/transfroms/functional.py index 11991a7..f4734cd 100644 --- a/torchlm/transfroms/functional.py +++ b/torchlm/transfroms/functional.py @@ -31,61 +31,6 @@ def to_numpy(x: Union[np.ndarray, Tensor]) -> np.ndarray: return x.detach().cpu().numpy() -def draw_bbox(im: np.ndarray, bboxes: np.ndarray) -> np.ndarray: - im = im[:, :, :].copy() - - bboxes = bboxes[:, :4] - bboxes = bboxes.reshape(-1, 4) - - # draw bbox - for box in bboxes: - im = cv2.rectangle(im, (int(box[0]), int(box[1])), - (int(box[2]), int(box[3])), (0, 255, 0), 2) - - return im.astype(np.uint8) - - -def draw_landmarks(im: np.ndarray, landmarks: np.ndarray, - font: float = 0.25, circle: int = 2, text: bool = False, - color: Optional[Tuple[int, int, int]] = None, - visibilities: Optional[np.ndarray] = None, - refine_circle: bool = False, - text_offset: int = 2, - thickness: int = 1) -> np.ndarray: - im = im.copy() - h, w, c = im.shape - refine = int(h * (1 / 512) + 1) - - if circle < refine and refine_circle: - circle = refine - - default_color = (0, 255, 0) - if color is None: - color = default_color - - for i in range(landmarks.shape[0]): - landmark_ = landmarks[i] - color_ = color - if visibilities is not None: - if visibilities[i] < 0.5: - color_ = (0, 0, 255) # red - - cv2.circle(im, (int(landmark_[0]), int(landmark_[1])), circle, color_, -1) - # if i == 33 or i == 48: - if text: - b = np.random.randint(0, 255) - g = np.random.randint(0, 255) - r = np.random.randint(0, 255) - # b = 0 - # g = 0 - # r = 255 - cv2.putText(im, '{}'.format(i), - (int(landmark_[0]), int(landmark_[1]) - text_offset), - cv2.FONT_ITALIC, font, (b, g, r), thickness) - - return im.astype(np.uint8) - - def bbox_area(bbox: np.ndarray) -> np.ndarray: return (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1]) diff --git a/torchlm/transfroms/transforms.py b/torchlm/transfroms/transforms.py index 8510349..de8d39a 100644 --- a/torchlm/transfroms/transforms.py +++ b/torchlm/transfroms/transforms.py @@ -7,7 +7,7 @@ import albumentations from torch import Tensor from abc import ABCMeta, abstractmethod -from typing import Tuple, Union, List, Optional, Callable +from typing import Tuple, Union, List, Optional, Callable, Any from . import functional as F from .autodtypes import ( @@ -48,9 +48,24 @@ "BindArrayCallable", "BindTensorCallable", "BindEnum", - "bind" + "bind", + "set_transforms_logging", + "set_transforms_debug" ] +TransformLoggingMode: bool = False +TransformDebugMode: bool = False + + +def set_transforms_logging(logging: bool = False): + global TransformLoggingMode + TransformLoggingMode = logging + + +def set_transforms_debug(debug: bool = False): + global TransformDebugMode + TransformDebugMode = debug + class LandmarksTransform(object): __metaclass__ = ABCMeta @@ -81,7 +96,7 @@ def __call__( raise NotImplementedError def __repr__(self) -> str: - return self.__class__.__name__ + '()' + return self.__class__.__name__ def apply_affine_to( self, @@ -150,20 +165,34 @@ def __init__( @autodtype(AutoDtypeEnum.Tensor_InOut) def __call__( self, - img: Image_InOutput_Type, - landmarks: Landmarks_InOutput_Type - ) -> Tuple[Image_InOutput_Type, Landmarks_InOutput_Type]: + img: Tensor, + landmarks: Tensor + ) -> Tuple[Tensor, Tensor]: # Image only transform from torchvision, - # just let the landmarks unchanged. + # just let the landmarks unchanged. Note (3,H,W) + # is need for torchvision try: + chw = img.size()[0] == 3 + if not chw: + img = img.permute((2, 0, 1)).contiguous() + img, landmarks = self.transform_internal(img), landmarks + # permute back + if not chw: + img = img.permute((1, 2, 0)).contiguous() + self.flag = True return img, landmarks except: self.flag = False return img, landmarks + def __repr__(self) -> str: + return self.__class__.__name__ + '(' \ + + self.transform_internal.__class__.__name__ \ + + '())' + Albumentations_Transform_Type = Union[ albumentations.ImageOnlyTransform, @@ -262,8 +291,8 @@ def __init__( )), f"The transform from albumentations must be one of:" \ f"\n{self._Supported_Image_Only_Transform_Set}, " \ f"\n{self._Supported_Dual_Transform_Set}" - - self.transform_internal = albumentations.Compose( + self.transform_internal = transform + self.compose_internal = albumentations.Compose( transforms=[transform], keypoint_params=albumentations.KeypointParams( format="xy", @@ -274,9 +303,9 @@ def __init__( @autodtype(AutoDtypeEnum.Array_InOut) def __call__( self, - img: Image_InOutput_Type, - landmarks: Landmarks_InOutput_Type - ) -> Tuple[Image_InOutput_Type, Landmarks_InOutput_Type]: + img: np.ndarray, + landmarks: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: # The landmarks format for albumentations should be a list of lists(tuple) # in xy format by default. Such as: # keypoints = [ @@ -299,7 +328,7 @@ def __call__( kps_num = len(keypoints) try: - transformed = self.transform_internal(image=img, keypoints=keypoints) + transformed = self.compose_internal(image=img, keypoints=keypoints) trans_img = transformed['image'] trans_kps = transformed['keypoints'] @@ -316,15 +345,20 @@ def __call__( self.flag = False return img.astype(np.uint8), landmarks.astype(np.float32) + def __repr__(self) -> str: + return self.__class__.__name__ + '(' \ + + self.transform_internal.__class__.__name__ \ + + '())' + Callable_Array_Func_Type = Union[ Callable[[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]], - Callable[[np.ndarray, np.ndarray, ...], Tuple[np.ndarray, np.ndarray]] + Callable[[np.ndarray, np.ndarray, Any], Tuple[np.ndarray, np.ndarray]] ] Callable_Tensor_Func_Type = Union[ Callable[[Tensor, Tensor], Tuple[Tensor, Tensor]], - Callable[[Tensor, Tensor, ...], Tuple[Tensor, Tensor]] + Callable[[Tensor, Tensor, Any], Tuple[Tensor, Tensor]] ] @@ -345,10 +379,10 @@ def __init__( @autodtype(AutoDtypeEnum.Array_InOut) def __call__( self, - img: Image_InOutput_Type, - landmarks: Landmarks_InOutput_Type, + img: np.ndarray, + landmarks: np.ndarray, **kwargs - ) -> Tuple[Image_InOutput_Type, Landmarks_InOutput_Type]: + ) -> Tuple[np.ndarray, np.ndarray]: try: img, landmarks = self.call_func(img, landmarks, **kwargs) @@ -360,6 +394,11 @@ def __call__( self.flag = False return img.astype(np.int32), landmarks.astype(np.float32) + def __repr__(self) -> str: + return self.__class__.__name__ + '(' \ + + self.call_func.__name__ \ + + '())' + class BindTensorCallable(LandmarksTransform): @@ -378,10 +417,10 @@ def __init__( @autodtype(AutoDtypeEnum.Tensor_InOut) def __call__( self, - img: Image_InOutput_Type, - landmarks: Landmarks_InOutput_Type, + img: Tensor, + landmarks: Tensor, **kwargs - ) -> Tuple[Image_InOutput_Type, Landmarks_InOutput_Type]: + ) -> Tuple[Tensor, Tensor]: try: img, landmarks = self.call_func(img, landmarks, **kwargs) @@ -391,6 +430,10 @@ def __call__( self.flag = False return img, landmarks + def __repr__(self) -> str: + return self.__class__.__name__ + '(' \ + + self.call_func.__name__ \ + + '())' Bind_Transform_Or_Callable_Input_Type = Union[ TorchVision_Transform_Type, @@ -450,11 +493,13 @@ class LandmarksCompose(object): def __init__( self, transforms: List[LandmarksTransform], - logging: bool = False + logging: bool = False, + debug: bool = False ): self.flags: List[bool] = [] self.transforms: List[LandmarksTransform] = transforms self.logging: bool = logging + self.debug: bool = debug assert self.check, "Wrong! Need LandmarksTransform !" \ f"But got {self.__repr__()}" @@ -473,10 +518,16 @@ def __call__( try: img, landmarks = t(img, landmarks) except Exception as e: - if self.logging: - print(f"Error at {t.__class__.__name__} Skip, " - f"Flag: {t.flag} Info: {e}") - continue + if self.logging or TransformLoggingMode: + print(f"Error at {t}() Skip, " + f"Flag: {t.flag} Error Info: {e}") + if self.debug or TransformDebugMode: + raise e + else: + continue + finally: + if self.logging or TransformLoggingMode: + print(f"{t}() Execution Flag: {t.flag}") self.flags.append(t.flag) return img, landmarks @@ -491,10 +542,15 @@ def apply_transform_to( if const_flag: other_img, other_landmarks = t(other_img, other_landmarks) except Exception as e: - if self.logging: - print(f"Error at {t.__class__.__name__} Skip, " - f"Flag: {t.flag} Info: {e}") + if self.logging or TransformLoggingMode: + print(f"Error at {t}() Skip, " + f"Flag: {t.flag} Error Info: {e}") + if self.debug or TransformDebugMode: + raise e continue + finally: + if self.logging or TransformLoggingMode: + print(f"{t}() Execution Flag: {t.flag}") return other_img, other_landmarks def apply_affine_to( @@ -516,10 +572,16 @@ def apply_affine_to( **kwargs ) except Exception as e: - if self.logging: - print(f"Error at {t.__class__.__name__} Skip, " - f"Flag: {t.flag} Info: {e}") - continue + if self.logging or TransformLoggingMode: + print(f"Error at {t} Skip, " + f"Flag: {t.flag} Error Info: {e}") + if self.debug or TransformDebugMode: + raise e + else: + continue + finally: + if self.logging or TransformLoggingMode: + print(f"{t}() Execution Flag: {t.flag}") return other_landmarks def clear_affine(self): @@ -1495,7 +1557,7 @@ class LandmarksRandomMask(LandmarksTransform): def __init__( self, - mask_ratio: float = 0.25, + mask_ratio: float = 0.1, prob: float = 0.5, trans_ratio: float = 0.5 ): @@ -1505,7 +1567,7 @@ def __init__( :param trans_ratio: control the random shape of masked area. """ super(LandmarksRandomMask, self).__init__() - assert 0.10 < mask_ratio < 1. + assert 0.02 < mask_ratio < 1. assert 0 < trans_ratio < 1. self._mask_ratio = mask_ratio self._trans_ratio = trans_ratio diff --git a/torchlm/utils/utils.py b/torchlm/utils/utils.py index e69de29..c3bedc1 100644 --- a/torchlm/utils/utils.py +++ b/torchlm/utils/utils.py @@ -0,0 +1,47 @@ +import cv2 +import numpy as np +from typing import Tuple, Optional + +__all__ = [ + "draw_bbox", + "draw_landmarks" +] + + +def draw_bbox(img: np.ndarray, bboxes: np.ndarray) -> np.ndarray: + im = img[:, :, :].copy() + + bboxes = bboxes[:, :4] + bboxes = bboxes.reshape(-1, 4) + + # draw bbox + for box in bboxes: + im = cv2.rectangle(im, (int(box[0]), int(box[1])), + (int(box[2]), int(box[3])), (0, 255, 0), 2) + + return im.astype(np.uint8) + + +def draw_landmarks( + img: np.ndarray, + landmarks: np.ndarray, + font: float = 0.25, + circle: int = 2, + text: bool = False, + color: Optional[Tuple[int, int, int]] = (0, 255, 0), + offset: int = 5, + thickness: int = 1 +) -> np.ndarray: + im = img.astype(np.uint8).copy() + + for i in range(landmarks.shape[0]): + x, y = landmarks[i].astype(int).tolist() + cv2.circle(im, (x, y), circle, color, -1) + if text: + b = np.random.randint(0, 255) + g = np.random.randint(0, 255) + r = np.random.randint(0, 255) + cv2.putText(im, '{}'.format(i), (x, y - offset), + cv2.FONT_ITALIC, font, (b, g, r), thickness) + + return im.astype(np.uint8)