├── LICENSE ├── README.md ├── main.py ├── .gitignore ├── requirements.txt ├── app.py └── app_discriminative.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Javier Esteve 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Detectron2 Instance Segmentation 2 | 3 | Streamlit App that performs object detection and instance segmentation powered by Detectron2. 4 | 5 | ![app_preview](https://user-images.githubusercontent.com/17023476/98106894-086ba000-1e9a-11eb-8479-a3679e36d64a.png) 6 | 7 | 8 | ## Dependencies 9 | 10 | Running the application can be done following the instructions above: 11 | 12 | 1. To create a Python Virtual Environment (virtualenv) to run the code, type: 13 | 14 | ```python3 -m venv my-env``` 15 | 16 | 2. Activate the new environment: 17 | * Windows: ```my-env\Scripts\activate.bat``` 18 | * macOS and Linux: ```source my-env/bin/activate``` 19 | 20 | 3. Install all the dependencies from *requirements.txt*: 21 | 22 | ```pip install -r requirements.txt``` 23 | 24 | ## Use 25 | 26 | Within the virtual environment: 27 | 28 | ```streamlit run app.py``` 29 | 30 | A web application will open in the prompted URL. The user should upload an image file (*jpg*, *jpeg*, *png*) with the button available. Then, the image will be fed to a model which will output tehe original image with the detections drawn on it. 31 | 32 | There's another app: 33 | 34 | ```streamlit run app_discriminative.py``` 35 | 36 | It behaves as the other one, but includes the following options: 37 | 38 | * **Select which classes to detect:** Multiselect to choose which of the classes that the model's been trained on are going to be used in the inference. 39 | 40 | ## Acknowledgments 41 | 42 | * [Detectron 2](https://github.com/facebookresearch/detectron2) 43 | 44 | ## License 45 | 46 | This project is licensed under the MIT License - see the [LICENSE.md](LICENSE.md) file for details -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | from detectron2 import model_zoo 5 | from detectron2.engine import DefaultPredictor 6 | from detectron2.config import get_cfg 7 | from detectron2.utils.visualizer import Visualizer 8 | from detectron2.data import MetadataCatalog, DatasetCatalog 9 | 10 | CLASSES = [0, 17] 11 | 12 | 13 | def load_cfg(): 14 | cfg = get_cfg() 15 | # Force model to operate within CPU, erase if CUDA compatible devices ara available 16 | cfg.MODEL.DEVICE = 'cpu' 17 | # add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library 18 | cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) 19 | cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model 20 | # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well 21 | cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") 22 | 23 | return cfg 24 | 25 | 26 | def inference(predictor, img): 27 | return predictor(img) 28 | 29 | 30 | def visualize_output(cfg, img, outputs): 31 | # We can use `Visualizer` to draw the predictions on the image. 32 | v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2) 33 | out = v.draw_instance_predictions(outputs["instances"].to("cpu")) 34 | cv2.imshow('kk', out.get_image()[:, :, ::-1]) 35 | cv2.waitKey(0) 36 | 37 | 38 | def discriminate(outputs): 39 | pred_classes = np.array(outputs['instances'].pred_classes) 40 | mask = np.isin(pred_classes, CLASSES) 41 | idx = np.nonzero(mask) 42 | 43 | # Get Instance values as a dict and leave only the desired ones 44 | out_fields = outputs['instances'].get_fields() 45 | for field in out_fields: 46 | out_fields[field] = out_fields[field][idx] 47 | 48 | return outputs 49 | 50 | 51 | def main(): 52 | #img = cv2.imread('img.png') 53 | img = cv2.imread('dog.jpg') 54 | # cv2.imshow('kk', img) 55 | # cv2.waitKey(0) 56 | cfg = load_cfg() 57 | predictor = DefaultPredictor(cfg) 58 | outputs = inference(predictor, img) 59 | # aaa = outputs.copy() 60 | # bbb = discriminate(outputs) 61 | visualize_output(cfg, img, outputs) 62 | 63 | 64 | if __name__ == "__main__": 65 | main() -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Visual Studio Code 2 | .vscode/ 3 | 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.so 11 | 12 | # Distribution / packaging 13 | .Python 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | wheels/ 26 | pip-wheel-metadata/ 27 | share/python-wheels/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | MANIFEST 32 | 33 | # PyInstaller 34 | # Usually these files are written by a python script from a template 35 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 36 | *.manifest 37 | *.spec 38 | 39 | # Installer logs 40 | pip-log.txt 41 | pip-delete-this-directory.txt 42 | 43 | # Unit test / coverage reports 44 | htmlcov/ 45 | .tox/ 46 | .nox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | *.py,cover 54 | .hypothesis/ 55 | .pytest_cache/ 56 | 57 | # Translations 58 | *.mo 59 | *.pot 60 | 61 | # Django stuff: 62 | *.log 63 | local_settings.py 64 | db.sqlite3 65 | db.sqlite3-journal 66 | 67 | # Flask stuff: 68 | instance/ 69 | .webassets-cache 70 | 71 | # Scrapy stuff: 72 | .scrapy 73 | 74 | # Sphinx documentation 75 | docs/_build/ 76 | 77 | # PyBuilder 78 | target/ 79 | 80 | # Jupyter Notebook 81 | .ipynb_checkpoints 82 | 83 | # IPython 84 | profile_default/ 85 | ipython_config.py 86 | 87 | # pyenv 88 | .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 98 | __pypackages__/ 99 | 100 | # Celery stuff 101 | celerybeat-schedule 102 | celerybeat.pid 103 | 104 | # SageMath parsed files 105 | *.sage.py 106 | 107 | # Environments 108 | .env 109 | .venv 110 | env/ 111 | venv/ 112 | ENV/ 113 | env.bak/ 114 | venv.bak/ 115 | 116 | # Spyder project settings 117 | .spyderproject 118 | .spyproject 119 | 120 | # Rope project settings 121 | .ropeproject 122 | 123 | # mkdocs documentation 124 | /site 125 | 126 | # mypy 127 | .mypy_cache/ 128 | .dmypy.json 129 | dmypy.json 130 | 131 | # Pyre type checker 132 | .pyre/ -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==0.11.0 2 | altair==4.1.0 3 | argon2-cffi==20.1.0 4 | astor==0.8.1 5 | async-generator==1.10 6 | attrs==20.2.0 7 | backcall==0.2.0 8 | base58==2.0.1 9 | bleach==3.2.1 10 | blinker==1.4 11 | boto3==1.16.10 12 | botocore==1.19.10 13 | cachetools==4.1.1 14 | certifi==2020.6.20 15 | cffi==1.14.3 16 | chardet==3.0.4 17 | click==7.1.2 18 | cloudpickle==1.6.0 19 | cycler==0.10.0 20 | Cython==0.29.21 21 | decorator==4.4.2 22 | defusedxml==0.6.0 23 | detectron2==0.2.1+cpu 24 | entrypoints==0.3 25 | enum-compat==0.0.3 26 | future==0.18.2 27 | fvcore==0.1.2.post20201030 28 | gitdb==4.0.5 29 | GitPython==3.1.11 30 | google-auth==1.23.0 31 | google-auth-oauthlib==0.4.2 32 | grpcio==1.33.2 33 | idna==2.10 34 | importlib-metadata==2.0.0 35 | ipykernel==5.3.4 36 | ipython==7.16.1 37 | ipython-genutils==0.2.0 38 | ipywidgets==7.5.1 39 | jedi==0.17.2 40 | Jinja2==2.11.2 41 | jmespath==0.10.0 42 | jsonschema==3.2.0 43 | jupyter-client==6.1.7 44 | jupyter-core==4.6.3 45 | jupyterlab-pygments==0.1.2 46 | kiwisolver==1.3.1 47 | Markdown==3.3.3 48 | MarkupSafe==1.1.1 49 | matplotlib==3.3.2 50 | mistune==0.8.4 51 | mock==4.0.2 52 | nbclient==0.5.1 53 | nbconvert==6.0.7 54 | nbformat==5.0.8 55 | nest-asyncio==1.4.2 56 | notebook==6.1.4 57 | numpy==1.19.4 58 | oauthlib==3.1.0 59 | opencv-python==4.4.0.46 60 | packaging==20.4 61 | pandas==1.1.4 62 | pandocfilters==1.4.3 63 | parso==0.7.1 64 | pathtools==0.1.2 65 | pexpect==4.8.0 66 | pickleshare==0.7.5 67 | Pillow==8.0.1 68 | pkg-resources==0.0.0 69 | portalocker==2.0.0 70 | prometheus-client==0.8.0 71 | prompt-toolkit==3.0.8 72 | protobuf==3.13.0 73 | ptyprocess==0.6.0 74 | pyarrow==2.0.0 75 | pyasn1==0.4.8 76 | pyasn1-modules==0.2.8 77 | pycocotools==2.0.2 78 | pycparser==2.20 79 | pydeck==0.5.0 80 | pydot==1.4.1 81 | Pygments==2.7.2 82 | pyparsing==2.4.7 83 | pyrsistent==0.17.3 84 | python-dateutil==2.8.1 85 | pytz==2020.4 86 | PyYAML==5.3.1 87 | pyzmq==19.0.2 88 | requests==2.24.0 89 | requests-oauthlib==1.3.0 90 | rsa==4.6 91 | s3transfer==0.3.3 92 | Send2Trash==1.5.0 93 | six==1.15.0 94 | smmap==3.0.4 95 | streamlit==0.70.0 96 | tabulate==0.8.7 97 | tensorboard==2.3.0 98 | tensorboard-plugin-wit==1.7.0 99 | termcolor==1.1.0 100 | terminado==0.9.1 101 | testpath==0.4.4 102 | toml==0.10.2 103 | toolz==0.11.1 104 | torch==1.6.0 105 | torchvision==0.7.0 106 | tornado==6.1 107 | tqdm==4.51.0 108 | traitlets==4.3.3 109 | tzlocal==2.1 110 | urllib3==1.25.11 111 | validators==0.18.1 112 | watchdog==0.10.3 113 | wcwidth==0.2.5 114 | webencodings==0.5.1 115 | Werkzeug==1.0.1 116 | widgetsnbextension==3.5.1 117 | yacs==0.1.8 118 | zipp==3.4.0 119 | -------------------------------------------------------------------------------- /app.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import streamlit as st 4 | 5 | from detectron2 import model_zoo 6 | from detectron2.engine import DefaultPredictor 7 | from detectron2.config import get_cfg 8 | from detectron2.utils.visualizer import Visualizer 9 | from detectron2.data import MetadataCatalog, DatasetCatalog 10 | 11 | 12 | @st.cache(persist=True) 13 | def initialization(): 14 | """Loads configuration and model for the prediction. 15 | 16 | Returns: 17 | cfg (detectron2.config.config.CfgNode): Configuration for the model. 18 | predictor (detectron2.engine.defaults.DefaultPredicto): Model to use. 19 | by the model. 20 | 21 | """ 22 | cfg = get_cfg() 23 | # Force model to operate within CPU, erase if CUDA compatible devices ara available 24 | cfg.MODEL.DEVICE = 'cpu' 25 | # Add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library 26 | cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) 27 | # Set threshold for this model 28 | cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 29 | # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well 30 | cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") 31 | # Initialize prediction model 32 | predictor = DefaultPredictor(cfg) 33 | 34 | return cfg, predictor 35 | 36 | 37 | @st.cache 38 | def inference(predictor, img): 39 | return predictor(img) 40 | 41 | 42 | @st.cache 43 | def output_image(cfg, img, outputs): 44 | v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2) 45 | out = v.draw_instance_predictions(outputs["instances"].to("cpu")) 46 | processed_img = out.get_image() 47 | 48 | return processed_img 49 | 50 | 51 | @st.cache 52 | def discriminate(outputs, classes_to_detect): 53 | """Select which classes to detect from an output. 54 | 55 | Get the dictionary associated with the outputs instances and modify 56 | it according to the given classes to restrict the detection to them 57 | 58 | Args: 59 | outputs (dict): 60 | instances (detectron2.structures.instances.Instances): Instance 61 | element which contains, among others, "pred_boxes", 62 | "pred_classes", "scores" and "pred_masks". 63 | classes_to_detect (list: int): Identifiers of the dataset on which 64 | the model was trained. 65 | 66 | Returns: 67 | ouputs (dict): Same dict as before, but modified to match 68 | the detection classes. 69 | 70 | """ 71 | pred_classes = np.array(outputs['instances'].pred_classes) 72 | # Take the elements matching *classes_to_detect* 73 | mask = np.isin(pred_classes, classes_to_detect) 74 | # Get the indexes 75 | idx = np.nonzero(mask) 76 | 77 | # Get the current Instance values 78 | pred_boxes = outputs['instances'].pred_boxes 79 | pred_classes = outputs['instances'].pred_classes 80 | pred_masks = outputs['instances'].pred_masks 81 | scores = outputs['instances'].scores 82 | 83 | # Get them as a dictionary and leave only the desired ones with the indexes 84 | out_fields = outputs['instances'].get_fields() 85 | out_fields['pred_boxes'] = pred_boxes[idx] 86 | out_fields['pred_classes'] = pred_classes[idx] 87 | out_fields['pred_masks'] = pred_masks[idx] 88 | out_fields['scores'] = scores[idx] 89 | 90 | return outputs 91 | 92 | 93 | def main(): 94 | # Initialization 95 | cfg, predictor = initialization() 96 | 97 | # Streamlit initialization 98 | st.title("Instance Segmentation") 99 | 100 | # Retrieve image 101 | uploaded_img = st.file_uploader("Choose an image...", type=['jpg', 'jpeg', 'png']) 102 | if uploaded_img is not None: 103 | file_bytes = np.asarray(bytearray(uploaded_img.read()), dtype=np.uint8) 104 | img = cv2.imdecode(file_bytes, 1) 105 | # Detection code 106 | outputs = inference(predictor, img) 107 | out_image = output_image(cfg, img, outputs) 108 | st.image(out_image, caption='Processed Image', use_column_width=True) 109 | 110 | 111 | if __name__ == '__main__': 112 | main() -------------------------------------------------------------------------------- /app_discriminative.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import streamlit as st 4 | 5 | from detectron2 import model_zoo 6 | from detectron2.engine import DefaultPredictor 7 | from detectron2.config import get_cfg 8 | from detectron2.utils.visualizer import Visualizer 9 | from detectron2.data import MetadataCatalog, DatasetCatalog 10 | 11 | 12 | @st.cache(persist=True) 13 | def initialization(): 14 | """Loads configuration and model for the prediction. 15 | 16 | Returns: 17 | cfg (detectron2.config.config.CfgNode): Configuration for the model. 18 | classes_names (list: str): Classes available for the model of interest. 19 | predictor (detectron2.engine.defaults.DefaultPredicto): Model to use. 20 | by the model. 21 | 22 | """ 23 | cfg = get_cfg() 24 | # Force model to operate within CPU, erase if CUDA compatible devices ara available 25 | cfg.MODEL.DEVICE = 'cpu' 26 | # Add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library 27 | cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) 28 | # Set threshold for this model 29 | cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 30 | # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well 31 | cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") 32 | # Get classes names for the dataset of interest 33 | classes_names = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes 34 | # Initialize prediction model 35 | predictor = DefaultPredictor(cfg) 36 | 37 | return cfg, classes_names, predictor 38 | 39 | 40 | def inference(predictor, img): 41 | return predictor(img) 42 | 43 | 44 | @st.cache 45 | def output_image(cfg, img, outputs): 46 | v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2) 47 | out = v.draw_instance_predictions(outputs["instances"].to("cpu")) 48 | processed_img = out.get_image() #[:, :, ::-1] 49 | 50 | return processed_img 51 | 52 | 53 | #@st.cache 54 | def retrieve_image(): 55 | uploaded_img = st.file_uploader("Choose an image...", type=['jpg', 'jpeg', 'png']) 56 | uploaded_img_cache = None 57 | if uploaded_img is not None and uploaded_img != uploaded_img_cache: 58 | uploaded_img_cache = uploaded_img 59 | print(uploaded_img) 60 | return image(uploaded_img) 61 | # file_bytes = np.asarray(bytearray(uploaded_img.read()), dtype=np.uint8) 62 | # try: 63 | # img = cv2.imdecode(file_bytes, 1) 64 | # except: 65 | # pass 66 | # return img 67 | 68 | 69 | @st.cache 70 | def image(uploaded_img): 71 | file_bytes = np.asarray(bytearray(uploaded_img.read()), dtype=np.uint8) 72 | try: 73 | img = cv2.imdecode(file_bytes, 1) 74 | return img 75 | except: 76 | return None 77 | return img 78 | 79 | 80 | 81 | @st.cache 82 | def discriminate(outputs, classes_to_detect): 83 | """Select which classes to detect from an output. 84 | 85 | Get the dictionary associated with the outputs instances and modify 86 | it according to the given classes to restrict the detection to them 87 | 88 | Args: 89 | outputs (dict): 90 | instances (detectron2.structures.instances.Instances): Instance 91 | element which contains, among others, "pred_boxes", 92 | "pred_classes", "scores" and "pred_masks". 93 | classes_to_detect (list: int): Identifiers of the dataset on which 94 | the model was trained. 95 | 96 | Returns: 97 | ouputs (dict): Same dict as before, but modified to match 98 | the detection classes. 99 | 100 | """ 101 | pred_classes = np.array(outputs['instances'].pred_classes) 102 | # Take the elements matching *classes_to_detect* 103 | mask = np.isin(pred_classes, classes_to_detect) 104 | # Get the indexes 105 | idx = np.nonzero(mask) 106 | 107 | # Get Instance values as a dict and leave only the desired ones 108 | out_fields = outputs['instances'].get_fields() 109 | for field in out_fields: 110 | out_fields[field] = out_fields[field][idx] 111 | 112 | return outputs 113 | 114 | 115 | def main(): 116 | # Initialization 117 | cfg, classes, predictor = initialization() 118 | 119 | # Streamlit initialization 120 | st.title("Instance Segmentation") 121 | st.sidebar.title("Options") 122 | ## Select classes to be detected by the model 123 | classes_to_detect = st.sidebar.multiselect( 124 | "Select which classes to detect", classes, ['person']) 125 | mask = np.isin(classes, classes_to_detect) 126 | class_idxs = np.nonzero(mask) 127 | 128 | # Define holder for the processed image 129 | img_placeholder = st.empty() 130 | 131 | # Retrieve image 132 | uploaded_img = st.file_uploader("Choose an image...", type=['jpg', 'jpeg', 'png']) 133 | if uploaded_img is not None: 134 | try: 135 | file_bytes = np.asarray(bytearray(uploaded_img.read()), dtype=np.uint8) 136 | img = cv2.imdecode(file_bytes, 1) 137 | # Detection code 138 | outputs = inference(predictor, img) 139 | outputs = discriminate(outputs, class_idxs) 140 | out_image = output_image(cfg, img, outputs) 141 | st.image(out_image, caption='Processed Image', use_column_width=True) 142 | except: 143 | st.subheader("Please, reupload an image to see the changes") 144 | 145 | 146 | if __name__ == '__main__': 147 | main() --------------------------------------------------------------------------------