A package for easy generation of binary semantic mask of different labels using multiple models easily. Moreover, supports operations on the mask created for image editing.
Added support for SAM.
- Deeplabv3 with pascal labels
- Maskrcnn with coco labels
- Bisnet with face labels
pip:
pip install seg-mask-modifs
pip install opencv-contrib-python>=4.5.4.60
# if you install opencv-python then inpainting won't work
Cloning repo then install requirements:
pip install -r requirements.txt
The documentation of the different classes and functions is available here
The models can be downloaded seperately or all of then can be downloaded at once.
from seg_mask_modifs import download_models
download_models.download_all() # download all models with default names which is highly recommended.
download_models.maskrcnn_coco(save_path='models/maskrcnn_restnet50_fpn.pt') # download maskrcnn model with coco labels
download_models.deeplab_pascal() # download deeplab model
download_models.face() # download bisnet face model
To see the list of labels supported by the package, this function can be used.
from seg_mask_modifs import print_labels
print_labels.all() # prints all labels
print_labels.deeplab_pascal() # prints pascal labels
print_labels.maskrcnn_coco() # prints coco labels
print_label.face() # prints face labels
Class to generate binary mask for any combination of labels. The models will be automatically used according to model preference and labels provided.
import cv2
from seg_mask_modifs import mask_generator
mask_gen = mask_generator.mask_generator(threshold=0.5, auto_init=True) # auto_init will only work if the models are saved to the default path.
# if auto_init is false or different path used to save model initialize them manually.
mask_gen.init_maskrcnn('maskrcnn.pt')
mask_gen.init_deeplab('deeplab.pt')
mask_gen.init_face('face.pth')
img = cv2.imread('images/city.jpg')
mask = mask_gen.generate(img=img, labels=['person', 'suitcase'])
In the example above the test image passed is:
This generates the following output, using deeplabv3 for person and maskrcnn for suitcase:
To generate mask use only one model, the use_model argument can be used.
mask = mask_gen.generate(img=img, labels=['person', 'suitcase'], use_model='maskrcnn')
By default, the deeplab model has the highest priority followed by maskrcnn and then bisnet face model. Any label which is supported by more than model will be generated by the one with higher preference. To check the model preference:
mask_gen.print_model_preference()
To change the model preference:
mask_gen.set_model_preference(model_list['maskrcnn', 'face', 'deeplab'])
mask_gen.set_model_preference(model='face', pos=0)
The model preference can be set using a list, or giving a particular model a particular position. Note: if model is not found in the list, it will be ignored.
Functions to combine, invert, dilate, etc. on multiple masks at once. Take a look at the documentation for more details.
Functions to modify masks. Includes various operations like blurring, pixelaing, replacing, inpainting the background and foreground among other operations. Take a look at the documentation for complete details. Some examples are given below.
