mantasu / face-crop-plus

Face aligner and cropper with quality enhancement and attribute parsing
https://mantasu.github.io/face-crop-plus/
MIT License
53 stars 6 forks source link
alignment computer-vision face face-attributes face-detection face-parsing pytorch segmentation super-resolution

Face Crop Plus

PyPI Python CUDA: yes Docs: passing DOI License: MIT

![Banner](https://raw.githubusercontent.com/mantasu/face-crop-plus/main/assets/banner.png)

About

Image preprocessing package for automatic face alignment and cropping with additional features. It provides the following functionality:

  1. Face cropping - face alignment and center-cropping using facial landmarks. Landmarks can be automatically predicted or, if they are already know, can be supplied through a separate file. It is possible to specify face factor, i.e., face area relative to the cropped image, and face extraction strategy, e.g., all faces or largest face per image.
  2. Face enhancement - face image quality enhancement. If images are blurry or contain many small faces, quality enhancement model can be used to make the images clearer. Small faces in the image are automatically checked and enhanced if desired.
  3. Face parsing - face attribute parsing and cropped image grouping to sub-directories. Face images can be grouped according to some facial attributes or some combination, such as glasses, earrings and necklace, hats. It is also possible to generate masks for facial attributes or some combination of them, for instance, glasses, nose, nose and eyes.

Please see References section for more details about which models are used for each feature.

Note: each feature can be used separately, e.g., if you just need to enhance the quality of blurry photos, or if you just need to generate attribute masks (like hats, glasses, face parts).

Installation

The packages requires at least Python 3.10. You may also want to set up PyTorch in advance from here.

To install the package simply run:

pip install face-crop-plus

Or, to install it from source, run:

git clone https://github.com/mantasu/face-crop-plus
cd face-crop-plus && pip install .

Quick Start

You can run the package from the command line:

face-crop-plus -i path/to/images

You can also use it in a Python script:

from face_crop_plus import Cropper

cropper = Cropper(face_factor=0.7, strategy="largest")
cropper.process_dir(input_dir="path/to/images")

For a quick demo, you can experiment with demo.py file:

git clone https://github.com/mantasu/face-crop-plus
cd face-crop-plus/demo
python demo.py

For more examples, see Examples section.

Features

Here, some of the main arguments are described that control the behavior of each of the features. These arguments can be specified via command line or when initializing the Cropper class. For further details about how the Cropper class works, please refer to the documentation.

Alignment and Cropping

The main feature is face alignment and cropping. The main arguments that control this feature:

Quality Enhancement

Quality enhancement feature allows to restore blurry faces. It has one main argument:

Quality enhancement can be used as a separate feature to enhance images that contain faces. For an end user, it is a useful feature to boost the quality of photos. It is not suggested to enhance ultra high resolution images (>2000) because your GPU will explode. See Pure Enhancement/Parsing section on how to run it as a stand-alone.

Attribute Parsing

Face parsing to attributes allows to group output images by category and generate attribute masks for that category. Categorized images are put to their corresponding sub-folders in the output directory.

If both attr_groups and mask_groups are specified, first images are grouped according to face attributes, then images in each groups are further sub-grouped to different mask groups (along with their masks).

Here are the 19 possible face attributes (with 0 representing the neutral category):

| | | | | ------------------- | ---------------- | ---------------- | | `1` - skin | `7` - left ear | `13` - lower lip | | `2` - left eyebrow | `8` - right ear | `14` - neck | | `3` - right eyebrow | `9` - earrings | `15` - necklace | | `4` - left eye | `10` - nose | `16` - clothes | | `5` - right eye | `11` - mouth | `17` - hair | | `6` - eyeglasses | `12` - upper lip | `18` - hat |

Examples

Running via Command Line

You can run the package via command line by providing the arguments as follows:

face-crop-plus -i path/to/images --output-size 200 300 --face-factor 0.75 -d cuda:0

You can specify the command-line arguments via JSON config file and provide the path to it. Further command-line arguments would overwrite the values taken from the JSON file.

face-crop-plus --config path/to/json --attr-groups '{"glasses": [6]}'

An example JSON config file is demo.json. If you've cloned the repository, you can run from it:

face-crop-plus --config demo/demo.json --device cuda:0 # overwrite device

For all the available command line arguments, just type (although refer to documentation for more details):

face-crop-plus -h

Note: you can use fcp as face-crop-plus alias , e.g., fcp -c config.json

Cleaning File Names

If your image files contain non-ascii symbols, lengthy names, os-reserved characters, it may be better to standardize them. To do so, it is possible to rename the image files before processing them:

face-crop-plus -i path/to/images --clean-names # --clean-names-inplace (avoids temp dir)

It is possible to specify more arguments via python script. The function can be used in general with any file types:

from face_crop_plus.utils import clean_names

clean_names(
    input_dir="path/to/input/dir",
    output_dir=None, # will rename in-place
    max_chars=250,
)

Pure Enhancement/Parsing

If you already have aligned and center-cropped face images, you can perform quality enhancement and face parsing without re-cropping them. Here is an example of enhancing quality of every face and parsing them to (note that none of the parameters described in Alignment and Cropping section have any affect here):

from face_crop_plus import Cropper

cropper = Cropper(
    det_threshold=None,
    enh_threshold=1, # enhance every image   
    attr_groups={"hats": [18], "no_hats": [-18]},
    mask_groups={"hats": [18], "ears": [7, 8, 9]},
    device="cuda:0",
)

cropper.crop(input_dir="path/to/images")

This would result in the following output directory structure:

└── path/to/images_faces
     ├── hats
     |    ├── hats       # Images with hats
     |    ├── hats_mask  # Hat masks for images in upper dir
     |    ├── ears       # Images with hats and visible ears
     |    └── ears_mask  # Ears masks for images in upper dir
     |
     └── no_hats
          ├── ears       # Masks with no hats and visible ears
          └── ears_mask  # Ears masks for images in upper dir

To just enhance the quality of images (e.g., if you have blurry photos), you can run enhancement feature separately:

face-crop-plus -i path/to/images -dt -1 -et 1 --device cuda:0

To just generate masks for images (e.g., as part of your research pipeline), you can run segmentation feature separately. This will only consider images for which the masks are actually present.

face-crop-plus -i path/to/images -dt -1 -et -1 -mg '{"glasses": [6]}'

Please beware of the following:

Preprocessing CelebA

Here is an example pipeline of how to pre-process CelebA dataset. It is useful if you want to customize the cropped face properties, e.g., face factor, output size. It only takes a few minutes to pre-process the whole dataset using multiple processors and the provided landmarks:

  1. Download the following files from Google Drive:
    • Download img_celeba.7z folder from here and put it under data/img_celeba.7z
    • Download annotations.zip file from here and put it under data/annotations.zip
  2. Unzip the data:
    7z x data/img_celeba.7z/img_celeba.7z.001 -o./data
    unzip data/annotations.zip -d data
  3. Create a script file, e.g., preprocess_celeba.py, in the same directory:

    from face_crop_plus import Cropper
    from multiprocessing import cpu_count
    
    cropper = Cropper(
        output_size=256,
        face_factor=0.7,
        landmarks="data/landmark.txt",
        enh_threshold=None,
        num_processes=cpu_count(),
    )
    
    cropper.process_dir("data/img_celeba")
  4. Run the script to pre-process the data:
    python preprocess_celeba.py
  5. Clean up the data dir (remove the original images and the annotations):
    rm -r data/img_celeba.7z data/img_celeba
    rm data/annotations.zip data/*.txt

Tips

  1. When using num_processes, only set it to a larger value if you have enough GPU memory, or reduce batch_size. Unless you only perform face cropping with already known landmarks and don't perform quality enhancement nor face parsing, in which case set it to the number of CPU cores you have.
  2. If you experience any of the following:

    • RuntimeError: CUDA error: an illegal memory access was encountered.
    • torch.cuda.OutOfMemoryError: CUDA out of memory.
    • cuDNN error: CUDNN_STATUS_MAPPING_ERROR.

    This is likely because you are processing images on too many processes or have a large batch size. If you run all 3 models on GPU, it may be helpful to just run on a single process with a larger batch size.

References

This package uses the code and the pre-trained models from the following repositories:

Citation

If you find this package helpful in your research, you can cite the following:

@misc{face-crop-plus,
  author = {Mantas Birškus},
  title = {Face Crop Plus},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/mantasu/face-crop-plus}},
  doi = {10.5281/zenodo.7856749}
}