├── .gitignore ├── LICENSE ├── README.md ├── background └── bg.jpg ├── classes └── products.txt ├── coco_dataset_generator ├── extras │ ├── combine_json_files.py │ ├── compute_dataset_statistics.py │ ├── create_binary_dataset.py │ ├── cut_objects.py │ ├── extract_frames.py │ ├── move_dataset_to_single_folder.py │ ├── occlusion_transforms.py │ ├── sample.py │ └── split_json_file.py ├── gui │ ├── contours.py │ ├── poly_editor.py │ ├── segment.py │ └── segment_bbox_only.py └── utils │ ├── create_json_file.py │ ├── delete_images.py │ ├── pascal_to_coco.py │ ├── visualize_dataset.py │ └── visualize_json_file.py ├── gui.png ├── images ├── davi_0_1.jpg └── davi_2_71.jpg ├── requirements.txt ├── requirements_maskrcnn.txt ├── setup.py └── unit_tests ├── test_button.py ├── test_coco.py ├── test_poly_pasting.py └── test_zoom.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.pyc 3 | node_modules/ 4 | *.egg-info 5 | *.h5 6 | dist/ 7 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | Copyright 2019 coco_dataset_generator 179 | 180 | Licensed under the Apache License, Version 2.0 (the "License"); 181 | you may not use this file except in compliance with the License. 182 | You may obtain a copy of the License at 183 | 184 | http://www.apache.org/licenses/LICENSE-2.0 185 | 186 | Unless required by applicable law or agreed to in writing, software 187 | distributed under the License is distributed on an "AS IS" BASIS, 188 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 189 | See the License for the specific language governing permissions and 190 | limitations under the License. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # COCO-Style-Dataset-Generator-GUI 2 | 3 | This is a simple GUI-based Widget based on matplotlib in Python to facilitate quick and efficient crowd-sourced generation of annotation masks and bounding boxes using a simple interactive User Interface. Annotation can be in terms of polygon points covering all parts of an object (see instructions in README) or it can simply be a bounding box, for which you click and drag the mouse button. Optionally, one could choose to use a pretrained Mask RCNN model to come up with initial segmentations. This shifts the work load from painstakingly annotating all the objects in every image to altering wrong predictions made by the system which maybe simpler once an efficient model is learnt. 4 | 5 | #### Note: This repo only contains code to annotate every object using a single polygon figure. Support for multi-polygon objects and `iscrowd=True` annotations isn't available yet. Feel free to extend the repo as you wish. Also, the code uses xyxy bounding boxes while coco uses xywh; something to keep in mind if you intend to create a custom COCO dataset to plug into other models as COCO datasets. 6 | 7 | ### REQUIREMENTS: 8 | 9 | `Python 3.5+` is required to run the Mask RCNN code. If only the GUI tool is used, `Python2.7` or `Python3.5+` can be used. 10 | 11 | ###### NOTE: For python2.7, OpenCV needs to be installed from source and configured to be in the environment running the code. 12 | ###### Before installing, please upgrade setuptools using: pip install --upgrade setuptools 13 | ###### For Windows users, please install Visual Studio C++ 14 or higher if necessary using this link: http://go.microsoft.com/fwlink/?LinkId=691126&fixForIE=.exe. 14 | 15 | ### RUN THE SEGMENTOR GUI: 16 | 17 | Clone the repo. 18 | 19 | ``` 20 | git clone https://github.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI.git 21 | ``` 22 | 23 | #### Installing Dependencies: 24 | 25 | Before running the code, install required pre-requisite python packages using pip. 26 | 27 | If you wish to use Mask RCNN to prelabel based on a trained model, please use the environment variable `MASK_RCNN="y"`, otherwise there's no need to include it and you could just perform the install. 28 | 29 | ###### Without Mask RCNN 30 | 31 | ``` 32 | cd COCO-Style-Dataset-Generator-GUI/ 33 | python setup.py install 34 | ``` 35 | 36 | ###### With Mask RCNN 37 | 38 | ``` 39 | cd COCO-Style-Dataset-Generator-GUI/ 40 | MASK_RCNN="y" python3 setup.py install 41 | ``` 42 | 43 | #### Running the instance segmentation GUI without Mask RCNN pretrained predictions: 44 | 45 | In a separate text file, list the target labels/classes line-by-line to be displayed along with the dataset for class labels. For example, look at [classes/products.txt](https://github.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI/blob/master/classes/products.txt) 46 | 47 | ``` 48 | python3 -m coco_dataset_generator.gui.segment -i background/ -c classes/products.txt 49 | 50 | python3 -m coco_dataset_generator.gui.segment_bbox_only -i background/ -c classes/products.txt 51 | ``` 52 | 53 | #### Running the instance segmentation GUI augmented by initial Mask RCNN pretrained model predictions: 54 | 55 | To run the particular model for the demo, download the pretrained weights from [HERE!!!](https://drive.google.com/file/d/1GaKVP3BvTfMwPbhEm4nF7fLATV-eDkFQ/view?usp=sharing). Download and extract pretrained weights into the repository. 56 | 57 | ``` 58 | python3 -m coco_dataset_generator.gui.segment -i background/ -c classes/products.txt \ 59 | -w [--config ] 60 | 61 | python3 -m coco_dataset_generator.gui.segment_bbox_only -i background/ -c classes/products.txt \ 62 | -w [--config ] 63 | ``` 64 | 65 | The configuration file for Mask RCNN becomes relevant when you play around with the configuration parameters that make up the network. In order to seamlessly use the repository with multiple such Mask RCNN models for different types of datasets, you could create a single config file for every project and use them as you please. The base repository has been configured to work well with the demo model provided and so any change to the parameters should be followed by generation of its corresponding config file. 66 | 67 | HINT: Use `get_json_config.py` inside `Mask RCNN` to get config file wrt specific parameters of Mask RCNN. You could either clone [Mask_RCNN](https://www.github.com/hanskrupakar/Mask_RCNN), use `pip install -e Mask_RCNN/` to replace the mask_rcnn installed from this repo and then get access to `get_json_config.py` easily or you could find where pip installs `mask_rcnn` and find it directly from the source. 68 | 69 | `USAGE: segment.py [-h] -i IMAGE_DIR -c CLASS_FILE [-w WEIGHTS_PATH] [-x CONFIG_PATH]` 70 | 71 | `USAGE: segment_bbox_only.py [-h] -i IMAGE_FILE -c CLASSES_FILE [-j JSON_FILE] [--save_csv] [-w WEIGHTS_PATH] [-x CONFIG_PATH]` 72 | 73 | ##### Optional Arguments 74 | 75 | | Shorthand | Flag Name | Description | 76 | | --------------- | --------------------------- | ---------------------------------------------------------------------------------- | 77 | | -h | --help | Show this help message and exit | 78 | | -i IMAGE_DIR | --image_dir IMAGE_DIR | Path to the image dir | 79 | | -c CLASS_FILE | --class_file CLASS_FILE | Path to object labels | 80 | | -w WEIGHTS_PATH | --weights_path WEIGHTS_PATH | Path to Mask RCNN checkpoint save file | 81 | | -j JSON_FILE | --json_file JSON_FILE | Path of JSON file to append dataset to | 82 | | | --save_csv | Choose option to save dataset as CSV file | 83 | | -x CONFIG_FILE | --config_file CONFIG_FILE | Path of JSON file for training config; Use `get_json_config` script from Mask RCNN | 84 | 85 | ### POLYGON SEGMENTATION GUI CONTROLS: 86 | 87 | ![deepmagic](https://github.com/Deep-Magic/COCO-Style-Dataset-Generator-GUI/blob/master/gui.png) 88 | 89 | In this demo, all the green patches over the objects are the rough masks generated by a pretrained Mask RCNN network. 90 | 91 | Key-bindings/ 92 | Buttons 93 | 94 | EDIT MODE (when `a` is pressed and polygon is being edited) 95 | 96 | 'a' toggle vertex markers on and off. 97 | When vertex markers are on, you can move them, delete them 98 | 99 | 'd' delete the vertex under point 100 | 101 | 'i' insert a vertex at point near the boundary of the polygon. 102 | 103 | Left click Use on any point on the polygon boundary and move around 104 | by dragging to alter shape of polygon 105 | 106 | REGULAR MODE 107 | 108 | Scroll Up Zoom into image 109 | 110 | Scroll Down Zoom out of image 111 | 112 | Left Click Create a point for a polygon mask around an object 113 | 114 | Right Click Complete the polygon currently formed by connecting all selected points 115 | 116 | Left Click Drag Create a bounding box rectangle from point 1 to point 2 (works only 117 | when there are no polygon points on screen for particular object) 118 | 119 | 'a' Press key on top of overlayed polygon (from Mask RCNN or 120 | previous annotations) to select it for editing 121 | 122 | 'r' Press key on top of overlayed polygon (from Mask RCNN or 123 | previous annotations) to completely remove it 124 | 125 | BRING PREVIOUS ANNOTATIONS Bring back the annotations from the previous image to preserve 126 | similar annotations. 127 | 128 | SUBMIT To be clicked after Right click completes polygon! Finalizes current 129 | segmentation mask and class label picked. 130 | After this, the polygon cannot be edited. 131 | 132 | NEXT Save all annotations created for current file and move on to next image. 133 | 134 | PREV Goto previous image to re-annotate it. This deletes the annotations 135 | created for the file before the current one in order to 136 | rewrite the fresh annotations. 137 | 138 | RESET If when drawing the polygon using points, the polygon doesn't cover the 139 | object properly, reset will let you start fresh with the current polygon. 140 | This deletes all the points on the image. 141 | 142 | The green annotation boxes from the network can be edited by pressing on the Keyboard key `a` when the mouse pointer is on top of a particular such mask. Once you press `a`, the points making up that polygon will show up and you can then edit it using the key bindings specified. Once you're done editing the polygon, press `a` again to finalize the edits. At this point, it will become possible to submit that particular annotation and move on to the next one. 143 | 144 | Once the GUI tool has been used successfully and relevant txt files have been created for all annotated images, one can use `create_json_file.py` to create the COCO-Style JSON file. 145 | 146 | ``` 147 | python -m coco_dataset_generator.utils.create_json_file -i background/ -c classes/products.txt 148 | -o output.json -t jpg 149 | ``` 150 | 151 | ``` 152 | USAGE: create_json_file.py [-h] -i IMAGE_DIR -o FILE_PATH -c CLASS_FILE -t TYPE 153 | ``` 154 | 155 | ##### Optional Arguments 156 | 157 | | Shorthand | Flag Name | Description | 158 | | ------------- | ----------------------- | --------------------------------------- | 159 | | -i IMAGE_DIR | --image_dir IMAGE_DIR | Path to the image dir | 160 | | -o FILE_PATH | --file_path FILE_PATH | Path of output file | 161 | | -c CLASS_FILE | --class_file CLASS_FILE | Path of file with output classes | 162 | | -t TYPE | --type TYPE | Type of the image files (jpg, png etc.) | 163 | 164 | ### RECTANGULAR BOUNDING BOX GUI CONTROLS: 165 | 166 | The same GUI is designed slightly differently in case of rectangular bounding box annotations with speed of annotation in mind. Thus, most keys are keyboard bindings. Most ideally, this interface is very suited to serve to track objects across video by dragging around a box of similar size. Since the save button saves multiple frame results together, the JSON file is directly created instead of txt files for each image, which means there wouldn't be a need to use `create_json_file.py`. 167 | 168 | Key-bindings/ 169 | Buttons 170 | 171 | EDIT MODE (when `a` is pressed and rectangle is being edited) 172 | 173 | 'a' toggle vertex markers on and off. When vertex markers are on, 174 | you can move them, delete them 175 | 176 | 'i' insert rectangle in the list of final objects to save. 177 | 178 | Left click Use on any point on the rectangle boundary and move around by 179 | dragging to alter shape of rectangle 180 | 181 | REGULAR MODE 182 | 183 | Scroll Up Zoom into image 184 | 185 | Scroll Down Zoom out of image 186 | 187 | Left Click Drag Create a bounding box rectangle from point 1 to point 2. 188 | 189 | 'a' Press key on top of overlayed polygon (from Mask RCNN or 190 | previous annotations) to select it for editing 191 | 192 | 'r' Press key on top of overlayed polygon (from Mask RCNN or 193 | previous annotations) to completely remove it 194 | 195 | 'n' Press key to move on to next image after completing all 196 | rectangles in current image 197 | 198 | SAVE Save all annotated objects so far 199 | 200 | ### LIST OF FUNCTIONALITIES: 201 | 202 | FILE FUNCTIONALITY 203 | 204 | cut_objects.py Cuts objects based on bounding box annotations using dataset.json 205 | file and creates occlusion-based augmented images dataset. 206 | 207 | create_json_file.py Takes a directory of annotated images (use segment.py to annotate 208 | into text files) and returns a COCO-style JSON file. 209 | 210 | extract_frames.py Takes a directory of videos and extracts all the frames of all 211 | videos into a folder labeled adequately by the video name. 212 | 213 | pascal_to_coco.py Takes a PASCAL-style dataset directory with JPEGImages/ and 214 | Annotations/ folders and uses the bounding box as masks to 215 | create a COCO-style JSON file. 216 | 217 | segment.py Read the instructions above. 218 | 219 | segment_bbox_only.py Same functionality but optimized for easier annotation of 220 | bbox-only datasets. 221 | 222 | test_*.py Unit tests. 223 | 224 | visualize_dataset.py Visualize the annotations created using the tool. 225 | 226 | visualize_json_file.py Visualize the dataset JSON file annotations on the entire dataset. 227 | 228 | compute_dataset_statistics.py Find distribution of objects in the dataset by counts. 229 | 230 | combine_json_files.py Combine different JSON files together into a single dataset file. 231 | 232 | delete_images.py Delete necessary images from the JSON dataset. 233 | 234 | NOTE: Please use `python .py -h` for details on how to use each of the above files. 235 | -------------------------------------------------------------------------------- /background/bg.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI/35c095a815a6ef97d2e820d175bf46e32e59f40b/background/bg.jpg -------------------------------------------------------------------------------- /classes/products.txt: -------------------------------------------------------------------------------- 1 | background 2 | blue_perfume 3 | black_perfume 4 | double_speedstick 5 | blue_speedstick 6 | dove_blue 7 | dove_perfume 8 | dove_pink 9 | green_speedstick 10 | gear_deo 11 | dove_black 12 | grey_speedstick 13 | choc_blue 14 | choc_red 15 | choc_yellow 16 | black_cup 17 | nyu_cup 18 | ilny_white 19 | ilny_blue 20 | ilny_black 21 | human 22 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/combine_json_files.py: -------------------------------------------------------------------------------- 1 | ''' 2 | USAGE: 3 | python combine_json_files.py 4 | ''' 5 | import os 6 | import json 7 | import glob 8 | import sys 9 | import numpy as np 10 | import argparse 11 | import os 12 | 13 | def cleanup_utf8(array): 14 | arr = [x.encode('ascii', errors='ignore').decode('utf-8') for x in array] 15 | return list(map(lambda x: x.strip().strip('\"').strip('\''), arr)) 16 | 17 | def merge_json(files, outfile='merged_dataset.json', abspath=False): 18 | 19 | img_counter = 0 20 | ann_counter = 0 21 | 22 | images, annotations, classes = [], [], [] 23 | 24 | for file_path in files: 25 | 26 | with open(file_path, 'r') as f: 27 | obj = json.load(f) 28 | 29 | for img in obj["images"]: 30 | img["id"] += img_counter 31 | 32 | if not abspath: 33 | img['file_name'] = os.path.join( 34 | os.path.abspath(os.path.dirname(file_path)), 35 | img['file_name']) 36 | 37 | for ann in obj["annotations"]: 38 | ann["id"] += ann_counter 39 | ann["image_id"] += img_counter 40 | 41 | ann_counter += len(obj["annotations"]) 42 | img_counter += len(obj["images"]) 43 | 44 | if len(images)==0: 45 | images = obj['images'] 46 | annotations = obj['annotations'] 47 | classes = cleanup_utf8(obj['classes']) 48 | else: 49 | obj['classes'] = cleanup_utf8(obj['classes']) 50 | 51 | if classes != obj["classes"]: 52 | print ("CLASSES MISMATCH BETWEEN:") 53 | print (classes) 54 | print (obj['classes']) 55 | if len(obj['classes']) < len(classes): 56 | c1, c2 = obj['classes'], classes 57 | new = True 58 | else: 59 | c1, c2 = classes, obj['classes'] 60 | new = False 61 | 62 | mapping = {} 63 | for idx, c in enumerate(c1): 64 | try: 65 | mapping[idx] = c2.index(c) 66 | except Exception: 67 | c2.append(c) 68 | mapping[idx] = len(c2) - 1 69 | 70 | print ('MAPPING: ', mapping) 71 | if not new: 72 | for idx, ann in enumerate(annotations): 73 | annotations[idx]['category_id'] = mapping[ann['category_id']-1] + 1 74 | classes = obj['classes'] 75 | else: 76 | for idx, ann in enumerate(obj['annotations']): 77 | obj['annotations'][idx]['category_id'] = mapping[ann['category_id']-1] + 1 78 | obj['classes'] = classes 79 | 80 | print ("CHANGE IN NUMBER OF CLASSES HAS BEEN DETECTED BETWEEN JSON FILES") 81 | print ("NOW MAPPING OLD CLASSES TO NEW LIST BASED ON TEXTUAL MATCHING") 82 | 83 | for k, v in mapping.items(): 84 | print (c1[k], "==>", c2[v]) 85 | 86 | remaining = set(c2) - set(c1) 87 | for r in remaining: 88 | print ("NEW CLASS: ", r) 89 | 90 | images.extend(obj["images"]) 91 | annotations.extend(obj["annotations"]) 92 | 93 | with open(outfile, "w") as f: 94 | data = {'images': images, 'annotations':annotations, 'classes': classes, 'categories':[]} 95 | json.dump(data, f) 96 | 97 | if __name__=='__main__': 98 | 99 | if len(sys.argv) < 3: 100 | print ("Not enough input files to combine into a single dataset file") 101 | exit() 102 | 103 | ap = argparse.ArgumentParser() 104 | ap.add_argument('files', nargs='+', help='List of JSON files to combine into single JSON dataset file') 105 | ap.add_argument('--absolute', nargs='+', help='Flag to use absolute paths in JSON file') 106 | args = ap.parse_args() 107 | 108 | merge_json(args.files, 'merged_json.json', args.absolute) 109 | 110 | 111 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/compute_dataset_statistics.py: -------------------------------------------------------------------------------- 1 | import json 2 | import argparse 3 | import numpy as np 4 | import cv2 5 | import os 6 | 7 | class DatasetStatistics: 8 | 9 | def __init__(self, json_file): 10 | 11 | with open(json_file, 'r') as f: 12 | self.json_obj = json.load(f) 13 | 14 | class_list = sorted(self.json_obj['classes']) 15 | 16 | self.class_counts = np.zeros(len(class_list), dtype=int) 17 | for ann in self.json_obj['annotations']: 18 | self.class_counts[ann['category_id']-1]+=1 19 | 20 | def print_occurrences(self): 21 | 22 | class_counts, class_list = zip(*sorted(zip(self.class_counts, sorted(self.json_obj['classes'])))) 23 | print ("\n".join([x+": "+str(y)+" occurrences" for x,y in zip(class_list, class_counts)])) 24 | 25 | def compute_mean_rgb(self): 26 | 27 | cnt, img_sum = 0, np.array([0,0,0], dtype='float') 28 | 29 | for img in self.json_obj['images']: 30 | if os.path.exists(img['file_name']): 31 | print ("Processing ", img['file_name']) 32 | 33 | img = cv2.cvtColor(cv2.imread(img['file_name']), cv2.COLOR_BGR2RGB) 34 | cnt+= np.prod(img.shape[:2]) 35 | img_sum += np.sum(np.sum(img, axis=0), axis=0) 36 | 37 | img_sum /= cnt 38 | 39 | print ("Mean RGB: ", img_sum) 40 | 41 | if __name__=='__main__': 42 | 43 | ap = argparse.ArgumentParser() 44 | ap.add_argument('-j', '--json_file', required=True, help='Path to dataset JSON file created') 45 | args = ap.parse_args() 46 | 47 | stat = DatasetStatistics(args.json_file) 48 | stat.print_occurrences() 49 | stat.compute_mean_rgb() 50 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/create_binary_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | 4 | if __name__=='__main__': 5 | 6 | ap = argparse.ArgumentParser() 7 | ap.add_argument('json', help='Path to original multi-class JSON file') 8 | args = ap.parse_args() 9 | 10 | with open(args.json, 'r') as f: 11 | obj = json.load(f) 12 | 13 | obj['classes'] = ['object'] 14 | 15 | for idx in range(len(obj['annotations'])): 16 | obj['annotations'][idx]['category_id'] = 1 17 | 18 | with open('.'.join(args.json.split('.')[:-1])+'_binary.json', 'w') as f: 19 | json.dump(obj, f) 20 | 21 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/cut_objects.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | np.set_printoptions(threshold=sys.maxsize) 4 | 5 | import matplotlib.pyplot as plt 6 | from PIL import Image, ImageDraw 7 | import json 8 | from collections import defaultdict 9 | import random 10 | from matplotlib.path import Path 11 | import argparse 12 | import glob 13 | import os 14 | import time 15 | #from skimage.measure import find_contours 16 | from ..gui.contours import find_contours 17 | 18 | class Occlusion_Generator_Bbox(object): 19 | 20 | def __init__(self, json_file, bg_dir, max_objs, imgs_path, curve_factor): 21 | 22 | self.dataset = json.load(open(json_file)) 23 | self.max_objs = max_objs 24 | 25 | self.imgToAnns = defaultdict(list) 26 | for ann in self.dataset['annotations']: 27 | self.imgToAnns[ann['image_id']].append(ann) 28 | 29 | self.objToAnns = [[] for _ in range(len(self.dataset['classes'])+1)] 30 | for index in self.imgToAnns: 31 | for obj in self.imgToAnns[index]: 32 | self.objToAnns[obj['category_id']].append({'image': obj['image_id'], 'bbox':obj['bbox']}) 33 | 34 | self.bg_dir = bg_dir 35 | 36 | self.imgs_dir = imgs_path 37 | self.set_random_background() 38 | 39 | self.classes = ['BG'] + self.dataset['classes'] 40 | self.curve_factor = curve_factor 41 | 42 | def set_random_background(self): 43 | 44 | imgs = [x for x in glob.glob(os.path.join(self.bg_dir, '*')) if 'txt' not in x] 45 | print (imgs, self.bg_dir) 46 | bg_path = random.choice(imgs) 47 | self.img = Image.open(bg_path).convert("RGBA") 48 | self.mask_img = Image.new('L', self.img.size, 0) 49 | self.text = '' 50 | 51 | def cut_bbox(self, rect): # Takes a bounding box of the form [x_min, y_min, x_max, y_max] and splits it in 2 based on a sine wave and returns 2 PIL polygons 52 | 53 | x = np.linspace(rect[0], rect[2], num=50) 54 | y = (rect[3]+rect[1])/2 + 15*np.sin(x/(rect[3]/np.pi/self.curve_factor)) 55 | 56 | x1 = np.concatenate((x, np.array([rect[2], rect[0]]))) 57 | y1 = np.concatenate((y, np.array([rect[3], rect[3]]))) 58 | 59 | x2 = np.concatenate((x, np.array([rect[2], rect[0]]))) 60 | y2 = np.concatenate((y, np.array([rect[1], rect[1]]))) 61 | 62 | poly1 = [(x,y) for x,y in zip(x1, y1)] 63 | poly2 = [(x,y) for x,y in zip(x2, y2)] 64 | 65 | return random.choice([poly1, poly2]) 66 | 67 | def add_objects(self): #Adds enlarged versions of n_objs (RANDOM) objects to self.img at random locations without overlap 68 | 69 | self.text += '%d'%self.image_id + '\n' + os.path.abspath(os.path.join(self.imgs_dir, '%d.jpg'%self.image_id))+'\n'+' '.join([str(x) for x in self.img.size])+'\n\n' 70 | 71 | n_objs = random.randint(5, self.max_objs) 72 | 73 | for _ in range(n_objs): 74 | 75 | c1 = random.randint(1, len(self.objToAnns)-1) 76 | c2 = random.randint(0, len(self.objToAnns[c1])-1) 77 | 78 | obj = Image.open(next(item for item in self.dataset['images'] if item["id"] == self.objToAnns[c1][c2]['image'])['file_name']) 79 | obj_bbox = self.objToAnns[c1][c2]['bbox'] 80 | obj_bbox = (obj_bbox[2], obj_bbox[3], obj_bbox[0], obj_bbox[1]) 81 | 82 | obj_mask = Image.new('L', obj.size, 0) 83 | random_occ = self.cut_bbox(obj_bbox) 84 | ImageDraw.Draw(obj_mask).polygon(random_occ, outline=255, fill=255) 85 | 86 | obj = obj.crop(obj_bbox) 87 | obj_mask = obj_mask.crop(obj_bbox) 88 | 89 | obj = obj.resize(np.array(np.array(obj.size)*1.35, dtype=int)) 90 | obj_mask = obj_mask.resize(np.array(np.array(obj_mask.size)*1.35, dtype=int)) 91 | 92 | done_flag, timeout = False, False 93 | clk = time.time() 94 | 95 | while not done_flag: 96 | 97 | if time.time()-clk > 1: # One second timeout 98 | timeout = True 99 | 100 | randx = random.randint(0, self.img.size[0]-obj.size[0]-2) 101 | randy = random.randint(0, self.img.size[1]-obj.size[1]-2) 102 | 103 | temp_mask = self.mask_img.copy() 104 | temp_mask.paste(Image.new('L', obj_mask.size, 0), (randx, randy)) 105 | 106 | if (temp_mask == self.mask_img): 107 | 108 | self.img.paste(obj, (randx, randy), obj_mask) 109 | self.mask_img.paste(obj_mask, (randx, randy)) 110 | 111 | obj_ann = Image.new('L', self.mask_img.size, 0) 112 | obj_ann.paste(obj_mask, (randx, randy)) 113 | 114 | padded_mask = np.zeros((obj_ann.size[0] + 2, obj_ann.size[1] + 2), dtype=np.uint8) 115 | padded_mask[1:-1, 1:-1] = np.array(obj_ann) 116 | contours = find_contours(padded_mask, 0.5) 117 | contours = [np.fliplr(verts) - 1 for verts in contours] 118 | 119 | x, y = contours[0][:,0], contours[0][:,1] 120 | area = (0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))))/2 #shoelace algorithm 121 | 122 | self.text += self.classes[c1]+'\n'+'%.2f'%area+'\n'+np.array2string(contours[0].flatten(), max_line_width=np.inf, formatter={'float_kind':lambda x: "%.2f" % x})[1:-1]+'\n\n' 123 | 124 | done_flag=True 125 | 126 | if not done_flag and timeout: # Add timeout-based object-preferencing 127 | 128 | print ('Object Timeout') 129 | timeout = False 130 | c2 = random.randint(0, len(self.objToAnns[c1])-1) 131 | 132 | obj = Image.open(next(item for item in self.dataset['images'] if item["id"] == self.objToAnns[c1][c2]['image'])['file_name']) 133 | obj_bbox = self.objToAnns[c1][c2]['bbox'] 134 | obj_bbox = (obj_bbox[2], obj_bbox[3], obj_bbox[0], obj_bbox[1]) 135 | 136 | obj_mask = Image.new('L', obj.size, 0) 137 | random_occ = self.cut_bbox(obj_bbox) 138 | ImageDraw.Draw(obj_mask).polygon(random_occ, outline=255, fill=255) 139 | 140 | obj = obj.crop(obj_bbox) 141 | obj_mask = obj_mask.crop(obj_bbox) 142 | 143 | obj = obj.resize(np.array(np.array(obj.size)*1.35, dtype=int)) 144 | obj_mask = obj_mask.resize(np.array(np.array(obj_mask.size)*1.35, dtype=int)) 145 | 146 | with open(os.path.join(self.imgs_dir, '%d.txt'%self.image_id), 'w') as f: 147 | f.write(self.text) 148 | self.img.convert('RGB').save(os.path.join(self.imgs_dir, '%d.jpg'%self.image_id)) 149 | 150 | def generate_images(self, num_imgs): 151 | 152 | self.image_id = 0 153 | 154 | for i in range(num_imgs): 155 | 156 | self.set_random_background() 157 | self.add_objects() 158 | self.image_id += 1 159 | print ('Image %d/%d created successfully!!!'%(i+1, num_imgs)) 160 | 161 | if __name__=='__main__': 162 | 163 | parser = argparse.ArgumentParser( 164 | description='Create occluded dataset.') 165 | parser.add_argument('--json_file', required=True, 166 | metavar="/path/to/json_file/", 167 | help='Path to JSON file', default='../pascal_dataset.json') 168 | parser.add_argument('--bg_dir', required=True, 169 | metavar="/path/to/possible/background/images", 170 | help="Path to Background Images", default='background/') 171 | parser.add_argument('--new_dir', required=True, 172 | help="Path to the new dataset directory", default='10') 173 | parser.add_argument('--max_objs', required=True, 174 | help="Maximum number of objects in an image (min=5)", default='10') 175 | parser.add_argument('--curve_factor', required=False, 176 | help="Amount of curvature of the sine wave (>2 values lead to high freq cuts)", default='1.4') 177 | parser.add_argument('--num_imgs', required=True, 178 | help="Total number of images in the synthetic dataset", default='50') 179 | args = parser.parse_args() 180 | 181 | occ = Occlusion_Generator_Bbox(args.json_file, args.bg_dir, int(args.max_objs), args.new_dir, float(args.curve_factor)) 182 | occ.generate_images(int(args.num_imgs)) 183 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/extract_frames.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import glob 4 | import argparse 5 | import imutils 6 | import os 7 | 8 | if __name__=='__main__': 9 | 10 | ap = argparse.ArgumentParser() 11 | ap.add_argument("-v", "--videos_dir", required=True, help="Path to the dir of videos") 12 | ap.add_argument("-o", "--frames_dir", required=True, help="Path to the output dir") 13 | ap.add_argument("-r", "--rotation", required=False, help="Angle of rotation") 14 | args = ap.parse_args() 15 | 16 | if not os.path.isdir(args.frames_dir): 17 | os.mkdir(args.frames_dir) 18 | 19 | for f in glob.glob(os.path.join(args.videos_dir, '*')): 20 | 21 | print (f) 22 | cap = cv2.VideoCapture(f) 23 | if not os.path.isdir(os.path.join(args.frames_dir, f.split('/')[-1][:-4])): 24 | os.mkdir(os.path.join(args.frames_dir, f.split('/')[-1][:-4])) 25 | 26 | i = 0 27 | ret, frame = cap.read() 28 | while(ret): 29 | 30 | i+=1 31 | 32 | if args.rotation is not None: 33 | frame = imutils.rotate_bound(frame, int(args.rotation)) 34 | 35 | if os.path.exists(os.path.join(*[args.frames_dir, f.split('/')[-1][:-4], f.split('/')[-1][:-4]+'_%d.jpg'%(i)])): 36 | continue 37 | else: 38 | cv2.imwrite(os.path.join(*[args.frames_dir, f.split('/')[-1][:-4], f.split('/')[-1][:-4]+'_%d.jpg'%(i)]), frame) 39 | 40 | ret, frame = cap.read() 41 | 42 | cap.release() 43 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/move_dataset_to_single_folder.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import shutil 3 | import json 4 | import os 5 | 6 | if __name__=='__main__': 7 | 8 | ap = argparse.ArgumentParser() 9 | ap.add_argument('dir', help='Path to folder to put all images in the dataset') 10 | ap.add_argument('json', help='Path to folder to put all images in the dataset') 11 | args = ap.parse_args() 12 | 13 | with open(args.json, 'r') as f: 14 | obj = json.load(f) 15 | 16 | try: 17 | os.makedirs(args.dir) 18 | except Exception: 19 | pass 20 | 21 | for idx, img in enumerate(obj['images']): 22 | 23 | path = img['file_name'] 24 | newpath = os.path.join(args.dir, '%s.'%(str(idx).zfill(5))+path.split('.')[-1]) 25 | 26 | shutil.copyfile(path, newpath) 27 | 28 | print ("Moving %s to %s"%(path, newpath)) 29 | 30 | obj['images'][idx]['file_name'] = newpath 31 | 32 | print ("Writing new JSON file!") 33 | 34 | base, direc = os.path.basename(args.dir), os.path.dirname(args.dir) 35 | 36 | with open(os.path.join(direc, '%s_dataset.json'%(base)), 'w') as f: 37 | json.dump(obj, f) 38 | print ("JSON file written!") 39 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/occlusion_transforms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import os 4 | import argparse 5 | import scipy.interpolate 6 | import time 7 | from shapely.geometry import Polygon 8 | 9 | #from skimage.measure import find_contours 10 | from ..gui.contours import find_contours 11 | 12 | from PIL import Image, ImageDraw 13 | 14 | class Occlusion_Generator(object): 15 | 16 | def __init__(self, strip_width): 17 | 18 | self.random_factor = 8 19 | self.distance = strip_width 20 | class Annotation(object): 21 | 22 | def __init__(self): 23 | 24 | self.objects = [] 25 | self.classes = [] 26 | 27 | self.all_images = [] 28 | 29 | self.images = [] 30 | 31 | self.polys = [] 32 | 33 | self.im_shape = np.asarray(Image.open(glob.glob(os.path.join(args["image_dir"], "*.jpg"))[0])).shape 34 | 35 | for ptr, f in enumerate(glob.glob(os.path.join(args["image_dir"], "*.jpg"))): 36 | 37 | print ("Processing Image %d/%d"%(ptr+1, len(glob.glob(os.path.join(args["image_dir"], "*.jpg"))))) 38 | 39 | im = Image.open(f).convert('RGBA') 40 | im.load() 41 | 42 | self.images.append(np.asarray(Image.open(f))) 43 | 44 | # convert to numpy (for convenience) 45 | imArray = np.asarray(im) 46 | 47 | lines = [x for x in range(50, imArray.shape[0], 100)] 48 | 49 | image_contents = Annotation() 50 | 51 | with open(f[:-3]+'txt', 'r') as f: 52 | txt = f.read().split('\n') 53 | 54 | for index in range(6, len(txt), 4): 55 | 56 | num = [float(x) for x in txt[index].split(' ')[:-1]] 57 | num = [(num[i], num[i+1]) for i in range(0, len(num), 2)] 58 | image_contents.objects.append([num]) 59 | image_contents.classes.append(txt[index-2]) 60 | 61 | strips = [Annotation() for _ in range(len(lines[2:]))] 62 | 63 | poly = [(imArray.shape[1], 0), (0, 0)] 64 | 65 | for pos, l in enumerate(lines[2:]): 66 | 67 | if ptr == 0: 68 | 69 | x, y = [0, imArray.shape[1]], [l, l+self.distance] 70 | y_interp = scipy.interpolate.interp1d(x, y) 71 | x_pts, y_pts = [x[0]], [y[0]] 72 | 73 | for p in range(0, imArray.shape[1], 5): 74 | yt = y_interp(p) + (2*np.random.random_sample()-1)*self.random_factor 75 | x_pts.append(p + (2*np.random.random_sample()-1)*self.random_factor) 76 | y_pts.append(yt) 77 | x_pts.append(x[1]) 78 | y_pts.append(y[1]) 79 | 80 | pts = [(x, y) for x, y in zip(x_pts, y_pts)] 81 | poly.extend(pts) 82 | 83 | self.polys.append(poly) 84 | 85 | else: 86 | 87 | poly = self.polys[pos] 88 | 89 | #ImageDraw.Draw(im).polygon(poly, fill="white", outline=None) 90 | 91 | #ImageDraw.Draw(im).line(pts, fill=128) 92 | 93 | #im.show() 94 | #time.sleep(.1) 95 | 96 | # create mask 97 | 98 | maskimg = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0) 99 | ImageDraw.Draw(maskimg).polygon(poly, outline=1, fill=1) 100 | mask = np.array(maskimg) 101 | #maskimg.show() 102 | 103 | for i in range(len(image_contents.classes)): 104 | 105 | obj_img = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0) 106 | ImageDraw.Draw(obj_img).polygon(image_contents.objects[i][0], outline=1, fill=1) 107 | obj = np.array(obj_img) 108 | logical_and = mask * obj 109 | 110 | if (np.sum(logical_and)>150): 111 | 112 | # Mask Polygon 113 | # Pad to ensure proper polygons for masks that touch image edges. 114 | padded_mask = np.zeros( 115 | (logical_and.shape[0] + 2, logical_and.shape[1] + 2), dtype=np.uint8) 116 | padded_mask[1:-1, 1:-1] = logical_and 117 | contours = find_contours(padded_mask, 0.5) 118 | 119 | strips[pos].objects.append([np.fliplr(verts) - 1 for verts in contours]) 120 | strips[pos].classes.append(image_contents.classes[i]) 121 | 122 | if ptr == 0: 123 | poly = list(map(tuple, np.flip(np.array(pts), 0))) 124 | self.all_images.append(strips) 125 | 126 | def polys_to_string(self, polys): 127 | 128 | ret = '' 129 | 130 | for poly in polys: 131 | for (x, y) in poly: 132 | ret+='%.2f %.2f '%(x, y) 133 | ret+='\n' 134 | return ret 135 | 136 | def find_poly_area(self, poly): 137 | 138 | x, y = np.zeros(len(poly)), np.zeros(len(poly)) 139 | for i, (xp, yp) in enumerate(poly): 140 | x[i] = xp 141 | y[i] = yp 142 | return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))) #shoelace algorithm 143 | 144 | def generate_samples(self, num, path): 145 | 146 | cumulative_mask = None 147 | text = '' 148 | 149 | if not os.path.exists(path): 150 | os.mkdir(path) 151 | 152 | for i in range(num): 153 | 154 | newImage = Image.new('RGBA', (self.im_shape[1], self.im_shape[0]), 0) 155 | 156 | text+="occ%d\n%s\n%d %d\n\n"%(i, os.path.join(path, 'occ_%d.jpg'%(i+1)), self.im_shape[0], self.im_shape[1]) 157 | 158 | for j in range(len(self.all_images[0])): 159 | 160 | rand = np.random.randint(len(self.all_images)) 161 | 162 | # create mask 163 | maskimg = Image.new('L', (self.im_shape[1], self.im_shape[0]), 0) 164 | 165 | ImageDraw.Draw(maskimg).polygon(self.polys[j], outline=1, fill=1) 166 | mask = np.array(maskimg) 167 | 168 | #Image.fromarray(mask*255, 'L').show() 169 | 170 | if cumulative_mask is None: 171 | cumulative_mask = mask 172 | else: 173 | cumulative_mask += mask 174 | 175 | #Image.fromarray(cumulative_mask*255, 'L').show() 176 | 177 | #time.sleep(.5) 178 | # assemble new image (uint8: 0-255) 179 | newImArray = np.empty(self.im_shape[:2]+(4,), dtype='uint8') 180 | 181 | # colors (three first columns, RGB) 182 | newImArray[:,:,:3] = self.images[rand][:,:,:3] 183 | 184 | # transparency (4th column) 185 | newImArray[:,:,3] = mask*255 186 | 187 | # back to Image from numpy 188 | 189 | newIm = Image.fromarray(newImArray, "RGBA") 190 | 191 | newImage.paste(newIm, (0, 0), newIm) 192 | 193 | for anns, cls in zip(self.all_images[rand][j].objects, self.all_images[rand][j].classes): 194 | text += cls+'\n' 195 | area = 0 196 | for poly in anns: 197 | area += self.find_poly_area(poly) 198 | text+='%.2f\n'%area 199 | text += self.polys_to_string(anns) 200 | text +='\n' 201 | 202 | background = Image.new("RGB", (newImArray.shape[1], newImArray.shape[0]), (0, 0, 0)) 203 | background.paste(newImage, mask=newImage.split()[3]) # 3 is the alpha channel 204 | 205 | background.save(os.path.join(path, 'occ_%d.jpg'%(i+1))) 206 | with open(os.path.join(path, 'occ_%d.txt'%(i+1)), 'w') as f: 207 | f.write(text) 208 | text = '' 209 | print ('Generated %d/%d Images: %s'%(i+1, num, os.path.join(path, 'occ_%d.jpg'%(i+1)))) 210 | 211 | if __name__=="__main__": 212 | 213 | ap = argparse.ArgumentParser() 214 | ap.add_argument("-i", "--image_dir", required=True, help="Path to the image dir") 215 | ap.add_argument("-o", "--output_dir", required=True, help="Path to the output dir") 216 | ap.add_argument("-s", "--strip_width", required=True, help="width of strip") 217 | ap.add_argument("-n", "--num_images", required=True, help="number of new images to generate") 218 | args = vars(ap.parse_args()) 219 | 220 | occlusion_gen = Occlusion_Generator(int(args['strip_width'])) 221 | 222 | occlusion_gen.generate_samples(int(args['num_images']), args['output_dir']) 223 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/sample.py: -------------------------------------------------------------------------------- 1 | import glob, os 2 | import shutil 3 | import argparse 4 | 5 | if __name__=='__main__': 6 | 7 | ap = argparse.ArgumentParser() 8 | ap.add_argument("-f", "--frames_dir", required=True, help="Path to the dir of videos") 9 | ap.add_argument("-n", "--frame_gap", required=True, help="Number of frames per sample") 10 | args = ap.parse_args() 11 | 12 | for folder in glob.glob(os.path.join(args.frames_dir, '*')): 13 | for i, f in enumerate(sorted(glob.glob(os.path.join(folder, '*')))): 14 | if(i%int(args.frame_gap)!=0): 15 | os.remove(f) 16 | -------------------------------------------------------------------------------- /coco_dataset_generator/extras/split_json_file.py: -------------------------------------------------------------------------------- 1 | import json 2 | import argparse 3 | 4 | def contains(splits): 5 | # Returns 1D binary map of images to take such that access is O(1) 6 | MAX, MIN = max([int(x.split('-')[-1]) for x in splits]), min([int(x.split('-')[0]) for x in splits]) 7 | A = [0 for _ in range(MAX-MIN+1)] 8 | for sp in splits: 9 | if '-' in sp: 10 | beg, end = [int(x) for x in sp.split('-')] 11 | else: 12 | beg = end = int(sp) 13 | for idx in range(beg-MIN, end+1-MIN): 14 | print (idx) 15 | A[idx] = 1 16 | 17 | return A, MIN, MAX 18 | 19 | if __name__=='__main__': 20 | 21 | ap = argparse.ArgumentParser() 22 | ap.add_argument('json', help='Path to JSON dataset file') 23 | ap.add_argument('split', nargs='+', help='Dataset split for splitting') 24 | ap.add_argument('--out', help='Path to output JSON file', default='cut_dataset.json') 25 | args = ap.parse_args() 26 | 27 | with open(args.json, 'r') as f: 28 | obj = json.load(f) 29 | 30 | A, MIN, MAX = contains(args.split) 31 | imgs, anns = [], [] 32 | for img in obj['images']: 33 | if img['id'] >= MIN and img['id'] <= MAX: 34 | if A[img['id']-MIN]: 35 | ANN = [ann for ann in obj['annotations'] if ann['image_id']==img['id']] 36 | anns.extend(ANN) 37 | imgs.append(img) 38 | 39 | with open(args.out, 'w') as f: 40 | 41 | json.dump({'images': imgs, 'annotations': anns, 'classes': obj['classes'], 'categories': []}, f) 42 | 43 | -------------------------------------------------------------------------------- /coco_dataset_generator/gui/contours.py: -------------------------------------------------------------------------------- 1 | from skimage.measure import find_contours as FC 2 | import numpy as np 3 | from simplification.cutil import simplify_coords 4 | 5 | def find_contours(*args): 6 | 7 | contours = FC(*args) 8 | 9 | simplified_contours = [np.array(simplify_coords(x, 1), dtype=np.int32) \ 10 | for x in contours] 11 | 12 | return simplified_contours 13 | 14 | -------------------------------------------------------------------------------- /coco_dataset_generator/gui/poly_editor.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from matplotlib.lines import Line2D 3 | from matplotlib.artist import Artist 4 | 5 | 6 | def dist(x, y): 7 | """ 8 | Return the distance between two points. 9 | """ 10 | d = x - y 11 | return np.sqrt(np.dot(d, d)) 12 | 13 | 14 | def dist_point_to_segment(p, s0, s1): 15 | """ 16 | Get the distance of a point to a segment. 17 | *p*, *s0*, *s1* are *xy* sequences 18 | This algorithm from 19 | http://geomalgorithms.com/a02-_lines.html 20 | """ 21 | v = s1 - s0 22 | w = p - s0 23 | c1 = np.dot(w, v) 24 | if c1 <= 0: 25 | return dist(p, s0) 26 | c2 = np.dot(v, v) 27 | if c2 <= c1: 28 | return dist(p, s1) 29 | b = c1 / c2 30 | pb = s0 + b * v 31 | return dist(p, pb) 32 | 33 | 34 | class PolygonInteractor(object): 35 | """ 36 | A polygon editor. 37 | 38 | Key-bindings 39 | 40 | 't' toggle vertex markers on and off. When vertex markers are on, 41 | you can move them, delete them 42 | 43 | 'd' delete the vertex under point 44 | 45 | 'i' insert a vertex at point. You must be within epsilon of the 46 | line connecting two existing vertices 47 | 48 | """ 49 | 50 | showverts = True 51 | epsilon = 5 # max pixel distance to count as a vertex hit 52 | 53 | def __init__(self, ax, poly): 54 | if poly.figure is None: 55 | raise RuntimeError('You must first add the polygon to a figure ' 56 | 'or canvas before defining the interactor') 57 | self.ax = ax 58 | canvas = poly.figure.canvas 59 | self.poly = poly 60 | 61 | x, y = zip(*self.poly.xy) 62 | self.line = Line2D(x, y, 63 | marker='o', markerfacecolor='r', 64 | animated=True) 65 | self.ax.add_line(self.line) 66 | 67 | self.cid = self.poly.add_callback(self.poly_changed) 68 | self._ind = None # the active vert 69 | 70 | canvas.mpl_connect('draw_event', self.draw_callback) 71 | canvas.mpl_connect('button_press_event', self.button_press_callback) 72 | canvas.mpl_connect('key_press_event', self.key_press_callback) 73 | canvas.mpl_connect('button_release_event', self.button_release_callback) 74 | canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) 75 | self.canvas = canvas 76 | 77 | def draw_callback(self, event): 78 | self.background = self.canvas.copy_from_bbox(self.ax.bbox) 79 | self.ax.draw_artist(self.poly) 80 | self.ax.draw_artist(self.line) 81 | # do not need to blit here, this will fire before the screen is 82 | # updated 83 | 84 | def poly_changed(self, poly): 85 | 'this method is called whenever the polygon object is called' 86 | # only copy the artist props to the line (except visibility) 87 | vis = self.line.get_visible() 88 | Artist.update_from(self.line, poly) 89 | self.line.set_visible(vis) # don't use the poly visibility state 90 | 91 | def get_ind_under_point(self, event): 92 | 'get the index of the vertex under point if within epsilon tolerance' 93 | 94 | # display coords 95 | xy = np.asarray(self.poly.xy) 96 | xyt = self.poly.get_transform().transform(xy) 97 | xt, yt = xyt[:, 0], xyt[:, 1] 98 | d = np.hypot(xt - event.x, yt - event.y) 99 | indseq, = np.nonzero(d == d.min()) 100 | ind = indseq[0] 101 | 102 | if d[ind] >= self.epsilon: 103 | ind = None 104 | 105 | return ind 106 | 107 | def button_press_callback(self, event): 108 | 'whenever a mouse button is pressed' 109 | if not self.showverts: 110 | return 111 | if event.inaxes is None: 112 | return 113 | if event.button != 1: 114 | return 115 | self._ind = self.get_ind_under_point(event) 116 | 117 | def button_release_callback(self, event): 118 | 'whenever a mouse button is released' 119 | if not self.showverts: 120 | return 121 | if event.button != 1: 122 | return 123 | self._ind = None 124 | 125 | def key_press_callback(self, event): 126 | 'whenever a key is pressed' 127 | if not event.inaxes: 128 | return 129 | if event.key == 't': 130 | self.showverts = not self.showverts 131 | self.line.set_visible(self.showverts) 132 | if not self.showverts: 133 | self._ind = None 134 | elif event.key == 'd': 135 | ind = self.get_ind_under_point(event) 136 | if ind is not None: 137 | self.poly.xy = np.delete(self.poly.xy, 138 | ind, axis=0) 139 | self.line.set_data(zip(*self.poly.xy)) 140 | elif event.key == 'i': 141 | xys = self.poly.get_transform().transform(self.poly.xy) 142 | p = event.x, event.y # display coords 143 | for i in range(len(xys) - 1): 144 | s0 = xys[i] 145 | s1 = xys[i + 1] 146 | d = dist_point_to_segment(p, s0, s1) 147 | if d <= self.epsilon: 148 | self.poly.xy = np.insert( 149 | self.poly.xy, i+1, 150 | [event.xdata, event.ydata], 151 | axis=0) 152 | self.line.set_data(zip(*self.poly.xy)) 153 | break 154 | if self.line.stale: 155 | self.canvas.draw_idle() 156 | 157 | def motion_notify_callback(self, event): 158 | 'on mouse movement' 159 | if not self.showverts: 160 | return 161 | if self._ind is None: 162 | return 163 | if event.inaxes is None: 164 | return 165 | if event.button != 1: 166 | return 167 | x, y = event.xdata, event.ydata 168 | 169 | self.poly.xy[self._ind] = x, y 170 | if self._ind == 0: 171 | self.poly.xy[-1] = x, y 172 | elif self._ind == len(self.poly.xy) - 1: 173 | self.poly.xy[0] = x, y 174 | self.line.set_data(zip(*self.poly.xy)) 175 | 176 | self.canvas.restore_region(self.background) 177 | self.ax.draw_artist(self.poly) 178 | self.ax.draw_artist(self.line) 179 | self.canvas.blit(self.ax.bbox) 180 | 181 | 182 | if __name__ == '__main__': 183 | import matplotlib.pyplot as plt 184 | from matplotlib.patches import Polygon 185 | 186 | theta = np.arange(0, 2*np.pi, 0.1) 187 | r = 1.5 188 | 189 | xs = r * np.cos(theta) 190 | ys = r * np.sin(theta) 191 | 192 | poly = Polygon(np.column_stack([xs, ys]), animated=True) 193 | 194 | fig, ax = plt.subplots() 195 | ax.add_patch(poly) 196 | p = PolygonInteractor(ax, poly) 197 | 198 | ax.set_title('Click and drag a point to move it') 199 | ax.set_xlim((-2, 2)) 200 | ax.set_ylim((-2, 2)) 201 | plt.show() 202 | 203 | -------------------------------------------------------------------------------- /coco_dataset_generator/gui/segment.py: -------------------------------------------------------------------------------- 1 | from matplotlib import pyplot as plt 2 | from matplotlib.collections import PatchCollection 3 | from matplotlib.patches import Polygon 4 | from matplotlib.widgets import RadioButtons 5 | from matplotlib.path import Path 6 | 7 | from PIL import Image 8 | import matplotlib 9 | 10 | import argparse 11 | import numpy as np 12 | import glob 13 | import os 14 | 15 | from matplotlib.widgets import Button 16 | from matplotlib.lines import Line2D 17 | from matplotlib.artist import Artist 18 | 19 | from .poly_editor import PolygonInteractor, dist_point_to_segment 20 | 21 | import sys 22 | from ..utils.visualize_dataset import return_info 23 | 24 | class COCO_dataset_generator(object): 25 | 26 | def __init__(self, fig, ax, args): 27 | 28 | self.ax = ax 29 | self.ax.set_yticklabels([]) 30 | self.ax.set_xticklabels([]) 31 | 32 | self.img_dir = args['image_dir'] 33 | self.index = 0 34 | self.fig = fig 35 | self.polys = [] 36 | self.zoom_scale, self.points, self.prev, self.submit_p, self.lines, self.circles = 1.2, [], None, None, [], [] 37 | 38 | self.zoom_id = fig.canvas.mpl_connect('scroll_event', self.zoom) 39 | self.click_id = fig.canvas.mpl_connect('button_press_event', self.onclick) 40 | self.clickrel_id = fig.canvas.mpl_connect('button_release_event', self.onclick_release) 41 | self.keyboard_id = fig.canvas.mpl_connect('key_press_event', self.onkeyboard) 42 | 43 | self.axradio = plt.axes([0.0, 0.0, 0.2, 1]) 44 | self.axbringprev = plt.axes([0.3, 0.05, 0.17, 0.05]) 45 | self.axreset = plt.axes([0.48, 0.05, 0.1, 0.05]) 46 | self.axsubmit = plt.axes([0.59, 0.05, 0.1, 0.05]) 47 | self.axprev = plt.axes([0.7, 0.05, 0.1, 0.05]) 48 | self.axnext = plt.axes([0.81, 0.05, 0.1, 0.05]) 49 | self.b_bringprev = Button(self.axbringprev, 'Bring Previous Annotations') 50 | self.b_bringprev.on_clicked(self.bring_prev) 51 | self.b_reset = Button(self.axreset, 'Reset') 52 | self.b_reset.on_clicked(self.reset) 53 | self.b_submit = Button(self.axsubmit, 'Submit') 54 | self.b_submit.on_clicked(self.submit) 55 | self.b_next = Button(self.axnext, 'Next') 56 | self.b_next.on_clicked(self.next) 57 | self.b_prev = Button(self.axprev, 'Prev') 58 | self.b_prev.on_clicked(self.previous) 59 | 60 | self.button_axes = [self.axbringprev, self.axreset, self.axsubmit, self.axprev, self.axnext, self.axradio] 61 | 62 | self.existing_polys = [] 63 | self.existing_patches = [] 64 | self.selected_poly = False 65 | self.objects = [] 66 | self.feedback = args['feedback'] 67 | 68 | self.right_click = False 69 | 70 | self.text = '' 71 | 72 | with open(args['class_file'], 'r') as f: 73 | self.class_names = [x.strip() for x in f.readlines() if x.strip() != ""] 74 | 75 | self.radio = RadioButtons(self.axradio, self.class_names) 76 | self.class_names = ('BG',) + tuple(self.class_names) 77 | 78 | self.img_paths = sorted(glob.glob(os.path.join(self.img_dir, '*.jpg'))) 79 | 80 | if len(self.img_paths)==0: 81 | self.img_paths = sorted(glob.glob(os.path.join(self.img_dir, '*.png'))) 82 | if os.path.exists(self.img_paths[self.index][:-3]+'txt'): 83 | self.index = len(glob.glob(os.path.join(self.img_dir, '*.txt'))) 84 | self.checkpoint = self.index 85 | 86 | try: 87 | im = Image.open(self.img_paths[self.index]) 88 | except IndexError: 89 | print ("Reached end of dataset! Delete some TXT files if you want to relabel some images in the folder") 90 | exit() 91 | 92 | width, height = im.size 93 | im.close() 94 | 95 | image = plt.imread(self.img_paths[self.index]) 96 | 97 | if args['feedback']: 98 | 99 | from mask_rcnn import model as modellib 100 | from mask_rcnn.get_json_config import get_demo_config 101 | 102 | #from skimage.measure import find_contours 103 | from .contours import find_contours 104 | 105 | from mask_rcnn.visualize_cv2 import random_colors 106 | 107 | config = get_demo_config(len(self.class_names)-2, True) 108 | 109 | if args['config_path'] is not None: 110 | config.from_json(args['config_path']) 111 | 112 | # Create model object in inference mode. 113 | model = modellib.MaskRCNN(mode="inference", model_dir='/'.join(args['weights_path'].split('/')[:-2]), config=config) 114 | 115 | # Load weights trained on MS-COCO 116 | model.load_weights(args['weights_path'], by_name=True) 117 | 118 | r = model.detect([image], verbose=0)[0] 119 | 120 | # Number of instances 121 | N = r['rois'].shape[0] 122 | 123 | masks = r['masks'] 124 | 125 | # Generate random colors 126 | colors = random_colors(N) 127 | 128 | # Show area outside image boundaries. 129 | height, width = image.shape[:2] 130 | 131 | class_ids, scores = r['class_ids'], r['scores'] 132 | 133 | for i in range(N): 134 | color = colors[i] 135 | 136 | # Label 137 | class_id = class_ids[i] 138 | score = scores[i] if scores is not None else None 139 | label = self.class_names[class_id] 140 | 141 | # Mask 142 | mask = masks[:, :, i] 143 | 144 | # Mask Polygon 145 | # Pad to ensure proper polygons for masks that touch image edges. 146 | padded_mask = np.zeros( 147 | (mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8) 148 | padded_mask[1:-1, 1:-1] = mask 149 | contours = find_contours(padded_mask, 0.5) 150 | 151 | for verts in contours: 152 | # Subtract the padding and flip (y, x) to (x, y) 153 | 154 | verts = np.fliplr(verts) - 1 155 | pat = PatchCollection([Polygon(verts, closed=True)], facecolor='green', linewidths=0, alpha=0.6) 156 | self.ax.add_collection(pat) 157 | self.objects.append(label) 158 | self.existing_patches.append(pat) 159 | self.existing_polys.append(Polygon(verts, closed=True, alpha=0.25, facecolor='red')) 160 | 161 | self.ax.imshow(image, aspect='auto') 162 | 163 | self.text+=str(self.index)+'\n'+os.path.abspath(self.img_paths[self.index])+'\n'+str(width)+' '+str(height)+'\n\n' 164 | 165 | def bring_prev(self, event): 166 | 167 | if not self.feedback: 168 | 169 | poly_verts, self.objects = return_info(self.img_paths[self.index-1][:-3]+'txt') 170 | 171 | for num in poly_verts: 172 | self.existing_polys.append(Polygon(num, closed=True, alpha=0.5, facecolor='red')) 173 | 174 | pat = PatchCollection([Polygon(num, closed=True)], facecolor='green', linewidths=0, alpha=0.6) 175 | self.ax.add_collection(pat) 176 | self.existing_patches.append(pat) 177 | 178 | def points_to_polygon(self): 179 | return np.reshape(np.array(self.points), (int(len(self.points)/2), 2)) 180 | 181 | def deactivate_all(self): 182 | self.fig.canvas.mpl_disconnect(self.zoom_id) 183 | self.fig.canvas.mpl_disconnect(self.click_id) 184 | self.fig.canvas.mpl_disconnect(self.clickrel_id) 185 | self.fig.canvas.mpl_disconnect(self.keyboard_id) 186 | 187 | def onkeyboard(self, event): 188 | 189 | if not event.inaxes: 190 | return 191 | elif event.key == 'a': 192 | 193 | if self.selected_poly: 194 | self.points = self.interactor.get_polygon().xy.flatten() 195 | self.interactor.deactivate() 196 | self.right_click = True 197 | self.selected_poly = False 198 | self.fig.canvas.mpl_connect(self.click_id, self.onclick) 199 | self.polygon.color = (0,255,0) 200 | self.fig.canvas.draw() 201 | else: 202 | for i, poly in enumerate(self.existing_polys): 203 | 204 | if poly.get_path().contains_point((event.xdata, event.ydata)): 205 | 206 | self.radio.set_active(self.class_names.index(self.objects[i])-1) 207 | self.polygon = self.existing_polys[i] 208 | self.existing_patches[i].set_visible(False) 209 | self.fig.canvas.mpl_disconnect(self.click_id) 210 | self.ax.add_patch(self.polygon) 211 | self.fig.canvas.draw() 212 | self.interactor = PolygonInteractor(self.ax, self.polygon) 213 | self.selected_poly = True 214 | self.existing_polys.pop(i) 215 | break 216 | 217 | elif event.key == 'r': 218 | 219 | for i, poly in enumerate(self.existing_polys): 220 | if poly.get_path().contains_point((event.xdata, event.ydata)): 221 | self.existing_patches[i].set_visible(False) 222 | self.existing_patches[i].remove() 223 | self.existing_patches.pop(i) 224 | self.existing_polys.pop(i) 225 | break 226 | self.fig.canvas.draw() 227 | 228 | def next(self, event): 229 | 230 | if len(self.text.split('\n'))>5: 231 | 232 | print (self.img_paths[self.index][:-3]+'txt') 233 | 234 | with open(self.img_paths[self.index][:-3]+'txt', "w") as text_file: 235 | text_file.write(self.text) 236 | 237 | self.ax.clear() 238 | 239 | self.ax.set_yticklabels([]) 240 | self.ax.set_xticklabels([]) 241 | 242 | if (self.indexself.checkpoint): 267 | self.index-=1 268 | #print (self.img_paths[self.index][:-3]+'txt') 269 | os.remove(self.img_paths[self.index][:-3]+'txt') 270 | 271 | self.ax.clear() 272 | 273 | self.ax.set_yticklabels([]) 274 | self.ax.set_xticklabels([]) 275 | 276 | image = plt.imread(self.img_paths[self.index]) 277 | self.ax.imshow(image, aspect='auto') 278 | 279 | im = Image.open(self.img_paths[self.index]) 280 | width, height = im.size 281 | im.close() 282 | 283 | self.reset_all() 284 | 285 | self.text+=str(self.index)+'\n'+os.path.abspath(self.img_paths[self.index])+'\n'+str(width)+' '+str(height)+'\n\n' 286 | 287 | def onclick(self, event): 288 | 289 | if not event.inaxes: 290 | return 291 | if not any([x.in_axes(event) for x in self.button_axes]): 292 | if event.button==1: 293 | self.points.extend([event.xdata, event.ydata]) 294 | #print (event.xdata, event.ydata) 295 | 296 | circle = plt.Circle((event.xdata,event.ydata),2.5,color='black') 297 | self.ax.add_artist(circle) 298 | self.circles.append(circle) 299 | 300 | if (len(self.points)<4): 301 | self.r_x = event.xdata 302 | self.r_y = event.ydata 303 | else: 304 | if len(self.points)>5: 305 | self.right_click=True 306 | self.fig.canvas.mpl_disconnect(self.click_id) 307 | self.click_id = None 308 | self.points.extend([self.points[0], self.points[1]]) 309 | #self.prev.remove() 310 | 311 | if (len(self.points)>2): 312 | line = self.ax.plot([self.points[-4], self.points[-2]], [self.points[-3], self.points[-1]], 'b--') 313 | self.lines.append(line) 314 | 315 | self.fig.canvas.draw() 316 | 317 | if len(self.points)>4: 318 | if self.prev: 319 | self.prev.remove() 320 | self.p = PatchCollection([Polygon(self.points_to_polygon(), closed=True)], facecolor='red', linewidths=0, alpha=0.4) 321 | self.ax.add_collection(self.p) 322 | self.prev = self.p 323 | 324 | self.fig.canvas.draw() 325 | 326 | #if len(self.points)>4: 327 | # print 'AREA OF POLYGON: ', self.find_poly_area(self.points) 328 | #print event.x, event.y 329 | 330 | def find_poly_area(self): 331 | coords = self.points_to_polygon() 332 | x, y = coords[:,0], coords[:,1] 333 | return (0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))))/2 #shoelace algorithm 334 | 335 | def onclick_release(self, event): 336 | 337 | if any([x.in_axes(event) for x in self.button_axes]) or self.selected_poly: 338 | return 339 | 340 | if hasattr(self, 'r_x') and hasattr(self, 'r_y') and None not in [self.r_x, self.r_y, event.xdata, event.ydata]: 341 | if np.abs(event.xdata - self.r_x)>10 and np.abs(event.ydata - self.r_y)>10: # 10 pixels limit for rectangle creation 342 | if len(self.points)<4: 343 | 344 | self.right_click=True 345 | self.fig.canvas.mpl_disconnect(self.click_id) 346 | self.click_id = None 347 | bbox = [np.min([event.xdata, self.r_x]), np.min([event.ydata, self.r_y]), np.max([event.xdata, self.r_x]), np.max([event.ydata, self.r_y])] 348 | self.r_x = self.r_y = None 349 | 350 | self.points = [bbox[0], bbox[1], bbox[0], bbox[3], bbox[2], bbox[3], bbox[2], bbox[1], bbox[0], bbox[1]] 351 | self.p = PatchCollection([Polygon(self.points_to_polygon(), closed=True)], facecolor='red', linewidths=0, alpha=0.4) 352 | self.ax.add_collection(self.p) 353 | self.fig.canvas.draw() 354 | 355 | def zoom(self, event): 356 | 357 | if not event.inaxes: 358 | return 359 | cur_xlim = self.ax.get_xlim() 360 | cur_ylim = self.ax.get_ylim() 361 | 362 | xdata = event.xdata # get event x location 363 | ydata = event.ydata # get event y location 364 | 365 | if event.button == 'down': 366 | # deal with zoom in 367 | scale_factor = 1 / self.zoom_scale 368 | elif event.button == 'up': 369 | # deal with zoom out 370 | scale_factor = self.zoom_scale 371 | else: 372 | # deal with something that should never happen 373 | scale_factor = 1 374 | print (event.button) 375 | 376 | new_width = (cur_xlim[1] - cur_xlim[0]) * scale_factor 377 | new_height = (cur_ylim[1] - cur_ylim[0]) * scale_factor 378 | 379 | relx = (cur_xlim[1] - xdata)/(cur_xlim[1] - cur_xlim[0]) 380 | rely = (cur_ylim[1] - ydata)/(cur_ylim[1] - cur_ylim[0]) 381 | 382 | self.ax.set_xlim([xdata - new_width * (1-relx), xdata + new_width * (relx)]) 383 | self.ax.set_ylim([ydata - new_height * (1-rely), ydata + new_height * (rely)]) 384 | self.ax.figure.canvas.draw() 385 | 386 | def reset(self, event): 387 | 388 | if not self.click_id: 389 | self.click_id = fig.canvas.mpl_connect('button_press_event', self.onclick) 390 | #print (len(self.lines)) 391 | #print (len(self.circles)) 392 | if len(self.points)>5: 393 | for line in self.lines: 394 | line.pop(0).remove() 395 | for circle in self.circles: 396 | circle.remove() 397 | self.lines, self.circles = [], [] 398 | self.p.remove() 399 | self.prev = self.p = None 400 | self.points = [] 401 | #print (len(self.lines)) 402 | #print (len(self.circles)) 403 | 404 | def print_points(self): 405 | 406 | ret = '' 407 | for x in self.points: 408 | ret+='%.2f'%x+' ' 409 | return ret 410 | 411 | def submit(self, event): 412 | 413 | if not self.right_click: 414 | print ('Right click before submit is a must!!') 415 | else: 416 | 417 | self.text+=self.radio.value_selected+'\n'+'%.2f'%self.find_poly_area()+'\n'+self.print_points()+'\n\n' 418 | self.right_click = False 419 | #print (self.points) 420 | 421 | self.lines, self.circles = [], [] 422 | self.click_id = fig.canvas.mpl_connect('button_press_event', self.onclick) 423 | 424 | self.polys.append(Polygon(self.points_to_polygon(), closed=True, color=np.random.rand(3), alpha=0.4, fill=True)) 425 | if self.submit_p: 426 | self.submit_p.remove() 427 | self.submit_p = PatchCollection(self.polys, cmap=matplotlib.cm.jet, alpha=0.4) 428 | self.ax.add_collection(self.submit_p) 429 | self.points = [] 430 | 431 | if __name__=='__main__': 432 | 433 | ap = argparse.ArgumentParser() 434 | ap.add_argument("-i", "--image_dir", required=True, help="Path to the image dir") 435 | ap.add_argument("-c", "--class_file", required=True, help="Path to the classes file of the dataset") 436 | ap.add_argument('-w', "--weights_path", default=None, help="Path to Mask RCNN checkpoint save file") 437 | ap.add_argument('-x', "--config_path", default=None, help="Path to Mask RCNN training config JSON file to load model based on specific parameters") 438 | args = vars(ap.parse_args()) 439 | 440 | args['feedback'] = args['weights_path'] is not None 441 | 442 | fig = plt.figure(figsize=(14, 14)) 443 | ax = plt.gca() 444 | 445 | gen = COCO_dataset_generator(fig, ax, args) 446 | 447 | plt.subplots_adjust(bottom=0.2) 448 | plt.show() 449 | 450 | gen.deactivate_all() 451 | -------------------------------------------------------------------------------- /coco_dataset_generator/gui/segment_bbox_only.py: -------------------------------------------------------------------------------- 1 | from matplotlib import pyplot as plt 2 | from matplotlib.collections import PatchCollection 3 | from matplotlib.patches import Polygon 4 | from matplotlib.widgets import RadioButtons 5 | from matplotlib.path import Path 6 | import matplotlib.patches as patches 7 | 8 | from PIL import Image 9 | import matplotlib 10 | 11 | import argparse 12 | import numpy as np 13 | import glob 14 | import os 15 | from matplotlib.widgets import RectangleSelector, Button, RadioButtons 16 | from matplotlib.lines import Line2D 17 | from matplotlib.artist import Artist 18 | 19 | from .poly_editor import PolygonInteractor 20 | 21 | from matplotlib.mlab import dist_point_to_segment 22 | import sys 23 | from ..utils.visualize_dataset import return_info 24 | 25 | import json 26 | 27 | from collections import defaultdict 28 | 29 | 30 | def read_JSON_file(f): 31 | 32 | with open(f, 'r') as g: 33 | d = json.loads(g.read()) 34 | 35 | img_paths = [x['file_name'] for x in d['images']] 36 | 37 | rects = [{'bbox': x['segmentation'][0], 'class': x['category_id'], 'image_id': x['image_id']} for x in d['annotations']] 38 | 39 | annotations = defaultdict(list) 40 | 41 | for rect in rects: 42 | r = rect['bbox'] 43 | x0, y0 = min(r[0], r[2], r[4], r[6]), min(r[1], r[3], r[5], r[7]) 44 | x1, y1 = max(r[0], r[2], r[4], r[6]), max(r[1], r[3], r[5], r[7]) 45 | 46 | r = patches.Rectangle((x0,y0),x1-x0,y1-y0,linewidth=1,edgecolor='g',facecolor='g', alpha=0.4) 47 | 48 | annotations[img_paths[rect['image_id']]].append({'bbox': r, 'cls': d['classes'][rect['class']-1]}) 49 | 50 | return d['classes'], img_paths, annotations 51 | 52 | 53 | class COCO_dataset_generator(object): 54 | 55 | def __init__(self, fig, ax, img_dir, classes, model_path, json_file): 56 | 57 | self.RS = RectangleSelector(ax, self.line_select_callback, 58 | drawtype='box', useblit=True, 59 | button=[1, 3], # don't use middle button 60 | minspanx=5, minspany=5, 61 | spancoords='pixels', 62 | interactive=True) 63 | 64 | ax.set_yticklabels([]) 65 | ax.set_xticklabels([]) 66 | 67 | #self.classes, self.img_paths, _ = read_JSON_file(json_file) 68 | with open(classes, 'r') as f: 69 | self.classes, img_paths = sorted([x.strip().split(',')[0] for x in f.readlines()]), glob.glob(os.path.abspath(os.path.join(img_dir, '*.jpg'))) 70 | plt.tight_layout() 71 | 72 | self.ax = ax 73 | self.fig = fig 74 | self.axradio = plt.axes([0.0, 0.0, 0.1, 1]) 75 | self.radio = RadioButtons(self.axradio, self.classes) 76 | self.zoom_scale = 1.2 77 | self.zoom_id = self.fig.canvas.mpl_connect('scroll_event', self.zoom) 78 | self.keyboard_id = self.fig.canvas.mpl_connect('key_press_event', self.onkeyboard) 79 | self.selected_poly = False 80 | self.axsave = plt.axes([0.81, 0.05, 0.1, 0.05]) 81 | self.b_save = Button(self.axsave, 'Save') 82 | self.b_save.on_clicked(self.save) 83 | self.objects, self.existing_patches, self.existing_rects = [], [], [] 84 | self.num_pred = 0 85 | if json_file is None: 86 | self.images, self.annotations = [], [] 87 | self.index = 0 88 | self.ann_id = 0 89 | else: 90 | with open(json_file, 'r') as g: 91 | d = json.loads(g.read()) 92 | self.images, self.annotations = d['images'], d['annotations'] 93 | self.index = len(self.images) 94 | self.ann_id = len(self.annotations) 95 | prev_files = [x['file_name'] for x in self.images] 96 | for i, f in enumerate(img_paths): 97 | im = Image.open(f) 98 | width, height = im.size 99 | dic = {'file_name': f, 'id': self.index+i, 'height': height, 'width': width} 100 | if f not in prev_files: 101 | self.images.append(dic) 102 | else: 103 | self.index+=1 104 | image = plt.imread(self.images[self.index]['file_name']) 105 | self.ax.imshow(image, aspect='auto') 106 | 107 | if not args['no_feedback']: 108 | from mask_rcnn.get_json_config import get_demo_config 109 | from mask_rcnn import model as modellib 110 | from mask_rcnn.visualize_cv2 import random_colors 111 | 112 | self.config = get_demo_config(len(self.classes)-1, True) 113 | 114 | if 'config_path' in args: 115 | self.config.from_json(args['config_path']) 116 | 117 | plt.connect('draw_event', self.persist) 118 | 119 | # Create model object in inference mode. 120 | self.model = modellib.MaskRCNN(mode="inference", model_dir='/'.join(args['weights_path'].split('/')[:-2]), config=self.config) 121 | 122 | # Load weights trained on MS-COCO 123 | self.model.load_weights(args['weights_path'], by_name=True) 124 | 125 | r = self.model.detect([image], verbose=0)[0] 126 | 127 | # Number of instances 128 | N = r['rois'].shape[0] 129 | 130 | masks = r['masks'] 131 | 132 | # Show area outside image boundaries. 133 | height, width = image.shape[:2] 134 | 135 | class_ids, scores, rois = r['class_ids'], r['scores'], r['rois'], 136 | 137 | for i in range(N): 138 | 139 | # Label 140 | class_id = class_ids[i] 141 | score = scores[i] if scores is not None else None 142 | label = self.classes[class_id-1] 143 | pat = patches.Rectangle((rois[i][1], rois[i][0]), rois[i][3]-rois[i][1], rois[i][2]-rois[i][0], linewidth=1, edgecolor='r',facecolor='r', alpha=0.4) 144 | rect = self.ax.add_patch(pat) 145 | 146 | self.objects.append(label) 147 | self.existing_patches.append(pat.get_bbox().get_points()) 148 | self.existing_rects.append(pat) 149 | self.num_pred = len(self.objects) 150 | 151 | def line_select_callback(self, eclick, erelease): 152 | 'eclick and erelease are the press and release events' 153 | x1, y1 = eclick.xdata, eclick.ydata 154 | x2, y2 = erelease.xdata, erelease.ydata 155 | 156 | def zoom(self, event): 157 | 158 | if not event.inaxes: 159 | return 160 | cur_xlim = self.ax.get_xlim() 161 | cur_ylim = self.ax.get_ylim() 162 | 163 | xdata = event.xdata # get event x location 164 | ydata = event.ydata # get event y location 165 | 166 | if event.button == 'down': 167 | # deal with zoom in 168 | scale_factor = 1 / self.zoom_scale 169 | elif event.button == 'up': 170 | # deal with zoom out 171 | scale_factor = self.zoom_scale 172 | else: 173 | # deal with something that should never happen 174 | scale_factor = 1 175 | print (event.button) 176 | 177 | new_width = (cur_xlim[1] - cur_xlim[0]) * scale_factor 178 | new_height = (cur_ylim[1] - cur_ylim[0]) * scale_factor 179 | 180 | relx = (cur_xlim[1] - xdata)/(cur_xlim[1] - cur_xlim[0]) 181 | rely = (cur_ylim[1] - ydata)/(cur_ylim[1] - cur_ylim[0]) 182 | 183 | self.ax.set_xlim([xdata - new_width * (1-relx), xdata + new_width * (relx)]) 184 | self.ax.set_ylim([ydata - new_height * (1-rely), ydata + new_height * (rely)]) 185 | self.ax.figure.canvas.draw() 186 | 187 | def save(self, event): 188 | 189 | data = {'images':self.images[:self.index+1], 'annotations':self.annotations, 'categories':[], 'classes': self.classes} 190 | 191 | with open('output.json', 'w') as outfile: 192 | json.dump(data, outfile) 193 | 194 | def persist(self, event): 195 | if self.RS.active: 196 | self.RS.update() 197 | 198 | def onkeyboard(self, event): 199 | 200 | if not event.inaxes: 201 | return 202 | elif event.key == 'a': 203 | for i, ((xmin, ymin), (xmax, ymax)) in enumerate(self.existing_patches): 204 | if xmin<=event.xdata<=xmax and ymin<=event.ydata<=ymax: 205 | self.radio.set_active(self.classes.index(self.objects[i])) 206 | self.RS.set_active(True) 207 | self.rectangle = self.existing_rects[i] 208 | self.rectangle.set_visible(False) 209 | coords = self.rectangle.get_bbox().get_points() 210 | self.RS.extents = [coords[0][0], coords[1][0], coords[0][1], coords[1][1]] 211 | self.RS.to_draw.set_visible(True) 212 | self.fig.canvas.draw() 213 | self.existing_rects.pop(i) 214 | self.existing_patches.pop(i) 215 | self.objects.pop(i) 216 | fig.canvas.draw() 217 | break 218 | 219 | elif event.key == 'i': 220 | b = self.RS.extents # xmin, xmax, ymin, ymax 221 | b = [int(x) for x in b] 222 | if b[1]-b[0]>0 and b[3]-b[2]>0: 223 | poly = [b[0], b[2], b[0], b[3], b[1], b[3], b[1], b[2], b[0], b[2]] 224 | area = (b[1]-b[0])*(b[3]-b[2]) 225 | bbox = [b[0], b[2], b[1], b[3]] 226 | dic2 = {'segmentation': [poly], 'area': area, 'iscrowd':0, 'image_id':self.index, 'bbox':bbox, 'category_id': self.classes.index(self.radio.value_selected)+1, 'id': self.ann_id} 227 | if dic2 not in self.annotations: 228 | self.annotations.append(dic2) 229 | self.ann_id+=1 230 | rect = patches.Rectangle((b[0],b[2]),b[1]-b[0],b[3]-b[2],linewidth=1,edgecolor='g',facecolor='g', alpha=0.4) 231 | self.ax.add_patch(rect) 232 | 233 | self.RS.set_active(False) 234 | 235 | self.fig.canvas.draw() 236 | self.RS.set_active(True) 237 | elif event.key in ['N', 'n']: 238 | self.ax.clear() 239 | self.index+=1 240 | if (len(self.objects)==self.num_pred): 241 | self.images.pop(self.index-1) 242 | self.index-=1 243 | if self.index==len(self.images): 244 | exit() 245 | image = plt.imread(self.images[self.index]['file_name']) 246 | self.ax.imshow(image) 247 | self.ax.set_yticklabels([]) 248 | self.ax.set_xticklabels([]) 249 | r = self.model.detect([image], verbose=0)[0] 250 | 251 | # Number of instances 252 | N = r['rois'].shape[0] 253 | 254 | masks = r['masks'] 255 | 256 | # Show area outside image boundaries. 257 | height, width = image.shape[:2] 258 | 259 | class_ids, scores, rois = r['class_ids'], r['scores'], r['rois'], 260 | self.existing_rects, self.existing_patches, self.objects = [], [], [] 261 | for i in range(N): 262 | 263 | # Label 264 | class_id = class_ids[i] 265 | score = scores[i] if scores is not None else None 266 | label = self.classes[class_id-1] 267 | pat = patches.Rectangle((rois[i][1], rois[i][0]), rois[i][3]-rois[i][1], rois[i][2]-rois[i][0], linewidth=1, edgecolor='r',facecolor='r', alpha=0.4) 268 | rect = self.ax.add_patch(pat) 269 | 270 | self.objects.append(label) 271 | 272 | self.existing_patches.append(pat.get_bbox().get_points()) 273 | self.existing_rects.append(pat) 274 | 275 | self.num_pred = len(self.objects) 276 | self.fig.canvas.draw() 277 | 278 | elif event.key in ['q','Q']: 279 | exit() 280 | 281 | if __name__=='__main__': 282 | 283 | ap = argparse.ArgumentParser() 284 | ap.add_argument("-i", "--image_file", required=True, help="Path to the images dir") 285 | ap.add_argument("-c", "--classes_file", required=True, help="Path to classes file") 286 | ap.add_argument("-j", "--json_file", required=False, help="Path of JSON file to append dataset to", default=None) 287 | ap.add_argument("--save_csv", required=False, action="store_true", help="Choose option to save dataset as CSV file annotations.csv") 288 | ap.add_argument('-w', "--weights_path", default=None, help="Path to Mask RCNN checkpoint save file") 289 | ap.add_argument('-x', "--config_path", default=None, help="Path to Mask RCNN JSON config file") 290 | args = vars(ap.parse_args()) 291 | 292 | args["no_feedback"] = 'weights_path' not in args 293 | 294 | fig = plt.figure(figsize=(14, 14)) 295 | ax = plt.gca() 296 | 297 | gen = COCO_dataset_generator(fig, ax, args['image_file'], args['classes_file'], args['weights_path'], args['json_file']) 298 | 299 | plt.subplots_adjust(bottom=0.2) 300 | plt.show() 301 | -------------------------------------------------------------------------------- /coco_dataset_generator/utils/create_json_file.py: -------------------------------------------------------------------------------- 1 | #coding: utf8 2 | import xml.etree.cElementTree as ET 3 | import glob 4 | import argparse 5 | import os 6 | import numpy as np 7 | import json 8 | import unicodedata 9 | 10 | from PIL import Image 11 | from ..gui.segment import COCO_dataset_generator as cocogen 12 | 13 | if __name__=='__main__': 14 | 15 | ap = argparse.ArgumentParser() 16 | ap.add_argument("-i", "--image_dir", required=True, help="Path to the image dir") 17 | ap.add_argument("-o", "--file_path", required=True, help="Path of output file") 18 | ap.add_argument("-c", "--class_file", required=True, help="Path of file with output classes") 19 | ap.add_argument("-t", "--type", required=True, help="Type of the image files") 20 | args = vars(ap.parse_args()) 21 | 22 | with open(args['class_file'], 'r') as f: 23 | classes = sorted([unicodedata.normalize('NFKD', x).strip() for x in f.readlines()]) 24 | 25 | images, anns = [], [] 26 | 27 | img_paths = [x for x in glob.glob(os.path.join(args['image_dir'], '*.'+args['type'])) if os.path.exists(x[:-3]+'txt')] 28 | 29 | num_imgs = len(img_paths) 30 | i=0 31 | for f in sorted(img_paths): 32 | 33 | img = Image.open(f) 34 | width, height = img.size 35 | dic = {'file_name': f, 'id': i, 'height': height, 'width': width} 36 | images.append(dic) 37 | 38 | ann_index = 0 39 | 40 | for i, f in enumerate(sorted(glob.glob(os.path.join(os.path.abspath(args['image_dir']), '*.txt')))): 41 | ptr = 0 42 | with open(f, 'r', encoding='utf-8-sig') as g: 43 | s = g.read() 44 | s = s.split('\n')[2:-1] 45 | 46 | width, height = [int(x) for x in s[0].split(' ')] 47 | s = s[2:] 48 | print (s) 49 | while(ptrptr+3 and s[ptr+3] != '': 58 | ind = ptr + 3 59 | while (indx2: 76 | x2 = points.max(0)[0] 77 | if points.max(0)[1]>y2: 78 | y2 = points.max(0)[1] 79 | 80 | bbox = [x2, y2, x1, y1] 81 | dic2 = {'segmentation': poly, 'area': area, 'iscrowd':0, 'image_id':i, 'bbox':bbox, 'category_id': cat_id, 'id': ann_index} 82 | ann_index+=1 83 | ptr+=4 84 | anns.append(dic2) 85 | 86 | data = {'images':images, 'annotations':anns, 'categories':[], 'classes': classes} 87 | 88 | with open(args['file_path']+'.json', 'w') as outfile: 89 | json.dump(data, outfile) 90 | -------------------------------------------------------------------------------- /coco_dataset_generator/utils/delete_images.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | import matplotlib.patches as patches 6 | 7 | def press(event): 8 | 9 | if event.key.lower() == 'q': 10 | exit() 11 | 12 | if event.key.lower() == 'd': 13 | for ann in anns: 14 | new_anns.remove(ann) 15 | print ('Deleted image:', new_imgs[ptr]["file_name"], "from the dataset!") 16 | del new_imgs[ptr] 17 | 18 | if event.key.lower() == 'j': 19 | print ("Saving dataset to file! Please wait!") 20 | 21 | # Account for deletions by changing label space 22 | 23 | id_list = [int(img['id']) for img in new_imgs] 24 | ann_list = [int(ann['id']) for ann in new_anns] 25 | 26 | full_img, full_ann = [x for x in range(len(id_list))], [x for x in range(len(ann_list))] 27 | 28 | free_img, free_ann = list(set(full_img)-set(id_list)), list(set(full_ann)-set(ann_list)) 29 | change_img, change_ann = list(set(id_list)-set(full_img)), list(set(ann_list)-set(full_ann)) 30 | 31 | for f, c in zip(free_img, change_img): 32 | for img in new_imgs: 33 | if img['id']==c: 34 | img['id']=f 35 | for ann in new_anns: 36 | if ann['image_id']==c: 37 | ann['image_id']=f 38 | 39 | for f, c in zip(free_ann, change_ann): 40 | for ann in new_anns: 41 | if ann['id']==c: 42 | ann['id']=f 43 | 44 | data = {'images': new_imgs, 'annotations': new_anns, 'categories':[], 'classes':classes} 45 | with open('deleted_dataset.json', 'w') as f: 46 | json.dump(data, f) 47 | print ("Dataset saved!") 48 | else: 49 | plt.close() 50 | 51 | 52 | if __name__=='__main__': 53 | 54 | ap = argparse.ArgumentParser() 55 | ap.add_argument('--json_file', required=True, help='Path to JSON file') 56 | args = ap.parse_args() 57 | 58 | with open(args.json_file, 'r') as f: 59 | obj = json.load(f) 60 | 61 | images, annotations = obj["images"], obj["annotations"] 62 | classes = obj["classes"] 63 | 64 | print ("Total number of images in dataset: ", len(images)) 65 | 66 | new_imgs, new_anns = images, annotations 67 | 68 | for ptr, img in enumerate(images): 69 | 70 | fig, ax = plt.subplots() 71 | plt.tick_params(axis='both', which='both', bottom='off', top='off', 72 | labelbottom='off', right='off', left='off', labelleft='off') 73 | 74 | fig.canvas.mpl_connect('key_press_event', press) 75 | ax.set_title('d - Delete image; j - Save dataset; q - Exit; others - Next image') 76 | 77 | anns = [ann for ann in annotations if ann["image_id"]==img["id"]] 78 | image = plt.imread(img["file_name"]) 79 | plt.imshow(image) 80 | for ann in anns: 81 | s = [int(x) for x in ann['bbox']] 82 | rect = patches.Rectangle((s[0],s[1]),s[2]-s[0],s[3]-s[1],linewidth=1,edgecolor='r',facecolor='none') 83 | ax = plt.gca() 84 | ax.add_patch(rect) 85 | plt.text(s[0]-10, s[1]+10, classes[ann['category_id']-1]) 86 | plt.show() 87 | 88 | print ("Saving dataset to file! Please wait!") 89 | data = {'images': new_imgs, 'annotations': new_anns, 'categories':[], 'classes':classes} 90 | with open('deleted_dataset.json', 'w') as f: 91 | json.dump(data, f) 92 | print ("Dataset saved!") 93 | -------------------------------------------------------------------------------- /coco_dataset_generator/utils/pascal_to_coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import glob 4 | import xml.etree.ElementTree as ET 5 | import json 6 | 7 | from json import encoder 8 | encoder.FLOAT_REPR = lambda o: format(o, '.2f') 9 | 10 | import cv2 11 | 12 | if __name__=='__main__': 13 | 14 | 15 | ap = argparse.ArgumentParser(description='Convert PASCAL VOC format dataset to COCO style dataset') 16 | ap.add_argument("-d", "--pascal_dir", required=True, help="Path to the PASCAL VOC style dataset") 17 | ap.add_argument("-c", "--class_file", required=True, help="Path to the classes in the dataset") 18 | ap.add_argument("-n", "--file_name", required=True, help="Name of output JSON file") 19 | args = vars(ap.parse_args()) 20 | 21 | with open(args['class_file'], 'r') as f: 22 | classes = [x.strip() for x in f.readlines()] 23 | 24 | images, anns = [], [] 25 | ann_index = 0 26 | 27 | for i, f in enumerate(glob.glob(os.path.join(args['pascal_dir'], 'JPEGImages/*.png'))): 28 | annot = os.path.join(args['pascal_dir'], 'Annotations', f.split('/')[-1][:-3]+'xml') 29 | tree = ET.parse(annot) 30 | root = tree.getroot() 31 | 32 | img = cv2.imread(f) 33 | height, width, _ = img.shape 34 | dic = {'file_name': os.path.abspath(f), 'id': i, 'height': height, 'width': width} 35 | images.append(dic) 36 | 37 | for obj in root.findall('object'): 38 | 39 | cls_id = classes.index(obj.find('name').text)+1 40 | bx = [int(obj.find('bndbox').find('xmax').text), int(obj.find('bndbox').find('ymax').text), int(obj.find('bndbox').find('xmin').text), int(obj.find('bndbox').find('ymin').text)] 41 | pts = [bx[0], bx[1], bx[2], bx[1], bx[2], bx[3], bx[0], bx[3], bx[0], bx[1]] # Create mask as the bounding box itself 42 | area = (bx[0]-bx[2])*(bx[1]-bx[3]) 43 | dic2 = {'segmentation': [pts], 'area': area, 'iscrowd':0, 'image_id':i, 'bbox':bx, 'category_id': cls_id, 'id': ann_index} 44 | ann_index+=1 45 | anns.append(dic2) 46 | 47 | data = {'images':images, 'annotations':anns, 'categories':[], 'classes':classes} 48 | 49 | with open(args['file_name']+'.json', 'w') as outfile: 50 | json.dump(data, outfile) 51 | -------------------------------------------------------------------------------- /coco_dataset_generator/utils/visualize_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | from matplotlib import pyplot as plt 3 | import sys 4 | import numpy as np 5 | import matplotlib 6 | from matplotlib.patches import Polygon 7 | from matplotlib.collections import PatchCollection 8 | from PIL import Image 9 | import glob 10 | import argparse 11 | 12 | def return_info(filename): 13 | 14 | polys, objects = [], [] 15 | with open(filename, 'r') as f: 16 | txt = f.read().split('\n') 17 | 18 | index = 6 19 | #for index in range(6, len(txt), 4): 20 | while (index < len(txt)): 21 | #print (txt[index]) 22 | numbers = txt[index].split(' ') if txt[index].split(' ')[-1]!='' else txt[index].split(' ')[:-1] 23 | num = [float(x) for x in numbers] 24 | 25 | num = np.reshape(np.array(num), (int(len(num)/2), 2)) 26 | polys.append(num) 27 | objects.append(txt[index-2]) 28 | 29 | while index+1 25: 56 | q = 113 57 | 58 | if q == 113: # if q == 'q' 59 | exit() 60 | -------------------------------------------------------------------------------- /gui.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI/35c095a815a6ef97d2e820d175bf46e32e59f40b/gui.png -------------------------------------------------------------------------------- /images/davi_0_1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI/35c095a815a6ef97d2e820d175bf46e32e59f40b/images/davi_0_1.jpg -------------------------------------------------------------------------------- /images/davi_2_71.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI/35c095a815a6ef97d2e820d175bf46e32e59f40b/images/davi_2_71.jpg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | certifi==2019.11.28 2 | cycler==0.10.0 3 | Cython==0.29.14 4 | decorator==4.4.1 5 | imageio==2.6.1 6 | kiwisolver==1.1.0 7 | matplotlib==3.1.3 8 | networkx==2.4 9 | numpy 10 | opencv-python 11 | Pillow 12 | pyparsing==2.4.6 13 | python-dateutil==2.8.1 14 | PyWavelets==1.1.1 15 | scikit-image 16 | scipy 17 | six==1.14.0 18 | -------------------------------------------------------------------------------- /requirements_maskrcnn.txt: -------------------------------------------------------------------------------- 1 | imutils==0.5.2 2 | -e git+https://www.github.com/hanskrupakar/Mask_RCNN@master#egg=mask_rcnn 3 | pycocotools==2.0.0 4 | simplification 5 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | import os 3 | import sys 4 | import subprocess 5 | 6 | def install(package): 7 | subprocess.call([sys.executable, "-m", "pip", "install", package]) 8 | 9 | if os.getenv('MASK_RCNN'): 10 | fl = 'requirements_maskrcnn.txt' 11 | else: 12 | fl = 'requirements.txt' 13 | 14 | with open(fl, 'r') as f: 15 | packs = [x.strip() for x in f.readlines()] 16 | 17 | for p in packs: 18 | install(p) 19 | 20 | dependencies = [] 21 | 22 | packages = [ 23 | package for package in find_packages() if package.startswith('coco_dataset_generator') 24 | ] 25 | 26 | setup(name='coco_dataset_generator', 27 | version='1.0', 28 | description='COCO Style Dataset Generator GUI', 29 | author='hanskrupakar', 30 | author_email='hansk@nyu.edu', 31 | license='Open-Source', 32 | url='https://www.github.com/hanskrupakar/COCO-Style-Dataset-Generator-GUI', 33 | packages=packages, 34 | install_requires=dependencies, 35 | test_suite='unit_tests', 36 | ) 37 | -------------------------------------------------------------------------------- /unit_tests/test_button.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from matplotlib.widgets import Button 4 | 5 | freqs = np.arange(2, 20, 3) 6 | 7 | fig, ax = plt.subplots() 8 | plt.subplots_adjust(bottom=0.2) 9 | t = np.arange(0.0, 1.0, 0.001) 10 | s = np.sin(2*np.pi*freqs[0]*t) 11 | l, = plt.plot(t, s, lw=2) 12 | 13 | 14 | class Index(object): 15 | ind = 0 16 | 17 | def next(self, event): 18 | self.ind += 1 19 | i = self.ind % len(freqs) 20 | ydata = np.sin(2*np.pi*freqs[i]*t) 21 | l.set_ydata(ydata) 22 | plt.draw() 23 | 24 | def prev(self, event): 25 | self.ind -= 1 26 | i = self.ind % len(freqs) 27 | ydata = np.sin(2*np.pi*freqs[i]*t) 28 | l.set_ydata(ydata) 29 | plt.draw() 30 | 31 | callback = Index() 32 | axprev = plt.axes([0.7, 0.05, 0.1, 0.075]) 33 | axnext = plt.axes([0.81, 0.05, 0.1, 0.075]) 34 | bnext = Button(axnext, 'Next') 35 | bnext.on_clicked(callback.next) 36 | bprev = Button(axprev, 'Previous') 37 | bprev.on_clicked(callback.prev) 38 | 39 | plt.show() 40 | -------------------------------------------------------------------------------- /unit_tests/test_coco.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | from matplotlib.collections import PatchCollection 3 | from matplotlib.patches import Polygon 4 | 5 | polygons = [Polygon([[510.66, 423.01], [511.72, 420.03], [510.45, 416], [510.34, 413.02], [510.77, 410.26], [510.77, 407.5], [510.34, 405.16], [511.51, 402.83], [511.41, 400.49], [510.24, 398.16], [509.39, 397.31], [504.61, 399.22], [502.17, 399.64], [500.89, 401.66], [500.47, 402.08], [499.09, 401.87], [495.79, 401.98], [490.59, 401.77], [488.79, 401.77], [485.39, 398.58], [483.9, 397.31], [481.56, 396.35], [478.48, 395.93], [476.68, 396.03], [475.4, 396.77], [473.92, 398.79], [473.28, 399.96], [473.49, 401.87], [474.56, 403.47], [473.07, 405.59], [473.39, 407.71], [476.68, 409.41], [479.23, 409.73], [481.56, 410.69], [480.4, 411.85], [481.35, 414.93], [479.86, 418.65], [477.32, 420.03], [476.04, 422.58], [479.02, 422.58], [480.29, 423.01], [483.79, 419.93], [486.66, 416.21], [490.06, 415.57], [492.18, 416.85], [491.65, 420.24], [492.82, 422.9], [493.56, 424.39], [496.43, 424.6], [498.02, 423.01], [498.13, 421.31], [497.07, 420.03], [497.07, 415.15], [496.33, 414.51], [501.1, 411.96], [502.06, 411.32], [503.02, 415.04], [503.33, 418.12], [501.1, 420.24], [498.98, 421.63], [500.47, 424.39], [505.03, 423.32], [506.2, 421.31], [507.69, 419.5], [506.31, 423.32], [510.03, 423.01], [510.45, 423.01]])] 6 | 7 | fig = plt.figure() 8 | ax = plt.gca() 9 | 10 | img = plt.imread('/home/hans/Datasets/COCO/val2017/000000289343.jpg') 11 | ax.imshow(img) 12 | 13 | p = PatchCollection(polygons, facecolor='red', linewidths=0, alpha=0.4) 14 | ax.add_collection(p) 15 | 16 | p = PatchCollection(polygons, facecolor='none', edgecolors='blue', linewidths=2) 17 | ax.add_collection(p) 18 | 19 | #plt.plot([1.,2.], [2.,3.]) 20 | plt.show() 21 | -------------------------------------------------------------------------------- /unit_tests/test_poly_pasting.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | from PIL import Image, ImageDraw 3 | 4 | # read image as RGB and add alpha (transparency) 5 | im = Image.open("images/davi_0_1.jpg").convert("RGBA") 6 | 7 | # convert to numpy (for convenience) 8 | imArray = numpy.asarray(im) 9 | 10 | # create mask 11 | polygon2 = [(10, 10), (10, 150), (550, 150), (150, 10)] 12 | polygon1 = [(1920, 0), (0, 0), (0, 250), (-0.10210674635236394, 249.77657923687673), (-2.1094657996870403, 255.03379664713123), (17.043357362773406, 256.0158194361004), (19.743585566230458, 257.24042105198055), (15.950897656114186, 254.2993740345912), (27.747904611277534, 250.51951337704952), (27.961274007008967, 253.63518034864092), (29.81425998511483, 255.41115415383845), (37.3667088625533, 256.4373058829165), (46.53380792054146, 243.57254718554756), (53.61176555130789, 244.36644634912406), (49.558319074426926, 245.93258750036037), (60.33955511927042, 255.83254114007315), (65.73851977273596, 255.91269387461898), (76.73575480091465, 247.02704108128736), (78.34805771376494, 248.17030289093188), (86.05632662995261, 258.41364572544853), (87.85472059685605, 251.9630660275433), (89.80188097115189, 244.4785932629154), (92.31760033899944, 253.8264647999193), (103.48680386603161, 259.5970120020665), (108.79236707446952, 260.23761745739506), (107.4937793054668, 254.30466090801215), (108.39358602655307, 246.14079063076383), (126.52486800617761, 257.098422070827), (126.89943745269167, 250.1957195389658), (132.02691150565755, 256.48412988536154), (128.58195282301338, 252.79775783620423), (136.47572160588913, 254.4441395133314), (146.75086407853183, 260.1350693035284), (155.073813275849, 254.85850113146446), (151.97302312432186, 251.33687712540114), (157.34863597011744, 256.9711603194696), (165.3621223224169, 259.44430066082975), (169.10946867961437, 248.41407201625222), (181.09443326356882, 254.04993377735954), (175.70858911880558, 258.2380301719833), (192.00679423643666, 252.91711095845895), (189.1673649545458, 252.21181387583448), (195.83368997955725, 259.6731833553506), (194.61216526235634, 257.67985807616327), (199.59668919652594, 256.86696909739544), (204.89929126562438, 256.3404560825), (209.62014250508463, 255.07011440400703), (212.07706747647595, 261.4557434304978), (229.96134579284376, 255.64160167491366), (232.43428089908423, 251.83975180535887), (229.90018481610127, 257.0135552875774), (234.80672832203157, 248.76310236447264), (243.25042358036447, 261.17495400508494), (249.68932632608966, 264.0681882162521), (255.74368446785104, 262.8189134037273), (255.5330016557221, 254.57174601013634), (260.2653123071998, 262.97142577117904), (264.09858509159926, 258.1544814168108), (275.24555974843247, 256.19194127041214), (286.24564412575853, 259.7104454066601), (284.71754162529976, 253.71317013201232), (293.4762533829063, 253.6635985653496), (290.6949362189203, 265.3365355964029), (294.1761306148136, 262.35245809288574), (310.2766731170838, 256.2427436783204), (317.2399732326495, 257.19080664472693), (318.39042410523496, 265.2841360754873), (319.912104059165, 264.8216512714212), (322.3480129044327, 261.86773581034964), (335.9036887305486, 264.88653554479106), (332.97228370177675, 259.16014609412684), (345.1352331564961, 266.03635464878687), (343.0437539721002, 260.35448821806597), (345.6416745691777, 260.1597059242849), (358.9043367621447, 258.522277196598), (355.00847219311873, 267.1244254191321), (367.51940391677755, 262.3069070377452), (369.4275875033949, 259.7016001436432), (371.9099373501995, 257.8257422068458), (374.5192933499681, 262.0863919259926), (383.172995439416, 265.1345723510826), (396.01114440021024, 263.90263357113616), (390.24183439053525, 267.98438605429095), (394.50876485297516, 257.40461928296986), (404.1138037652838, 267.63878806317155), (407.84839944051527, 259.36014224577974), (410.919296661234, 261.85406598934696), (413.3293143874036, 263.91645457293714), (419.7131771063948, 259.54884010662346), (427.7662396620981, 255.74355444558074), (438.56674757484996, 267.1465872195995), (438.1531997489801, 255.3327554352958), (441.19632661544046, 267.60841321814763), (450.43222343773203, 267.4037487393895), (455.0688973027098, 263.37794907150123), (453.4182917377982, 263.8641406049655), (469.5434250354823, 255.67725466245275), (470.0239369093112, 263.1880594155271), (482.66468791948625, 264.4754668271869), (476.56821722377003, 260.42230020262724), (482.40435887795024, 257.6951553927151), (487.11101097139226, 269.3935423402255), (491.99594703611876, 268.42034465075585), (505.8730055363312, 269.9189989091135), (511.36354389993227, 268.8852381895841), (508.8427022427122, 255.49268088650217), (513.1562881258143, 268.7633911410157), (526.2263542918532, 258.19196773685854), (526.0073515650495, 267.1479755341373), (529.7346706077001, 255.86130326839367), (536.4154971134869, 267.2665393334687), (544.2924475692166, 268.4048756073015), (550.7572294528655, 263.18181636674143), (542.5747604062316, 266.969882094905), (553.2932763962233, 258.29260061906234), (561.568113536864, 261.6030899809526), (560.7561201526694, 261.4482858902076), (564.4975123154929, 262.25141631808634), (581.2948318079701, 265.66724825588926), (581.9996036242798, 273.0651590172597), (589.3393708500435, 263.18587292520607), (590.827050456112, 272.29054248597595), (601.3159591714563, 258.876416064304), (605.6895264461832, 263.77240113762724), (598.5479944983696, 264.6687239762898), (615.5179812848138, 265.3864130446329), (608.9819231550235, 263.2508931106562), (620.421177359739, 270.54726731145485), (625.3075827317724, 259.75011186505895), (622.98436600077, 262.85296860715397), (639.2744468627925, 264.78305853530514), (643.1987841328669, 271.51223837781464), (649.2456961792811, 270.7200332338785), (650.2218274586423, 267.27606166920185), (647.56809791391, 270.2774386070881), (656.0993885240471, 266.34473794767604), (661.2246165387108, 261.55653405025583), (666.745638042081, 270.64838369739886), (673.1379136890687, 275.33916998409853), (682.9043535840688, 274.86915990335575), (678.5437019799954, 261.9133566885317), (688.5923356626867, 269.0599110304343), (693.7563354659698, 274.336248899017), (693.0202606374799, 263.60496018648337), (698.4000849399749, 261.7219308217077), (712.7875555474909, 272.6578667172132), (708.0655270091651, 263.1472040695353), (722.8961634797217, 270.80034882757894), (728.2438636119088, 270.86391665040594), (732.8865069126186, 262.13229555563095), (734.9628651080236, 261.5314805455938), (743.5079064112757, 269.6097211440626), (740.626661435203, 265.6040164357563), (756.0919992495631, 268.85730703537445), (760.3331171117211, 273.987114566834), (752.3195033731063, 267.630395860822), (762.1406114005654, 271.1688795410745), (768.4437436990225, 272.39100423802495), (769.6389614278111, 270.4043033219978), (786.1774454092401, 277.41811756327735), (792.5531569050898, 275.2998134104807), (788.4294583537404, 263.976887084107), (789.5846895631001, 273.34417608726073), (799.3040642927632, 272.98107600129674), (805.3373408398285, 271.2588716882814), (806.4432370898028, 270.665206353944), (813.921755369561, 274.6760175601746), (814.7106002905416, 274.9489436528596), (822.7174022115418, 275.84736676695366), (822.7715307130898, 276.91376426612067), (835.3152010060019, 266.61597337847604), (839.0251903726219, 277.75101004140157), (837.3187189218118, 272.2690890341042), (842.7875201664002, 278.7262473668461), (847.6349544981093, 268.3231832424176), (863.3217136149369, 277.02232939469394), (859.7262946299107, 276.7703367556846), (863.7880789134827, 274.60428393906585), (870.3671267554245, 276.2733166819531), (877.1124027696003, 265.3715964818472), (879.2350545295685, 280.27982157897986), (896.8258910971084, 273.5642812777275), (898.8418404955643, 266.02345907389315), (900.5798761970904, 273.36036361952847), (899.300040800521, 272.7243713886492), (917.3163420965044, 279.4676551787424), (921.2251894050358, 269.57946855637545), (914.8786340829956, 271.6213646027582), (929.7815950364192, 273.59603216135093), (935.8990995172119, 267.12332074367674), (936.770914819462, 268.86161663514883), (946.9510772677363, 268.5616084016788), (945.1090418541772, 271.63002043863486), (942.6135746981644, 279.0405218802981), (947.4031691602536, 267.64844980366576), (963.308020723541, 281.4072838592242), (958.526277391257, 267.4322053687551), (969.3113280992739, 280.4159148353725), (975.5933386199479, 271.54884438561976), (973.5351087221619, 280.477159190593), (992.5091920918983, 275.134531859925), (997.5962738460184, 274.16389099925846), (998.2185045124364, 274.0892680535451), (1000.3270383443931, 277.6329263772447), (998.242967122664, 273.5531216987363), (1007.2733979562806, 274.1074891524873), (1007.6674906228562, 272.7188177589547), (1016.050636562134, 269.4085065957457), (1019.318297961289, 284.20426178533575), (1024.7839983400452, 282.27450945417013), (1041.4561525467411, 274.48930825320787), (1032.8974521000166, 285.0532235514212), (1051.391282669397, 274.90552549133025), (1049.8697645586783, 277.4922759095924), (1053.2509316910198, 279.4522734710789), (1052.4291462714991, 285.1237425296172), (1066.8577125250429, 278.1120294093625), (1073.949309620567, 274.6044886492283), (1075.8234437321937, 283.0205704762114), (1079.430319815726, 276.53601428079725), (1090.8896131595122, 277.72749578809254), (1094.6113858780475, 276.6308238133251), (1092.7122113823248, 283.81246537338643), (1101.1044383069905, 271.29485215064744), (1100.3982100593685, 283.6342384084046), (1103.17977159758, 279.3013875999353), (1110.6335799905085, 276.9264175278924), (1124.300727191119, 273.7476083350614), (1121.0250558379662, 275.88672370900724), (1136.7781977290326, 273.94777811196263), (1133.0345416992195, 272.66019532390925), (1135.8288733295954, 286.0490023877359), (1145.7834905455275, 277.8800304939227), (1142.0127304070766, 282.8126173213673), (1150.060747167361, 287.7084882287928), (1158.9967507454728, 280.45919830340813), (1167.3888976122614, 274.9379811976714), (1177.2876731662539, 283.7513561928875), (1172.2685221274166, 278.9040120377208), (1182.8736182492496, 280.02716567750986), (1192.9272423641341, 285.80274365563747), (1196.1934987498205, 287.8499386559872), (1192.7912850084322, 287.529547577409), (1206.2510251715257, 284.5882565239575), (1210.8118796783185, 279.9138402489226), (1216.6532202520598, 276.39380107378867), (1219.9432418558035, 283.4282179356329), (1221.11315006636, 285.81077597021374), (1225.6503551627452, 278.6946843620867), (1236.8940345337091, 287.58710541234984), (1228.4538537412402, 283.920342409159), (1237.2036598007837, 279.63062606428446), (1246.498489884122, 280.4379130902805), (1245.3329414253124, 281.3118098936571), (1257.5265714248037, 283.8058866457418), (1263.5881925092442, 282.6220767968771), (1266.8598904640346, 285.14621747402725), (1276.7455211932797, 290.25488018060753), (1273.8391422339619, 285.439870633802), (1275.6736181570277, 285.1375300702421), (1284.2027740722367, 286.3601863989344), (1283.1242784124656, 276.5288355052286), (1292.5177446485825, 288.77064795998206), (1304.6926802901105, 281.43896552957716), (1309.057868537286, 288.89946980189853), (1305.3713265667188, 282.7528420745641), (1309.4653494633342, 280.84322328766433), (1322.9115501168376, 276.65302855386403), (1324.8654233610569, 291.11524960850863), (1322.1192315439378, 279.28119906645867), (1335.4087081297594, 281.67118812554446), (1337.0719393538936, 287.0839224725788), (1349.7228791617363, 289.27894206619345), (1344.7657802156589, 292.010851903016), (1362.9789806266958, 287.9043158064125), (1365.9896768742876, 284.37043215347114), (1372.1332682139991, 279.4534150799918), (1375.7265635344704, 281.9876869897247), (1371.3695760201686, 288.0756012824648), (1375.4939186809445, 288.57430510672964), (1380.9666508828682, 282.2553694995741), (1393.3197830470465, 278.42515166918975), (1394.8169398019932, 283.56381229701884), (1399.201394063257, 288.4848431829683), (1399.9447905367658, 290.91489498583076), (1412.672356789721, 290.9590207989381), (1421.356419152092, 280.74100684539275), (1424.1674702693583, 284.9364279120107), (1429.682595243649, 292.6247906281175), (1436.9760777419249, 283.1912917687539), (1433.683576693816, 290.51768763691723), (1436.0683610984079, 293.0306217767802), (1448.5553795135531, 282.92297954164394), (1446.1985620420835, 283.12630583151844), (1450.5202520711025, 290.4031533402797), (1462.6324935173911, 290.88736618830944), (1467.495958316436, 285.3245725146651), (1470.4196615644912, 282.70927041991484), (1481.323401999719, 289.51578162582916), (1474.125546028661, 284.38425821911903), (1477.359337062149, 287.33188332209596), (1495.6871346220776, 283.0084533239059), (1495.1967068366064, 286.6635997461865), (1494.282982591961, 290.9184415991366), (1498.502144661068, 289.729253689415), (1511.9546484153993, 284.91162433195103), (1512.2389915420836, 285.440015049375), (1517.5420372756412, 282.26062929527075), (1520.538978516502, 282.5392933333763), (1537.534718617824, 286.86628766611597), (1541.1790157598793, 283.22042060480385), (1544.9362054947, 289.5430833806965), (1544.5217005930315, 286.8878809321911), (1546.3449269904704, 295.01341303724485), (1554.0950025749494, 290.42936342993204), (1560.7834935378316, 291.58256803660987), (1560.4324260320532, 298.13414497465857), (1573.5680831395728, 287.8754887578058), (1580.2112327730063, 296.5453439728025), (1579.9807184968845, 296.4978355205458), (1588.7015904617056, 288.90533644821056), (1595.004199831141, 290.30364617957713), (1597.4577072571958, 298.842138679482), (1605.5290216544904, 290.8060208305865), (1606.217957717263, 298.679073807415), (1617.1247493110718, 287.66355422112275), (1615.274428374421, 291.826794738184), (1624.995711240025, 285.86270742184126), (1625.8300625069712, 294.2077965620291), (1635.7938155181246, 285.80215080016546), (1629.3824611853986, 293.83575974021926), (1639.4529851967375, 289.1994259157376), (1644.9171304552904, 287.50951965509694), (1651.9198923037907, 287.4369470569843), (1659.9851520495888, 291.8711234057065), (1656.5222000432416, 293.07889481132315), (1670.0845726834505, 290.24854973804145), (1671.6859459192694, 289.4573640368861), (1676.4955800036932, 298.36299349992686), (1681.880228863793, 298.2671644192603), (1690.1428702092433, 300.1661873054405), (1683.8456219262723, 298.84772300797295), (1690.9271590592998, 287.662736632823), (1696.2400535531317, 291.88496181030945), (1701.2826058990988, 298.8245602145944), (1704.736487075138, 298.26836857352964), (1721.0531893058453, 299.4216774561701), (1718.53678200437, 292.54395550999243), (1718.6633946885545, 289.7981210184704), (1730.5638565242414, 289.7495077872403), (1740.94042541695, 296.8618871897551), (1742.3248433604049, 292.9443844497491), (1745.416834279689, 295.52678160560856), (1747.2516925589669, 290.93170277293234), (1752.3705235851935, 299.923226694993), (1760.7469568668346, 299.03897827117413), (1761.6144278014483, 301.97364381509107), (1767.9651279305258, 302.7913027221788), (1781.555955824845, 293.7486432755097), (1775.5000635124534, 301.59974624308654), (1782.1689391025593, 290.3418087678853), (1783.652384810308, 300.21792792368984), (1798.7131852375373, 295.4314337053134), (1802.2228982750771, 289.7641118421078), (1805.3105204153078, 290.02724296684556), (1805.845877507224, 296.4778780267972), (1810.6111802027297, 297.7420281629228), (1822.1409852624865, 299.1687968534061), (1827.9047336886124, 299.2892194278894), (1831.5408160691823, 290.75071975313386), (1833.0152989585984, 303.9953009268141), (1836.3764654755892, 295.2842209080913), (1850.7326810430084, 304.5325490782085), (1842.4261841321693, 299.4364064365735), (1860.2057746757157, 301.81160557086326), (1853.4000174515243, 291.16866060095055), (1858.9376726537282, 303.0488109078991), (1871.2607687357843, 295.65438225902045), (1882.256519602529, 302.4389157332479), (1873.7838776904755, 300.2554259849878), (1890.6467814871096, 292.5088482516275), (1884.2715434390298, 298.91862896607904), (1895.1004363469867, 292.1821017569861), (1907.3055130276714, 298.65266239456304), (1911.0265565149812, 291.834461859917), (1906.0452170044769, 292.50877937261305), (1911.6141336497378, 307.13252747915385), (1920, 300)] 13 | 14 | maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0) 15 | ImageDraw.Draw(maskIm).polygon(polygon1, outline=1, fill=1) 16 | mask1 = numpy.array(maskIm) 17 | ImageDraw.Draw(maskIm).polygon(polygon2, outline=1, fill=1) 18 | mask = numpy.array(maskIm) 19 | 20 | print (numpy.allclose(mask, mask1)) 21 | 22 | # assemble new image (uint8: 0-255) 23 | newImArray = numpy.empty(imArray.shape,dtype='uint8') 24 | 25 | # colors (three first columns, RGB) 26 | newImArray[:,:,:3] = imArray[:,:,:3] 27 | 28 | # transparency (4th column) 29 | newImArray[:,:,3] = mask*255 30 | 31 | 32 | # back to Image from numpy 33 | newIm = Image.fromarray(newImArray, "RGBA") 34 | newIm.save("out.png") 35 | 36 | 37 | 38 | -------------------------------------------------------------------------------- /unit_tests/test_zoom.py: -------------------------------------------------------------------------------- 1 | from matplotlib import pyplot as plt 2 | from matplotlib.collections import PatchCollection 3 | from matplotlib.patches import Polygon 4 | 5 | import argparse 6 | import numpy as np 7 | 8 | zoom_scale, points, objects, prev = 1.2, [], [], None 9 | 10 | ap = argparse.ArgumentParser() 11 | ap.add_argument("-i", "--image", required=True, help="Path to the image") 12 | args = vars(ap.parse_args()) 13 | 14 | fig = plt.figure(figsize=(14, 14)) 15 | ax = plt.gca() 16 | image = plt.imread(args["image"]) 17 | 18 | def zoom(event): 19 | 20 | global zoom_scale, ax, fig 21 | 22 | cur_xlim = ax.get_xlim() 23 | cur_ylim = ax.get_ylim() 24 | 25 | xdata = event.xdata # get event x location 26 | ydata = event.ydata # get event y location 27 | 28 | if event.button == 'down': 29 | # deal with zoom in 30 | scale_factor = 1 / zoom_scale 31 | elif event.button == 'up': 32 | # deal with zoom out 33 | scale_factor = zoom_scale 34 | else: 35 | # deal with something that should never happen 36 | scale_factor = 1 37 | #print event.button 38 | 39 | new_width = (cur_xlim[1] - cur_xlim[0]) * scale_factor 40 | new_height = (cur_ylim[1] - cur_ylim[0]) * scale_factor 41 | 42 | relx = (cur_xlim[1] - xdata)/(cur_xlim[1] - cur_xlim[0]) 43 | rely = (cur_ylim[1] - ydata)/(cur_ylim[1] - cur_ylim[0]) 44 | 45 | ax.set_xlim([xdata - new_width * (1-relx), xdata + new_width * (relx)]) 46 | ax.set_ylim([ydata - new_height * (1-rely), ydata + new_height * (rely)]) 47 | ax.figure.canvas.draw() 48 | 49 | 50 | plt.imshow(image, aspect='auto') 51 | zoom_id = fig.canvas.mpl_connect('scroll_event', zoom) 52 | #click_id = fig.canvas.mpl_connect('button_press_event', onclick) 53 | 54 | #plt.subplots_adjust(bottom=0.2) 55 | plt.axis('off') 56 | plt.show() 57 | 58 | fig.canvas.mpl_disconnect(zoom_id) 59 | #fig.canvas.mpl_disconnect(click_id) 60 | --------------------------------------------------------------------------------