├── prepare_records ├── generate_pbtxt.py ├── label_map.pbtxt ├── xml_to_csv.py ├── partition_dataset.py └── generate_tfrecord.py ├── docs ├── _config.yml └── index.md ├── img └── PlasticNetProjectArchitectureDiagram.png ├── .gitignore ├── out └── label_map.pbtxt ├── download_training.py ├── override_pipeline.py ├── requirements.txt ├── prepare_training.py ├── setup.py ├── train_model.py ├── test_model ├── testTensorflow.py ├── testTensorWebcam.py └── testTensorVideo.py ├── PlasticNet ├── README.md ├── ModelZoo.md └── PlasticNet.py /prepare_records/generate_pbtxt.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/_config.yml: -------------------------------------------------------------------------------- 1 | theme: jekyll-theme-tactile -------------------------------------------------------------------------------- /img/PlasticNetProjectArchitectureDiagram.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/IBM/PlasticNet/HEAD/img/PlasticNetProjectArchitectureDiagram.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | models/ 2 | darknet/ 3 | test_model/images 4 | test_model/videos 5 | checkpoint/ 6 | detections/ 7 | *.config 8 | *.tar.gz 9 | prepare_records/images 10 | prepare_records/train.record 11 | prepare_records/test.record 12 | test_model/detections/* 13 | !test_model/detections/.gitkeep 14 | test_model/images/* 15 | !test_model/images/.gitkeep 16 | test_model/videos/* 17 | !test_model/videos/.gitkeep 18 | out/ 19 | !out/label_map.pbtxt 20 | .DS_Store 21 | __pycache__/ 22 | hi.py 23 | cmd.py.zip 24 | -------------------------------------------------------------------------------- /out/label_map.pbtxt: -------------------------------------------------------------------------------- 1 | item { 2 | name: "plastic-bag" 3 | id: 1 4 | } 5 | item { 6 | name: "plastic-bottle" 7 | id: 2 8 | } 9 | item { 10 | name: "potato-chip-bag" 11 | id: 3 12 | } 13 | item { 14 | name: "plastic-jug" 15 | id: 4 16 | } 17 | item { 18 | name: "styrafoam" 19 | id: 5 20 | } 21 | item { 22 | name: "metal-can" 23 | id: 6 24 | } 25 | item { 26 | name: "fishing-net" 27 | id: 7 28 | } 29 | item { 30 | name: "plastic-piece" 31 | id: 8 32 | } 33 | item { 34 | name: "face-mask" 35 | id: 9 36 | } -------------------------------------------------------------------------------- /prepare_records/label_map.pbtxt: -------------------------------------------------------------------------------- 1 | item { 2 | name: "plastic-bag" 3 | id: 1 4 | } 5 | item { 6 | name: "plastic-bottle" 7 | id: 2 8 | } 9 | item { 10 | name: "potato-chip-bag" 11 | id: 3 12 | } 13 | item { 14 | name: "plastic-jug" 15 | id: 4 16 | } 17 | item { 18 | name: "styrafoam" 19 | id: 5 20 | } 21 | item { 22 | name: "metal-can" 23 | id: 6 24 | } 25 | item { 26 | name: "fishing-net" 27 | id: 7 28 | } 29 | item { 30 | name: "plastic-piece" 31 | id: 8 32 | } 33 | item { 34 | name: "face-mask" 35 | id: 9 36 | } -------------------------------------------------------------------------------- /download_training.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | import tarfile 17 | import sys 18 | 19 | import six.moves.urllib as urllib 20 | 21 | def download_checkpoint(model): 22 | # For PlasticNet Model Zoo 23 | download_base = 'https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/' 24 | # Download the checkpoint 25 | opener = urllib.request.URLopener() 26 | opener.retrieve(download_base + model + '.tar.gz', model) 27 | 28 | tar = tarfile.open(model) 29 | tar.extractall('./out') # pipeline.config file will be included 30 | tar.close() 31 | 32 | os.remove(model) 33 | 34 | 35 | if __name__ == '__main__': 36 | print(sys.argv[1]) 37 | download_checkpoint(sys.argv[1]) -------------------------------------------------------------------------------- /override_pipeline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | from google.protobuf import text_format 18 | 19 | from object_detection.utils import config_util 20 | from object_detection.protos import preprocessor_pb2 21 | 22 | def override_pipeline(pipeline, override_dict, num_classes=0): 23 | configs = config_util.get_configs_from_pipeline_file(pipeline) 24 | 25 | meta_arch = configs["model"].WhichOneof("model") 26 | override_dict['model.{}.num_classes'.format(meta_arch)] = num_classes 27 | configs = config_util.merge_external_params_with_configs(configs, kwargs_dict=override_dict) 28 | pipeline_config = config_util.create_pipeline_proto_from_configs(configs) 29 | config_util.save_pipeline_config(pipeline_config, "./training") # replace with where pipeline file is to be stored 30 | -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | ## Welcome to GitHub Pages 2 | 3 | You can use the [editor on GitHub](https://github.com/IBM/PlasticNet/edit/main/docs/index.md) to maintain and preview the content for your website in Markdown files. 4 | 5 | Whenever you commit to this repository, GitHub Pages will run [Jekyll](https://jekyllrb.com/) to rebuild the pages in your site, from the content in your Markdown files. 6 | 7 | ### Markdown 8 | 9 | Markdown is a lightweight and easy-to-use syntax for styling your writing. It includes conventions for 10 | 11 | ```markdown 12 | Syntax highlighted code block 13 | 14 | # Header 1 15 | ## Header 2 16 | ### Header 3 17 | 18 | - Bulleted 19 | - List 20 | 21 | 1. Numbered 22 | 2. List 23 | 24 | **Bold** and _Italic_ and `Code` text 25 | 26 | [Link](url) and ![Image](src) 27 | ``` 28 | 29 | For more details see [GitHub Flavored Markdown](https://guides.github.com/features/mastering-markdown/). 30 | 31 | ### Jekyll Themes 32 | 33 | Your Pages site will use the layout and styles from the Jekyll theme you have selected in your [repository settings](https://github.com/IBM/PlasticNet/settings/pages). The name of this theme is saved in the Jekyll `_config.yml` configuration file. 34 | 35 | ### Support or Contact 36 | 37 | Having trouble with Pages? Check out our [documentation](https://docs.github.com/categories/github-pages-basics/) or [contact support](https://support.github.com/contact) and we’ll help you sort it out. 38 | -------------------------------------------------------------------------------- /prepare_records/xml_to_csv.py: -------------------------------------------------------------------------------- 1 | 2 | # Copyright 2021 IBM Corporation 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | import os 16 | import glob 17 | import pandas as pd 18 | import xml.etree.ElementTree as ET 19 | 20 | 21 | def xml_to_csv(path): 22 | xml_list = [] 23 | for xml_file in glob.glob(path + '/*.xml'): 24 | tree = ET.parse(xml_file) 25 | root = tree.getroot() 26 | for member in root.findall('object'): 27 | value = (os.path.basename(root.find('filename').text), 28 | int(root.find('size')[0].text), 29 | int(root.find('size')[1].text), 30 | member[0].text, 31 | int(member[4][0].text), 32 | int(member[4][1].text), 33 | int(member[4][2].text), 34 | int(member[4][3].text) 35 | ) 36 | xml_list.append(value) 37 | column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] 38 | xml_df = pd.DataFrame(xml_list, columns=column_name) 39 | return xml_df 40 | 41 | 42 | def main(): 43 | for folder in ['train', 'test']: 44 | image_path = os.path.join(os.getcwd(), ('images/' + folder)) 45 | xml_df = xml_to_csv(image_path) 46 | xml_df.to_csv(('images/'+folder+'_labels.csv'), index=None) 47 | print('Successfully converted xml to csv.') 48 | 49 | 50 | main() 51 | 52 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==0.12.0 2 | apache-beam==2.30.0 3 | astunparse==1.6.3 4 | avro-python3==1.9.2.1 5 | cachetools==4.2.2 6 | certifi==2021.5.30 7 | charset-normalizer==2.0.3 8 | crcmod==1.7 9 | cycler==0.10.0 10 | Cython==0.29.24 11 | dill==0.3.1.1 12 | dm-tree==0.1.6 13 | docopt==0.6.2 14 | fastavro==1.4.2 15 | flatbuffers==1.12 16 | future==0.18.2 17 | gast==0.4.0 18 | gin-config==0.4.0 19 | google-api-core==1.30.0 20 | google-api-python-client==2.11.0 21 | google-auth==1.33.1 22 | google-auth-httplib2==0.1.0 23 | google-auth-oauthlib==0.4.4 24 | google-cloud-bigquery==2.20.0 25 | google-cloud-core==1.7.1 26 | google-crc32c==1.1.2 27 | google-pasta==0.2.0 28 | google-resumable-media==1.3.1 29 | googleapis-common-protos==1.53.0 30 | grpcio==1.34.1 31 | h5py==3.1.0 32 | hdfs==2.6.0 33 | httplib2==0.19.1 34 | idna==3.2 35 | joblib==1.0.1 36 | kaggle==1.5.12 37 | keras-nightly==2.5.0.dev2021032900 38 | Keras-Preprocessing==1.1.2 39 | kiwisolver==1.3.1 40 | lvis==0.5.3 41 | Markdown==3.3.4 42 | matplotlib==3.4.2 43 | numpy==1.19.5 44 | oauth2client==4.1.3 45 | oauthlib==3.1.1 46 | opencv-python==4.5.2.54 47 | opencv-python-headless==4.5.2.54 48 | opt-einsum==3.3.0 49 | Pillow==8.3.1 50 | portalocker==2.0.0 51 | promise==2.3 52 | proto-plus==1.19.0 53 | protobuf==3.17.3 54 | psutil==5.8.0 55 | py-cpuinfo==8.0.0 56 | pyarrow==3.0.0 57 | pyasn1==0.4.8 58 | pyasn1-modules==0.2.8 59 | pycocotools==2.0.2 60 | pydot==1.4.2 61 | pymongo==3.11.4 62 | pyparsing==2.4.7 63 | python-dateutil==2.8.2 64 | python-slugify==5.0.2 65 | PyYAML==5.4.1 66 | requests==2.26.0 67 | requests-oauthlib==1.3.0 68 | rsa==4.7.2 69 | sacrebleu==1.5.1 70 | scikit-learn==0.24.2 71 | scipy==1.7.0 72 | sentencepiece==0.1.96 73 | seqeval==1.2.2 74 | six==1.15.0 75 | tensorboard==2.5.0 76 | tensorboard-data-server==0.6.1 77 | tensorboard-plugin-wit==1.8.0 78 | tensorflow==2.5.0 79 | tensorflow-addons==0.13.0 80 | tensorflow-datasets==4.3.0 81 | tensorflow-estimator==2.5.0 82 | tensorflow-metadata==1.1.0 83 | tensorflow-model-optimization==0.6.0 84 | termcolor==1.1.0 85 | text-unidecode==1.3 86 | tf-slim==1.1.0 87 | threadpoolctl==2.1.0 88 | tqdm==4.61.1 89 | typeguard==2.12.1 90 | typing-extensions==3.7.4.3 91 | uritemplate==3.0.1 92 | urllib3==1.26.6 93 | Werkzeug==2.0.1 94 | wrapt==1.12.1 95 | -------------------------------------------------------------------------------- /prepare_training.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | import json 17 | 18 | from download_training import download_checkpoint 19 | from override_pipeline import override_pipeline 20 | from absl import app, flags 21 | import subprocess 22 | 23 | flags.DEFINE_string('arch_type', None, 'Darknet YOLO (y) or Tensorflow (t)') 24 | flags.DEFINE_string('yolo_weights', None, 'name of weights file') 25 | flags.DEFINE_string('model_name', None, 'Title of model') 26 | flags.DEFINE_string('model_config_name', None, 'Model config name') 27 | flags.DEFINE_integer('num_classes', None, "Number of classes") 28 | 29 | FLAGS = flags.FLAGS 30 | 31 | def main(argv): 32 | """ 33 | Downloads and overrides the pipeline.config file distributed with any base PlasticNet models. 34 | Any image augmentation parameters will have to be changed manually within training/pipeline.config 35 | """ 36 | #select YOLOv4 or tensorflow 37 | archType = FLAGS.arch_type 38 | yoloWeights = FLAGS.yolo_weights 39 | if archType == 't': 40 | 41 | MODEL_CHECKPOINT = FLAGS.model_name 42 | NUM_CLASSES = FLAGS.num_classes 43 | DIR_PATH = os.path.abspath(os.getcwd()) 44 | train_record_path = DIR_PATH + '/prepare_records/train.record' 45 | val_record_path = DIR_PATH + '/prepare_records/test.record' 46 | checkpoint_path = DIR_PATH + '/checkpoint/' + MODEL_CHECKPOINT + '/checkpoint/ckpt-0' 47 | label_map_path = DIR_PATH + '/prepare_records/label_map.pbtxt' 48 | checkpoint_type = "detection" 49 | 50 | override_dict = { 51 | 'train_input_path': train_record_path, 52 | 'eval_input_path': val_record_path, 53 | 'train_config.fine_tune_checkpoint': checkpoint_path, 54 | 'label_map_path': label_map_path, 55 | 'train_config.fine_tune_checkpoint_type': checkpoint_type 56 | } 57 | download_checkpoint(MODEL_CHECKPOINT) 58 | override_pipeline("./out/" + MODEL_CHECKPOINT + "/exported_model/" + "pipeline.config", override_dict, NUM_CLASSES) 59 | elif archType == 'y': 60 | print('yolo') 61 | download_checkpoint(yoloWeights) 62 | if __name__ == '__main__': 63 | app.run(main) 64 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import os, sys, platform 15 | 16 | os_name = platform.system() 17 | 18 | # Install required dependencies 19 | os.system("pip install -r requirements.txt") 20 | 21 | # Clone the tensorflow/models respository 22 | current_dir = os.path.dirname(os.path.abspath(__file__)) 23 | 24 | # Clone the PlasticNet YOLOv4 Darknet Forked Repository 25 | if not os.path.isdir(os.path.join(current_dir, "darknet")): 26 | os.system("git clone https://github.com/mattokc35/darknet.git") 27 | 28 | # Set up the darknet directory by running make 29 | os.chdir('darknet') 30 | os.system('make') 31 | os.chdir('../') 32 | 33 | # Clone the TensorFlow Model Garden 34 | if not os.path.isdir(os.path.join(current_dir, "models")): 35 | os.system("git clone https://github.com/tensorflow/models.git") 36 | 37 | 38 | # Change to research directory and complete TensorFlow Setup 39 | os.chdir('models/research') 40 | print(str(os.getcwd())) 41 | os.system("protoc -I=./ --python_out=./ ./object_detection/protos/*.proto") 42 | sys.path.append(str(current_dir) + '/sli') 43 | sys.path.append(str(current_dir) + "/models") 44 | print(str(current_dir) + "/models") 45 | os.chdir("../../") 46 | 47 | 48 | if (os_name == 'Darwin'): #Mac 49 | # Convert command line client file into an executable 50 | current_path = os.getcwd() 51 | print(current_path) 52 | os.system("cp PlasticNet.py PlasticNetCopy.py") 53 | os.system("chmod +x PlasticNetCopy.py") 54 | os.system("mv PlasticNetCopy.py PlasticNet") 55 | 56 | #permanently add PlasticNet to your bin $PATH 57 | os.system('mkdir -p ~/bin') 58 | os.system('cp PlasticNet ~/bin') 59 | 60 | current_dir = os.path.dirname(os.path.abspath(__file__)) 61 | ### WORKING FOR MACOS 62 | file_read = open(os.path.expanduser("~") + '/.zshrc', 'r+') 63 | readfile = file_read.read() 64 | if not 'PATH=$PATH":$HOME/bin"' in readfile: 65 | file = open(os.path.expanduser("~") + '/.zshrc', 'a') 66 | file.write('\nexport PATH=$PATH\":$HOME/bin\"') 67 | file.write('\nexport PYTHONPATH=$PYTHONPATH:' + os.path.join(current_dir, "models")) 68 | file.write('\nexport PYTHONPATH=$PYTHONPATH:' + os.path.join(current_dir, "models/research")) 69 | file.write('\nexport PYTHONPATH=$PYTHONPATH:' + os.path.join(current_dir, "models/research/sli")) 70 | file.write('\nexport PLASTICNET_PATH=' + current_dir) 71 | file.close() 72 | 73 | print("********************************") 74 | print("SUCCESS!!!!!") 75 | print("Dependencies installed. Please restart your terminal session and type PlasticNet from any location to run") 76 | print("********************************") 77 | 78 | if (os_name != 'Darwin'): 79 | print("********************************") 80 | print("SUCCESS!!!!!") 81 | print("Dependencies installed. Please follow the documentation on what scripts you need to run to use the project.") 82 | print("********************************") -------------------------------------------------------------------------------- /prepare_records/partition_dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import glob, os, shutil 15 | 16 | ### NOTE: IF YOU HAVE MULTIPLE FILE TYPES, YOU WILL HAVE TO RUN THIS SCRIPT ONCE FOR EACH FILE TYPE AND CHANGE THE EXTENSION 17 | 18 | # Current directory 19 | current_dir = os.path.dirname(os.path.abspath(__file__)) 20 | # Percentage of images to be used for the test set 21 | percentage_test = 20 22 | counter=1 23 | current_dir = os.path.join(current_dir, "images/") 24 | if not os.path.isdir(os.path.join(current_dir, "train")): 25 | os.mkdir(os.path.join(current_dir, "train")) 26 | os.mkdir(os.path.join(current_dir, "test")) 27 | index_test = round(100 / percentage_test) 28 | 29 | # move over jpg 30 | for pathAndFilename in glob.iglob(os.path.join(current_dir, "*.jpg")): 31 | title, ext = os.path.splitext(os.path.basename(pathAndFilename)) 32 | if counter == index_test: 33 | counter = 1 34 | shutil.move(str(pathAndFilename), current_dir + "test/" + str(title) + str(ext)) 35 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "test/" + str(title) + ".xml") 36 | else: 37 | shutil.move(str(pathAndFilename), current_dir + "train/" + str(title) + str(ext)) 38 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "train/" + str(title) + ".xml") 39 | counter = counter + 1 40 | 41 | # move over png 42 | counter = 1 43 | for pathAndFilename in glob.iglob(os.path.join(current_dir, "*.png")): 44 | title, ext = os.path.splitext(os.path.basename(pathAndFilename)) 45 | if counter == index_test: 46 | counter = 1 47 | shutil.move(str(pathAndFilename), current_dir + "test/" + str(title) + str(ext)) 48 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "test/" + str(title) + ".xml") 49 | else: 50 | shutil.move(str(pathAndFilename), current_dir + "train/" + str(title) + str(ext)) 51 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "train/" + str(title) + ".xml") 52 | counter = counter + 1 53 | 54 | # move over jpeg 55 | counter = 1 56 | for pathAndFilename in glob.iglob(os.path.join(current_dir, "*.jpeg")): 57 | title, ext = os.path.splitext(os.path.basename(pathAndFilename)) 58 | if counter == index_test: 59 | counter = 1 60 | shutil.move(str(pathAndFilename), current_dir + "test/" + str(title) + str(ext)) 61 | shutil.move(str(pathAndFilename)[:-4] + "xml", current_dir + "test/" + str(title) + ".xml") 62 | else: 63 | shutil.move(str(pathAndFilename), current_dir + "train/" + str(title) + str(ext)) 64 | shutil.move(str(pathAndFilename)[:-4] + "xml", current_dir + "train/" + str(title) + ".xml") 65 | counter = counter + 1 66 | 67 | # move over JPG 68 | counter = 1 69 | for pathAndFilename in glob.iglob(os.path.join(current_dir, "*.JPG")): 70 | title, ext = os.path.splitext(os.path.basename(pathAndFilename)) 71 | if counter == index_test: 72 | counter = 1 73 | shutil.move(str(pathAndFilename), current_dir + "test/" + str(title) + str(ext)) 74 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "test/" + str(title) + ".xml") 75 | else: 76 | shutil.move(str(pathAndFilename), current_dir + "train/" + str(title) + str(ext)) 77 | shutil.move(str(pathAndFilename)[:-3] + "xml", current_dir + "train/" + str(title) + ".xml") 78 | counter = counter + 1 79 | -------------------------------------------------------------------------------- /prepare_records/generate_tfrecord.py: -------------------------------------------------------------------------------- 1 | 2 | # Copyright 2021 IBM Corporation 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | from __future__ import division 17 | from __future__ import print_function 18 | from __future__ import absolute_import 19 | 20 | import os 21 | import io 22 | import pandas as pd 23 | 24 | from tensorflow.python.framework.versions import VERSION 25 | if VERSION >= "2.0.0a0": 26 | import tensorflow.compat.v1 as tf 27 | else: 28 | import tensorflow as tf 29 | 30 | from PIL import Image 31 | from object_detection.utils import dataset_util 32 | from collections import namedtuple, OrderedDict 33 | 34 | flags = tf.app.flags 35 | flags.DEFINE_string('csv_input', '', 'Path to the CSV input') 36 | flags.DEFINE_string('output_path', '', 'Path to output TFRecord') 37 | flags.DEFINE_string('image_dir', '', 'Path to images') 38 | FLAGS = flags.FLAGS 39 | 40 | 41 | def class_text_to_int(row_label): 42 | if row_label == 'plastic-bag': 43 | return 1 44 | elif row_label == 'plastic-bottle': 45 | return 2 46 | elif row_label == 'potato-chip-bag': 47 | return 3 48 | elif row_label == 'plastic-jug': 49 | return 4 50 | elif row_label == 'styrafoam': 51 | return 5 52 | elif row_label == 'metal-can': 53 | return 6 54 | elif row_label == 'soda-can-metal': 55 | return 7 56 | elif row_label == 'fishing-net': 57 | return 8 58 | elif row_label == 'plastic-piece': 59 | return 9 60 | elif row_label == 'face-mask': 61 | return 10 62 | else: 63 | return None 64 | def split(df, group): 65 | data = namedtuple('data', ['filename', 'object']) 66 | gb = df.groupby(group) 67 | return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)] 68 | 69 | 70 | def create_tf_example(group, path): 71 | with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid: 72 | encoded_jpg = fid.read() 73 | encoded_jpg_io = io.BytesIO(encoded_jpg) 74 | image = Image.open(encoded_jpg_io) 75 | width, height = image.size 76 | 77 | filename = group.filename.encode('utf8') 78 | image_format = b'jpg' 79 | xmins = [] 80 | xmaxs = [] 81 | ymins = [] 82 | ymaxs = [] 83 | classes_text = [] 84 | classes = [] 85 | 86 | for index, row in group.object.iterrows(): 87 | xmins.append(row['xmin'] / width) 88 | xmaxs.append(row['xmax'] / width) 89 | ymins.append(row['ymin'] / height) 90 | ymaxs.append(row['ymax'] / height) 91 | classes_text.append(row['class'].encode('utf8')) 92 | classes.append(class_text_to_int(row['class'])) 93 | 94 | tf_example = tf.train.Example(features=tf.train.Features(feature={ 95 | 'image/height': dataset_util.int64_feature(height), 96 | 'image/width': dataset_util.int64_feature(width), 97 | 'image/filename': dataset_util.bytes_feature(filename), 98 | 'image/source_id': dataset_util.bytes_feature(filename), 99 | 'image/encoded': dataset_util.bytes_feature(encoded_jpg), 100 | 'image/format': dataset_util.bytes_feature(image_format), 101 | 'image/object/bbox/xmin': dataset_util.float_list_feature(xmins), 102 | 'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs), 103 | 'image/object/bbox/ymin': dataset_util.float_list_feature(ymins), 104 | 'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs), 105 | 'image/object/class/text': dataset_util.bytes_list_feature(classes_text), 106 | 'image/object/class/label': dataset_util.int64_list_feature(classes), 107 | })) 108 | return tf_example 109 | 110 | 111 | def main(_): 112 | writer = tf.io.TFRecordWriter(FLAGS.output_path) 113 | path = os.path.join(FLAGS.image_dir) 114 | examples = pd.read_csv(FLAGS.csv_input) 115 | grouped = split(examples, 'filename') 116 | for group in grouped: 117 | tf_example = create_tf_example(group, path) 118 | writer.write(tf_example.SerializeToString()) 119 | 120 | writer.close() 121 | output_path = os.path.join(os.getcwd(), FLAGS.output_path) 122 | print('Successfully created the TFRecords: {}'.format(output_path)) 123 | 124 | 125 | if __name__ == '__main__': 126 | tf.app.run() 127 | -------------------------------------------------------------------------------- /train_model.py: -------------------------------------------------------------------------------- 1 | # train_model script from TensorFlow/models/research/object_detection: 2 | # https://github.com/tensorflow/models 3 | # Modified by IBM 4 | 5 | r"""Creates and runs TF2 object detection models. 6 | For local training/evaluation run: 7 | PIPELINE_CONFIG_PATH=path/to/pipeline.config 8 | MODEL_DIR=/tmp/model_outputs 9 | NUM_TRAIN_STEPS=10000 10 | SAMPLE_1_OF_N_EVAL_EXAMPLES=1 11 | python train_model.py -- \ 12 | --model_dir=$MODEL_DIR --num_train_steps=$NUM_TRAIN_STEPS \ 13 | --sample_1_of_n_eval_examples=$SAMPLE_1_OF_N_EVAL_EXAMPLES \ 14 | --pipeline_config_path=$PIPELINE_CONFIG_PATH \ 15 | --alsologtostderr 16 | """ 17 | from absl import flags 18 | import tensorflow.compat.v2 as tf 19 | from object_detection import model_lib_v2 20 | 21 | flags.DEFINE_string('pipeline_config_path', None, 'Path to pipeline config ' 22 | 'file.') 23 | flags.DEFINE_integer('num_train_steps', None, 'Number of train steps.') 24 | flags.DEFINE_bool('eval_on_train_data', False, 'Enable evaluating on train ' 25 | 'data (only supported in distributed training).') 26 | flags.DEFINE_integer('sample_1_of_n_eval_examples', None, 'Will sample one of ' 27 | 'every n eval input examples, where n is provided.') 28 | flags.DEFINE_integer('sample_1_of_n_eval_on_train_examples', 5, 'Will sample ' 29 | 'one of every n train input examples for evaluation, ' 30 | 'where n is provided. This is only used if ' 31 | '`eval_training_data` is True.') 32 | flags.DEFINE_string( 33 | 'model_dir', None, 'Path to output model directory ' 34 | 'where event and checkpoint files will be written.') 35 | flags.DEFINE_string( 36 | 'checkpoint_dir', None, 'Path to directory holding a checkpoint. If ' 37 | '`checkpoint_dir` is provided, this binary operates in eval-only mode, ' 38 | 'writing resulting metrics to `model_dir`.') 39 | 40 | flags.DEFINE_integer('eval_timeout', 3600, 'Number of seconds to wait for an' 41 | 'evaluation checkpoint before exiting.') 42 | 43 | flags.DEFINE_bool('use_tpu', False, 'Whether the job is executing on a TPU.') 44 | flags.DEFINE_string( 45 | 'tpu_name', 46 | default=None, 47 | help='Name of the Cloud TPU for Cluster Resolvers.') 48 | flags.DEFINE_integer( 49 | 'num_workers', 1, 'When num_workers > 1, training uses ' 50 | 'MultiWorkerMirroredStrategy. When num_workers = 1 it uses ' 51 | 'MirroredStrategy.') 52 | flags.DEFINE_integer( 53 | 'checkpoint_every_n', 1000, 'Integer defining how often we checkpoint.') 54 | flags.DEFINE_boolean('record_summaries', True, 55 | ('Whether or not to record summaries defined by the model' 56 | ' or the training pipeline. This does not impact the' 57 | ' summaries of the loss values which are always' 58 | ' recorded.')) 59 | 60 | FLAGS = flags.FLAGS 61 | 62 | 63 | def main(unused_argv): 64 | flags.mark_flag_as_required('model_dir') 65 | flags.mark_flag_as_required('pipeline_config_path') 66 | tf.config.set_soft_device_placement(True) 67 | gpus = tf.config.experimental.list_physical_devices('GPU') 68 | for gpu in gpus: 69 | tf.config.experimental.set_memory_growth(gpu, True) 70 | 71 | if FLAGS.checkpoint_dir: 72 | model_lib_v2.eval_continuously( 73 | pipeline_config_path=FLAGS.pipeline_config_path, 74 | model_dir=FLAGS.model_dir, 75 | train_steps=FLAGS.num_train_steps, 76 | sample_1_of_n_eval_examples=FLAGS.sample_1_of_n_eval_examples, 77 | sample_1_of_n_eval_on_train_examples=( 78 | FLAGS.sample_1_of_n_eval_on_train_examples), 79 | checkpoint_dir=FLAGS.checkpoint_dir, 80 | wait_interval=300, timeout=FLAGS.eval_timeout) 81 | else: 82 | if FLAGS.use_tpu: 83 | # TPU is automatically inferred if tpu_name is None and 84 | # we are running under cloud ai-platform. 85 | resolver = tf.distribute.cluster_resolver.TPUClusterResolver( 86 | FLAGS.tpu_name) 87 | tf.config.experimental_connect_to_cluster(resolver) 88 | tf.tpu.experimental.initialize_tpu_system(resolver) 89 | strategy = tf.distribute.experimental.TPUStrategy(resolver) 90 | elif FLAGS.num_workers > 1: 91 | strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() 92 | else: 93 | strategy = tf.compat.v2.distribute.MirroredStrategy() 94 | 95 | with strategy.scope(): 96 | model_lib_v2.train_loop( 97 | pipeline_config_path=FLAGS.pipeline_config_path, 98 | model_dir=FLAGS.model_dir, 99 | train_steps=FLAGS.num_train_steps, 100 | use_tpu=FLAGS.use_tpu, 101 | checkpoint_every_n=FLAGS.checkpoint_every_n, 102 | record_summaries=FLAGS.record_summaries) 103 | 104 | if __name__ == '__main__': 105 | tf.compat.v1.app.run() -------------------------------------------------------------------------------- /test_model/testTensorflow.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | """ 17 | Object Detection (On Image) From PlasticNet Tensorflow Saved Model 18 | ===================================== 19 | """ 20 | import os 21 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # Suppress TensorFlow logging (1) 22 | import pathlib 23 | import tensorflow as tf 24 | import cv2 25 | from matplotlib import pyplot as plt 26 | from matplotlib.pyplot import figure 27 | from absl import app, flags 28 | import numpy as np 29 | from PIL import Image 30 | import matplotlib.pyplot as plt 31 | import time 32 | import warnings 33 | warnings.filterwarnings('ignore') # Suppress Matplotlib warnings 34 | 35 | # Flag provides name of trained model we want to test 36 | flags.DEFINE_string('trained_model', None, 'Trained Model') 37 | # Flag provides the minimum confidence threshold 38 | flags.DEFINE_string('score', None, 'Score') 39 | 40 | FLAGS = flags.FLAGS 41 | 42 | 43 | def load_image_into_numpy_array(path): 44 | """Load an image from file into a numpy array. 45 | Puts image into numpy array to feed into tensorflow graph. 46 | Note that by convention we put it into a numpy array with shape 47 | (height, width, channels), where channels=3 for RGB. 48 | Args: 49 | path: the file path to the image 50 | Returns: 51 | uint8 numpy array with shape (img_height, img_width, 3) 52 | """ 53 | return np.array(Image.open(path)) 54 | 55 | def main(argv): 56 | 57 | #from google.colab.patches import cv2_imshow 58 | # Enable GPU dynamic memory allocation 59 | gpus = tf.config.experimental.list_physical_devices('GPU') 60 | for gpu in gpus: 61 | tf.config.experimental.set_memory_growth(gpu, True) 62 | 63 | # Path to image directory is inside PlasticNet/test_model/images (drop all your images here you want to test) 64 | LIST_IMAGE_PATHS = [] 65 | for root, dirs, files in os.walk(os.path.abspath("images/")): 66 | for file in files: 67 | LIST_IMAGE_PATHS.append(os.path.join(root, file)) 68 | 69 | # Provide path to model directory 70 | trained_model = FLAGS.trained_model 71 | # minimum confidence threshold 72 | score = FLAGS.score 73 | 74 | #Provide path to model directory 75 | PATH_TO_MODEL_DIR = '../out/' + str(trained_model) + '/exported_model/saved_model' 76 | 77 | #Provide the path to the label map file (label_map.pbtxt), change if you want your own label map 78 | PATH_TO_LABELS = '../out/label_map.pbtxt' 79 | 80 | # LOAD THE MODEL 81 | from object_detection.utils import label_map_util 82 | from object_detection.utils import visualization_utils as viz_utils 83 | PATH_TO_SAVED_MODEL = PATH_TO_MODEL_DIR 84 | print('Loading model...', end='') 85 | start_time = time.time() 86 | 87 | # Load saved model and build detection function 88 | detect_fn = tf.saved_model.load(PATH_TO_SAVED_MODEL) 89 | end_time = time.time() 90 | elapsed_time = end_time - start_time 91 | print('Done! Took {} seconds'.format(elapsed_time)) 92 | 93 | # List of the strings that is used to add correct label for each box. 94 | category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True) 95 | 96 | 97 | # Run detection test on each image in the images folder, and draw boxes around each object detected. 98 | print('Running inference for images') 99 | counter = 0 100 | for image_path in LIST_IMAGE_PATHS: 101 | print(image_path) 102 | image = cv2.imread(image_path) 103 | image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 104 | image_expanded = np.expand_dims(image_rgb, axis=0) 105 | # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. 106 | input_tensor = tf.convert_to_tensor(image) 107 | # The model expects a batch of images, so add an axis with `tf.newaxis`. 108 | input_tensor = input_tensor[tf.newaxis, ...] 109 | # input_tensor = np.expand_dims(image_np, 0) 110 | detections = detect_fn(input_tensor) 111 | # All outputs are batches tensors. 112 | # Convert to numpy arrays, and take index [0] to remove the batch dimension. 113 | # We're only interested in the first num_detections. 114 | num_detections = int(detections.pop('num_detections')) 115 | detections = {key: value[0, :num_detections].numpy() 116 | for key, value in detections.items()} 117 | detections['num_detections'] = num_detections 118 | # detection_classes should be ints. 119 | detections['detection_classes'] = detections['detection_classes'].astype(np.int64) 120 | image_with_detections = image.copy() 121 | # SET MIN_SCORE_THRESH BASED ON YOU MINIMUM THRESHOLD FOR DETECTIONS 122 | viz_utils.visualize_boxes_and_labels_on_image_array( 123 | image_with_detections, 124 | detections['detection_boxes'], 125 | detections['detection_classes'], 126 | detections['detection_scores'], 127 | category_index, 128 | use_normalized_coordinates=True, 129 | max_boxes_to_draw=200, 130 | min_score_thresh=float(score), 131 | agnostic_mode=False) 132 | 133 | # DISPLAYS OUTPUT IMAGE 134 | plt.figure(figsize=(35,17)) 135 | #figure(figsize=(8, 6), dpi = 100) 136 | plt.imshow(cv2.cvtColor(image_with_detections, cv2.COLOR_BGR2RGB)) 137 | plt.savefig('detections/' + str(counter) + "_result.jpg") 138 | print(str(counter)) 139 | print('Done') 140 | counter += 1 141 | 142 | if __name__ == '__main__': 143 | app.run(main) 144 | -------------------------------------------------------------------------------- /test_model/testTensorWebcam.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | """ 17 | Object Detection (On Webcam) From PlasticNet Tensorflow Saved Model 18 | ===================================== 19 | """ 20 | 21 | import numpy as np 22 | import os 23 | import six.moves.urllib as urllib 24 | import sys 25 | import tensorflow as tf 26 | import pathlib 27 | import cv2 28 | from collections import defaultdict 29 | from io import StringIO 30 | from matplotlib import pyplot as plt 31 | from PIL import Image 32 | from IPython.display import display 33 | from object_detection.utils import ops as utils_ops, label_map_util, visualization_utils as vis_util 34 | from absl import app, flags 35 | 36 | # Flag provides name of trained model we want to test 37 | flags.DEFINE_string('trained_model', None, 'Trained Model') 38 | # Flag provides the minimum confidence threshold 39 | flags.DEFINE_string('score', None, 'Score') 40 | 41 | FLAGS = flags.FLAGS 42 | 43 | def load_model(PATH_TO_MODEL_DIR): 44 | """ 45 | Function will load model that is saved in the PATH_TO_MODEL_DIR 46 | """ 47 | model_dir = pathlib.Path(PATH_TO_MODEL_DIR) 48 | model = tf.saved_model.load(str(model_dir)) 49 | return model 50 | 51 | def run_inference_for_single_image(model, image): 52 | """ 53 | Function tests model for single image, and draws boxes of detections on the frame 54 | """ 55 | image = np.asarray(image) 56 | # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. 57 | input_tensor = tf.convert_to_tensor(image) 58 | # The model expects a batch of images, so add an axis with `tf.newaxis`. 59 | input_tensor = input_tensor[tf.newaxis,...] 60 | 61 | # Run inference 62 | model_fn = model.signatures['serving_default'] 63 | output_dict = model_fn(input_tensor) 64 | 65 | # All outputs are batches tensors. 66 | # Convert to numpy arrays, and take index [0] to remove the batch dimension. 67 | # We're only interested in the first num_detections. 68 | num_detections = int(output_dict.pop('num_detections')) 69 | output_dict = {key:value[0, :num_detections].numpy() 70 | for key,value in output_dict.items()} 71 | output_dict['num_detections'] = num_detections 72 | 73 | # detection_classes should be ints. 74 | output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64) 75 | 76 | # Handle models with masks: 77 | if 'detection_masks' in output_dict: 78 | # Reframe the the bbox mask to the image size. 79 | detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks( 80 | output_dict['detection_masks'], output_dict['detection_boxes'], 81 | image.shape[0], image.shape[1]) 82 | detection_masks_reframed = tf.cast(detection_masks_reframed > 0.5, 83 | tf.uint8) 84 | output_dict['detection_masks_reframed'] = detection_masks_reframed.numpy() 85 | 86 | return output_dict 87 | 88 | def show_inference(model, image_np, score, category_index): 89 | """ 90 | Function outputs the visualized results of a detection on a frame. 91 | """ 92 | # Run single detection 93 | output_dict = run_inference_for_single_image(model, image_np) 94 | # Visualization of the results of a detection. 95 | final_img=vis_util.visualize_boxes_and_labels_on_image_array( 96 | image_np, 97 | output_dict['detection_boxes'], 98 | output_dict['detection_classes'], 99 | output_dict['detection_scores'], 100 | category_index, 101 | instance_masks=output_dict.get('detection_masks_reframed', None), 102 | use_normalized_coordinates=True, 103 | min_score_thresh=float(score), #minimum confidence threshold 104 | line_thickness=8) 105 | return(final_img) 106 | 107 | 108 | 109 | 110 | def main(argv): 111 | 112 | # provide path to model directory 113 | trained_model = FLAGS.trained_model 114 | 115 | # Provide minimum confidence threshold 116 | score = FLAGS.score 117 | PATH_TO_MODEL_DIR = '../out/' + str(trained_model) + '/exported_model/saved_model' 118 | 119 | #Provide the path to the label map file (label_map.pbtxt), change if you want your own label map 120 | PATH_TO_LABELS = '../out/label_map.pbtxt' 121 | # List of the strings that is used to add correct label for each box. 122 | category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True) 123 | 124 | 125 | # Enable GPU dynamic memory allocation 126 | gpus = tf.config.experimental.list_physical_devices('GPU') 127 | for gpu in gpus: 128 | tf.config.experimental.set_memory_growth(gpu, True) 129 | 130 | # Load model 131 | detection_model = load_model(PATH_TO_MODEL_DIR) 132 | detection_model.signatures['serving_default'].output_dtypes 133 | detection_model.signatures['serving_default'].output_shapes 134 | 135 | # Start video capture on the webcam 136 | cap = cv2.VideoCapture(0) 137 | 138 | # Set size for video capture 139 | frame_width = int(cap.get(3)) 140 | frame_height = int(cap.get(4)) 141 | 142 | size = (frame_width, frame_height) 143 | 144 | # save webcam result as mp4 video file 145 | result = cv2.VideoWriter('detections/webcam_result.mp4', cv2.VideoWriter_fourcc(*'MP4V'),20.0, size) 146 | 147 | while 1: 148 | # Keep looping until user hits 'q'. Will loop through frames of live video in webcam and run detection tests on each frame. 149 | _,img = cap.read() 150 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 151 | final_img = show_inference(detection_model, img, score, category_index) 152 | 153 | final_img = cv2.cvtColor(final_img, cv2.COLOR_RGB2BGR) 154 | cv2.imshow('img', final_img) 155 | result.write(final_img) 156 | 157 | 158 | #break loop on q key 159 | if cv2.waitKey(1) == ord('q'): 160 | break 161 | cap.release() 162 | cv2.destroyAllWindows() 163 | 164 | 165 | if __name__ == '__main__': 166 | app.run(main) 167 | -------------------------------------------------------------------------------- /test_model/testTensorVideo.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 IBM Corporation 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | """ 17 | Object Detection (On Video) From PlasticNet Tensorflow Saved Model 18 | ===================================== 19 | """ 20 | import numpy as np 21 | import os 22 | import six.moves.urllib as urllib 23 | import sys 24 | import tensorflow as tf 25 | import pathlib 26 | import cv2 27 | from collections import defaultdict 28 | from io import StringIO 29 | from matplotlib import pyplot as plt 30 | from PIL import Image 31 | from IPython.display import display 32 | from object_detection.utils import ops as utils_ops 33 | from object_detection.utils import label_map_util 34 | from object_detection.utils import visualization_utils as vis_util 35 | from absl import app, flags 36 | 37 | # Flag provides name of trained model we want to test 38 | flags.DEFINE_string('trained_model', None, 'Trained Model') 39 | # Flag provides name of video we are testing on 40 | flags.DEFINE_string('video_name', None, 'Video name') 41 | # Flag provides minimum confidence threshold 42 | flags.DEFINE_string('score', None, 'Score') 43 | 44 | FLAGS = flags.FLAGS 45 | 46 | def load_model(PATH_TO_MODEL_DIR): 47 | """ 48 | Function will load model that is saved in the PATH_TO_MODEL_DIR 49 | """ 50 | model_dir = pathlib.Path(PATH_TO_MODEL_DIR) 51 | model = tf.saved_model.load(str(model_dir)) 52 | return model 53 | 54 | def run_inference_for_single_image(model, image): 55 | """ 56 | Function tests model for single image, and draws boxes of detections on the frame 57 | """ 58 | image = np.asarray(image) 59 | # The input needs to be a tensor, convert it using `tf.convert_to_tensor`. 60 | input_tensor = tf.convert_to_tensor(image) 61 | # The model expects a batch of images, so add an axis with `tf.newaxis`. 62 | input_tensor = input_tensor[tf.newaxis,...] 63 | 64 | # Run inference 65 | model_fn = model.signatures['serving_default'] 66 | output_dict = model_fn(input_tensor) 67 | 68 | # All outputs are batches tensors. 69 | # Convert to numpy arrays, and take index [0] to remove the batch dimension. 70 | # We're only interested in the first num_detections. 71 | num_detections = int(output_dict.pop('num_detections')) 72 | output_dict = {key:value[0, :num_detections].numpy() 73 | for key,value in output_dict.items()} 74 | output_dict['num_detections'] = num_detections 75 | 76 | # detection_classes should be ints. 77 | output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64) 78 | 79 | # Handle models with masks: 80 | if 'detection_masks' in output_dict: 81 | # Reframe the the bbox mask to the image size. 82 | detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks( 83 | output_dict['detection_masks'], output_dict['detection_boxes'], 84 | image.shape[0], image.shape[1]) 85 | detection_masks_reframed = tf.cast(detection_masks_reframed > 0.5, 86 | tf.uint8) 87 | output_dict['detection_masks_reframed'] = detection_masks_reframed.numpy() 88 | 89 | return output_dict 90 | 91 | def show_inference(model, image_np, score, category_index): 92 | """ 93 | Function outputs the visualized results of a detection on a frame. 94 | """ 95 | # Run detection on single frame 96 | output_dict = run_inference_for_single_image(model, image_np) 97 | # Visualization of the results of a detection. 98 | final_img=vis_util.visualize_boxes_and_labels_on_image_array( 99 | image_np, 100 | output_dict['detection_boxes'], 101 | output_dict['detection_classes'], 102 | output_dict['detection_scores'], 103 | category_index, 104 | instance_masks=output_dict.get('detection_masks_reframed', None), 105 | use_normalized_coordinates=True, 106 | min_score_thresh=float(score), 107 | line_thickness=8) 108 | return(final_img) 109 | 110 | def main(argv): 111 | 112 | # PROVIDE PATH TO MODEL DIRECTORY 113 | trained_model = FLAGS.trained_model 114 | PATH_TO_MODEL_DIR = '../out/' + str(trained_model) + '/exported_model/saved_model' 115 | 116 | 117 | #Provide the path to the label map file (label_map.pbtxt), change if you want your own label map 118 | PATH_TO_LABELS = '../out/label_map.pbtxt' 119 | # List of the strings that is used to add correct label for each box. 120 | category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True) 121 | 122 | # Enable GPU dynamic memory allocation 123 | gpus = tf.config.experimental.list_physical_devices('GPU') 124 | for gpu in gpus: 125 | tf.config.experimental.set_memory_growth(gpu, True) 126 | 127 | # Load model 128 | detection_model = load_model(PATH_TO_MODEL_DIR) 129 | detection_model.signatures['serving_default'].output_dtypes 130 | detection_model.signatures['serving_default'].output_shapes 131 | # Provide video name 132 | video_name = FLAGS.video_name 133 | # Provide minimum confidence threshold 134 | score = FLAGS.score 135 | # Start video capture with given video file 136 | cap = cv2.VideoCapture(r'videos/' + str(video_name)) 137 | 138 | # Set video size 139 | frame_width = int(cap.get(3)) 140 | frame_height = int(cap.get(4)) 141 | 142 | size = (frame_width, frame_height) 143 | # Split the video title from the extension type 144 | video_title, video_ext = os.path.splitext(str(video_name)) 145 | # Save video file renamed to be a MP4 file 146 | result = cv2.VideoWriter('detections/result_' + video_title + '.mp4', cv2.VideoWriter_fourcc(*'MP4V'),20.0, size) 147 | 148 | while 1: 149 | # Will loop through frames of given video and run detection tests on each frame, stops if user presses 'q'. 150 | _,img = cap.read() 151 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 152 | final_img = show_inference(detection_model, img, score, category_index) 153 | 154 | final_img = cv2.cvtColor(final_img, cv2.COLOR_RGB2BGR) 155 | cv2.imshow('img', final_img) 156 | result.write(final_img) 157 | 158 | 159 | #break loop on q key 160 | if cv2.waitKey(1) == ord('q'): 161 | break 162 | cap.release() 163 | cv2.destroyAllWindows() 164 | 165 | 166 | if __name__ == '__main__': 167 | app.run(main) 168 | -------------------------------------------------------------------------------- /PlasticNet: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | from cmd import Cmd 4 | import sys 5 | import os 6 | import subprocess 7 | 8 | global currentModel 9 | global currentPath 10 | class MyPrompt(Cmd): 11 | 12 | print(os.popen("mdfind kind:folder “PlasticNet”").read()) 13 | def __init__(self): 14 | super(MyPrompt, self).__init__() 15 | self.currentModel = "default" 16 | proc = subprocess.Popen(["mdfind kind:folder 'PlasticNet'", "kind:folder 'PlasticNet'"], stdout=subprocess.PIPE, shell=True) 17 | (out, err) = proc.communicate() 18 | final_out = str(out).split("\\") 19 | self.currentPath = final_out[0][2:] 20 | 21 | def do_set_model(self, args): 22 | print("SETTING MODEL") 23 | args = args.split() 24 | if (len(args) == 1): 25 | self.currentModel = args[0] 26 | else: 27 | print("ERROR: Invalid syntax. Use help set_model for help.") 28 | return 29 | print('current model set to: ' + str(self.currentModel)) 30 | def help_set_model(self): 31 | print('--sets current model global variable that can be used for preparing, training, exporting, and testing') 32 | print('syntax: set_model [model name]') 33 | def do_prepare_training(self, args): 34 | args = args.split() 35 | model = "" 36 | if (len(args) == 1): 37 | model = args[0] 38 | print("PREPARING TRAINING FOR MODEL: " + str(model)) 39 | elif (len(args) == 0): 40 | model = "" 41 | print("PREPARING TRAINING FOR MODEL: " + str(self.currentModel)) 42 | else: 43 | print("ERROR: Invalid syntax. Use help prepare_training for help.") 44 | return 45 | return_code = os.system("python " + self.currentPath + "/" + "prepare_training.py --model_name " + args[0]) 46 | if return_code == 0: # successful 47 | self.currentModel = args[0] 48 | print(self.currentModel) 49 | else: 50 | print("ERROR: Prepare training has failed. Check the name of the model you've specified.") 51 | def help_prepare_training(self): 52 | print("syntax: prepare_training [model_name]") 53 | print("generates all the necessary files for training to begin. After this is complete, run train_model") 54 | 55 | def do_train_model(self, args): 56 | print("TRAINING MODEL: " + str(self.currentModel)) 57 | os.system("python " + self.currentPath + "/" + "train_model.py --pipeline_config_path=training/pipeline.config/ " + 58 | "--model_dir=out/" + str(self.currentModel)) 59 | def help_train_model(self): 60 | print ("syntax: train_model (no arguments needed, will use global variable \'current_model\'") 61 | print ("trains model using pre-trained base model specified by user") 62 | 63 | 64 | def do_export_model(self, args): 65 | args = args.split() 66 | if (len(args) == 1): 67 | model = args[0] 68 | else: 69 | print("ERROR: Invalid syntax. Use help export_model for help.") 70 | return 71 | print("EXPORTING MODEL: " + str(model)) 72 | os.chdir(self.currentPath) 73 | os.chdir('models/research/object_detection') 74 | os.system("python exporter_main_v2.py --input_type image_tensor --pipeline_config_path " + 75 | "../../../training/pipeline.config --trained_checkpoint_dir ../../../out/" + model + " " + 76 | "--output_directory ../../../out/" + model + "/exported_model") 77 | os.chdir('../../../') 78 | def help_export_model(self): 79 | print("syntax: export_model [model_name]") 80 | print("Exports a model after training. Specify the same model name as you used for training") 81 | def do_test_model(self, args): 82 | args = args.split() 83 | if (len(args) >= 1): 84 | mode = args[0] 85 | if(mode == 'w'): 86 | score = args[1] 87 | print('Testing on WEBCAM: ' + str(self.currentModel)) 88 | current_path = self.currentPath 89 | os.chdir(current_path + '/test_model') 90 | current_path = os.getcwd() 91 | os.system("python " + self.currentPath + "/test_model/" + "testTensorWebcam.py --trained_model " + str(self.currentModel) + " --score " + str(score)) 92 | os.chdir('../') 93 | current_path = os.getcwd() 94 | print(current_path) 95 | elif(mode == 'v'): 96 | if (len(args) == 3): 97 | videoName = args[1] 98 | score = args[2] 99 | print('Testing on VIDEO: ' + str(self.currentModel)) 100 | current_path = self.currentPath 101 | os.chdir(current_path + '/test_model') 102 | current_path = os.getcwd() 103 | os.system("python " + self.currentPath + "/test_model/" + "testTensorVideo.py --video_name " + str(videoName) + " --trained_model " + str(self.currentModel) + " --score " + str(score)) 104 | os.chdir('../') 105 | current_path = os.getcwd() 106 | print(current_path) 107 | else: 108 | print("ERROR: Invalid syntax. Use help test_model for help.") 109 | return 110 | elif(mode == 'i'): 111 | score = args[1] 112 | print('Testing on IMAGES: ' + str(self.currentModel)) 113 | current_path = self.currentPath 114 | os.chdir(current_path + '/test_model') 115 | current_path = os.getcwd() 116 | os.system("python " + self.currentPath + "/test_model/" + "testTensorflow.py --trained_model" + str(self.currentModel) + " --score " + str(score)) 117 | os.chdir('../') 118 | current_path = os.getcwd() 119 | print(current_path) 120 | else: 121 | print("ERROR: Invalid mode. Use help test_model for help.") 122 | else: 123 | print("ERROR: Invalid syntax. Use help test_model for help.") 124 | return 125 | def help_test_model(self): 126 | print ("syntax: test_model [w (webcam) | v (video) | i (image directory)] [video name (only if you select to test on video] [score (minimum confidence threshold value)]") 127 | print ("tests model on either webcam (w), video (v), or a directory of images (i)") 128 | 129 | def do_generatetfrecord(self, args): 130 | print (os.getcwd()) 131 | current_path = os.getcwd() 132 | os.chdir(current_path + '/prepare_records/') 133 | print (os.getcwd()) 134 | args = args.split() 135 | if (len(args) == 3): 136 | csvPath = args[0] 137 | outputPath = args[1] 138 | imagePath = args[2] 139 | print(csvPath) 140 | 141 | #generate for train data 142 | #os.system("pwd") 143 | os.system("python " + self.currentPath + "/prepare_records/" + "generate_tfrecord.py --csv_input=images/train_labels.csv --output_path=train.record --image_dir=images/train") 144 | 145 | #generate for test data 146 | os.system("python " + self.currentPath + "/prepare_records/" + "generate_tfrecord.py --csv_input=images/test_labels.csv --output_path=test.record --image_dir=images/test") 147 | #subprocess.call(['python generate_tfrecord.py', csvPath, outputPath, imagePath]) 148 | os.chdir('../') 149 | else: 150 | print("ERROR: Invalid syntax. Use help generatetfrecord for help.") 151 | return 152 | def help_generatetfrecord(self): 153 | print ("syntax: generatetfrecord [csv_path] [images_path] [output tf.record path]") 154 | print ("--generates both train.record and test.record for your tensorflow model, make sure your train_labels.csv is saved in images, and your images directories for train and test are saved as \'train\' and \'test\' in the images directory") 155 | 156 | def do_quit(self, args): 157 | """Quits the program.""" 158 | print("Quitting.") 159 | raise SystemExit 160 | 161 | 162 | if __name__ == '__main__': 163 | prompt = MyPrompt() 164 | prompt.prompt = '***PLASTIC-NET-IBM*** ' 165 | prompt.cmdloop('WELCOME TO PLASTIC-NET! (developed by IBM Space Tech, 2021) \n Type \'help\' to see a list of documented commands!') -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Plastic-Net Trash Detection (an IBM Space Tech Team project) 2 | 3 | **PlasticNet** is an "**IBM Tech for Good**" open-source project developed by the IBM Space Tech team to build a repository of AI object detection models to classify types/brands of plastics, trash on beaches, trash in the ocean, and more. 4 | 5 | We can scale this effort with the global community of developers participating and contributing towards this noble effort, with long term goals to help with ocean cleanup and positively impact climate change. 6 | 7 | For more information on how to get started, check out the [**PlasticNet Wiki**](https://github.com/IBM/PlasticNet/wiki)! 8 | 9 | ## Goals 10 | 11 | The **Goals** for our project are listed below as the following: 12 | 13 | - Real-time detection of different types of trash (plastic in particular) in the ocean utilizing transfer learning on different machine learning object detection architectures 14 | - In the future, we would also like to be able to improve our model to be able to recognize logos/brands on trash, in order to detect and identify which company different types of ocean/beach trash come from 15 | - To build a fully functional **PlasticNet** machine learning pipeline that can be easily used to train and test object detection models based from architectures such as YOLOv4, Faster-RCNN, SSD-Resnet, Efficient-DET, Tensorflow, etc. (all accessible inside a command line client) 16 | - To provide a set of pretrained **PlasticNet** models that can be utilized for future development and improvement via transfer learning 17 | - Implement our models to work on real-time satellite and camera footage 18 | 19 | ## Basic Project Structure and Technologies Used (feat. diagram) 20 | 21 | ![img](/img/PlasticNetProjectArchitectureDiagram.png) 22 | 23 | The PlasticNet command line program combines YOLOv4 and Tensorflow Object Detection API technologies into a single, easily usable machine learning pipeline CLI. Collaborators can use the PlasticNet CLI to prepare models for training (via transfer learning from the provided pre-trained PlasticNet models), train custom detection models built upon pre-trained PlasticNet models, export the trained models, and finally test the trained models. The CLI was created so these steps can all be done with [a few simple commands](https://github.com/IBM/PlasticNet/wiki/Utilizing-the-PlasticNet-Command-Line-Client). 24 | 25 | Initially trained via transfer learning from [pre-trained YOLO weights](https://github.com/mattokc35/darknet#pre-trained-models), and pre-trained Tensorflow models from the [Tensorflow Detection Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md), our official [**PlasticNet Model Zoo**](https://github.com/IBM/PlasticNet/blob/main/ModelZoo.md) can be used by collaborators for the further improvement/development of new PlasticNet object detection models. For labeling images, we utilized [IBM's Cloud Annotations](https://github.com/IBM/PlasticNet/wiki/Creating-Your-Own-Dataset-for-Custom-Training). 26 | 27 | ## Demo of Object Detection 28 | 29 | **YOLOv4 9 class v3 PlasticNet Demo (**best YOLO model**, with face masks and fishing nets included):** 30 | 31 | [![PlasticNet Demo](https://img.youtube.com/vi/eQRkLSfv8CY/0.jpg)](https://youtu.be/eQRkLSfv8CY "YOLO PlasticNet Demo") 32 | 33 | **Tensorflow EfficientDET D1 9 Class Demo:** 34 | 35 | [![PlasticNet Demo](https://img.youtube.com/vi/jWkY3OoC7Rw/0.jpg)](https://youtu.be/jWkY3OoC7Rw "Tensorflow EfficientDET D1 PlasticNet Demo") 36 | 37 | ## Get Started 38 | 39 | To get started with PlasticNet, you first must clone the repository using the following command: 40 | 41 | ```sh 42 | git clone https://github.com/IBM/PlasticNet.git 43 | ``` 44 | 45 | For more detailed instructions on how to get started, check out the [**PlasticNet Wiki**](https://github.com/IBM/PlasticNet/wiki)! 46 | 47 | Once the repository is cloned, you can run the following command in the python evironment of your choice. (NOTE: It is recommended that this is done in a new python environment, to avoid any issues between package dependencies) 48 | 49 | The setup script currently only supports MacOS, but Windows and Linux support will be added soon. 50 | 51 | ```sh 52 | cd PlasticNet && python setup.py 53 | ``` 54 | 55 | Once the setup script has finished running, you should restart your terminal session and run the following command: 56 | 57 | ```sh 58 | PlasticNet 59 | ``` 60 | 61 | This will open the PlasticNet terminal, so you can easily download our models from our model zoo, test the models on videos, webcam, and images, or train on top of an existing model. 62 | 63 | A list of all commands can be found by typing `help`, and more detailed instructions about aguments for any command can be found with `help [command name]`. 64 | 65 | To exit the command line, type `quit`. 66 | 67 | ## Using YOLO Models 68 | 69 | If you intend to train YOLO Models, you, may have to make some changes to the Makefile depending on your system. The Makefile is located in ./darknet, and you will want to update these parameters to whatver your system has, 1 meaning enabled. 70 | 71 | ```make 72 | GPU=1 73 | CUDNN=0 74 | CUDNN_HALF=0 75 | OPENCV=1 76 | AVX=0 77 | OPENMP=0 78 | LIBSO=0 79 | ``` 80 | 81 | It is highly recommended you install CUDNN & OpenCV for use with Darknet, as it will expedite the training process. 82 | 83 | You can additionally update the `yolo-obj.cfg` file (located in /darknet/cfg/) with any parameters you choose, specifically the number of iterations, classes, and filters. [Here's a guide for setting these values](https://github.com/mattokc35/darknet#how-to-train-to-detect-your-custom-objects). 84 | 85 | After you have updated this makefile and/or the configuration file, run `make clean` and `make` to be able to start darknet training with the PlasticNet CLI. 86 | 87 | ## Test Results 88 | 89 | [See our spreadsheet documentating our test results from different trained models](https://docs.google.com/spreadsheets/d/1mcFC2HqjohRp2_G723D8_Xd19LuuUKD-frROiY5ksOQ/edit?usp=sharing). 90 | 91 | ## Resources 92 | 93 | - [Forked YOLOv4 PlasticNet Darknet repository](https://github.com/mattokc35/darknet) 94 | - [Labeling Images with IBM Cloud Annotations](https://cloud.annotations.ai/) 95 | - [Darknet YOLO]((https://pjreddie.com/darknet/yolo/)) 96 | - [Restoring Integrity to the Oceans (RIO)](https://www.oceansintegrity.com/) 97 | - [Pacific Whale Foundation](https://www.pacificwhale.org/) 98 | - [Tensorflow Object Detection API](https://github.com/tensorflow/models/tree/master/research/object_detection) 99 | - [Yolo with Tensorflow](https://github.com/theAIGuysCode/tensorflow-yolov4-tflite) 100 | 101 | ## Ocean Plastic Statistics 102 | 103 | - More than 1 million seabirds and 100,000 marine animals die from plastic pollution every year. 104 | - 100% of baby sea turtles have plastic in their stomachs. 105 | - There is now 5.25 trillion macro and micro pieces of plastic in our ocean & 46,000 pieces in every square mile of ocean, weighing up to 269,000 tonnes. 106 | - Every day around 8 million pieces of plastic makes their way into our oceans. 107 | - The Great Pacific Garbage Patch is around 1.6 million square kilometers – bigger than Texas. 108 | - The world produces 381 million tonnes in plastic waste yearly – this is set to double by 2034. 109 | - 50% of this is single-use plastic & only 9% has ever been recycled. 110 | - Over 2 million tonnes of plastic packaging are used in the UK each year. 111 | - 88% of the sea's surface is polluted by plastic waste. 112 | - Between 8 to 14 million tonnes enters our ocean every year. 113 | - Britain contributes an estimated 1.7 million tonnes of plastic annually. 114 | - The US contributes 38 million tonnes of plastic every year. 115 | - Plastic packaging is the biggest culprit, resulting in 80 million tonnes of waste yearly from the US alone. 116 | - On UK beaches there are 5000 pieces of plastic & 150 plastic bottles for each mile. 117 | - More than 1 million plastic bags end up in the trash every minute. 118 | - The world uses over 500 billion plastic bags a year – that’s 150 for each person on Earth. 119 | - 8.3 billion plastic straws pollute the world’s beaches, but only 1% of straws end up as waste in the ocean. 120 | - By 2020 the number of plastics in the sea will be higher than the number of fish. 121 | - 1 in 3 fish caught for human consumption contains plastic. 122 | - Plastic microbeads are estimated to be one million times more toxic than the seawater around it. 123 | - Products containing microbeads can release 100,000 tiny beads with just one squeeze. 124 | 125 | Source: [Condor Ferries: _Shocking Ocean Plastic Statistics_](https://www.condorferries.co.uk/plastic-in-the-ocean-statistics) 126 | -------------------------------------------------------------------------------- /ModelZoo.md: -------------------------------------------------------------------------------- 1 | # PlasticNet Model Zoo 2 | 3 | **For pre-trained weights to train a PlasticNet model from scratch, see the bottom of this page!** 4 | 5 | ## TensorFlow Models 6 | 7 | - [**PlasticNet EfficientDET D1 **9 classes**640x640**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/efficientdet_d1_640x640_9classes_v2.tar.gz) 8 | 9 | Model File Name: efficientdet_d1_640x640_9classes_v2 10 | 11 | Transfer Learning From: TensorFlow efficientdet_d1 12 | 13 | Average Loss: 14 | 15 | ```sh 16 | class_loss: 0.08758, 17 | localization_loss: 7.339e-4, 18 | regularization_loss: 0.08055, 19 | total_loss: 0.1686 20 | ``` 21 | 22 | - [**PlasticNet Faster-RCNN **9 classes** 640x640**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/faster_rcnn_640x640_9classes_v1.tar.gz) 23 | 24 | Model File Name: faster_rcnn_640x640_9classes_v1 25 | 26 | Transfer Learning From: PlasticNet Model: faster_rcnn_640x640_7classes_v2 27 | 28 | Average Loss: 29 | 30 | ```sh 31 | localization_loss: 0.16872177, 32 | localization_loss: 0.24726836, 33 | objectness_loss: 0.05993971, 34 | regularization_loss: 0.0, 35 | total_loss: 0.63053286, 36 | ``` 37 | 38 | - [**PlasticNet FasterRCNN 7 Classes 640x640 v2**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/faster_rcnn_640x640_7classes_v2.tar.gz) 39 | 40 | Model File Name: faster_rcnn_640x640_7classes_v2 41 | 42 | Transfer Learning From: PlasticNet Model: faster_rcnn_640x640_7classes_v1 43 | 44 | Average Loss: 45 | 46 | ```sh 47 | classification_loss: 0.1095 48 | localization_loss: 0.3035 49 | objectness_loss: 0.01122, 50 | regularization_loss: 0.0, 51 | total_loss: 0.4964 52 | ``` 53 | 54 | - [**PlasticNet FasterRCNN 7 Classes 640x640 v1**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/faster_rcnn_640x640_7classes_v1.tar.gz) 55 | 56 | Model File Name: faster_rcnn_640x640_7classes_v1 57 | 58 | Transfer Learning From: TensorFlow Base Model: Faster R-CNN ResNet101 V1 640x640 59 | 60 | Average Loss: 61 | 62 | ```sh 63 | classification_loss: 0.006, 64 | localization_loss: 0.0024, 65 | regularization_loss: 0.00, 66 | total_loss: 0.0224 67 | ``` 68 | 69 | - [**PlasticNet SSD-Resnet 7 Classes 640x640**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/ssd_resnet_640x640_7classes_v1.tar.gz) 70 | 71 | Model File Name: ssd_resnet_640x640_7classes_v1 72 | 73 | Transfer Learning From: TensorFlow base model: SSD ResNet152 V1 FPN 640x640 (RetinaNet152) 74 | 75 | Average Loss: 76 | 77 | ```sh 78 | classification_loss: 0.06852589 79 | localization_loss: 0.025546694, 80 | regularization_loss: 0.35246804, 81 | total_loss: 0.4465 82 | ``` 83 | 84 | ## YOLO Models 85 | 86 | - **_Best YOLO Model_** **PlasticNet YOLO 9 Classes 448x448 v3** 87 | 88 | Conv File Name (for transfer learning, click to download): [yolo9class448x448v3.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class448x448v3.conv.81.tar.gz) 89 | 90 | Weights File Name (for testing, click to download): [yolo9class448x448v3weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class448x448v3weights.tar.gz) 91 | 92 | Average Loss: 2.33 93 | 94 | - **PlasticNet YOLO 9 Classes 448x448 v4** 95 | 96 | Transfer Learning Base Model: PlasticNet YOLO 9 Classes 448x448 ([yolo9class448x448v3.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class448x448v3.conv.81.tar.gz)) 97 | 98 | Conv File Name (for transfer learning, click to download): [yolo9class448x448v4.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class448x448v4.conv.81.tar.gz) 99 | 100 | Weights File Name (for testing, click to download): [yolo9class448x448v4weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class448x448v4weights.tar.gz) 101 | 102 | Average Loss: 2.32 103 | 104 | - **PlasticNet YOLO 9 Classes 480x480 v2** 105 | 106 | Conv File Name (for transfer learning, click to download): [yolo9classv2iter27000.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9classv2iter27000.conv.81.tar.gz) 107 | 108 | Weights File Name (for testing, click to download): [yolo9classv2iter27000weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9classv2iter27000weights.tar.gz) 109 | 110 | Average Loss: 2.183 111 | 112 | - **PlasticNet YOLO 9 Classes 480x480 v1** 113 | 114 | Conv File Name (for transfer learning, click to download): [yolo9class480x480.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class480x480.conv.81.tar.gz) 115 | 116 | Weights File Name (for testing, click to download): [yolo9class480x480weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo9class480x480weights.tar.gz) 117 | 118 | Average Loss: 2.378 119 | 120 | - **PlasticNet YOLO 7 Classes v3 576x576** 121 | 122 | Conv File Name (for transfer learning, click to download): [yolo7class576x576.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo7class576x576.conv.81.tar.gz) 123 | 124 | Weights File Name (for testing, click to download): [yolo7class576x576weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo7class576x576weights.tar.gz) 125 | 126 | Average Loss: 2.45 127 | 128 | - **PlasticNet YOLO 7 Classes v2 416x416** 129 | 130 | Conv File Name (for transfer learning, click to download): [7classyoloV2.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/7classyoloV2.conv.81.tar.gz) 131 | 132 | Weights File Name (for testing, click to download): [7classyoloV2weights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/7classyoloV2weights.tar.gz) 133 | 134 | Average Loss: 2.47 135 | 136 | - **PlasticNet YOLO 4 Classes Image Augmentation 416x416** 137 | 138 | Conv File Name (for transfer learning, click to download): [yolo4classesimgaug.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo4classesimgaug.conv.81.tar.gz) 139 | 140 | Weights File Name (for testing, click to download): [yolo4classesimgaugweights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo4classesimgaugweights.tar.gz) 141 | 142 | Average Loss: 1.48 143 | 144 | - **PlasticNet YOLO 4 Classes v2 416x416** 145 | 146 | Conv File Name (for transfer learning, click to download): [yolo4classesupdated.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo4classesupdated.conv.81.tar.gz) 147 | 148 | Weights File Name (for testing, click to download): [yolo4classesupdatedweights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo4classesupdatedweights.tar.gz) 149 | 150 | Average Loss: N/A 151 | 152 | - **PlasticNet YOLO 4 Classes v1 8000 iterations 416x416** 153 | 154 | Conv File Name (for transfer learning, click to download): [4class8000iteration.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/4class8000iteration.conv.81.tar.gz) 155 | 156 | Weights File Name (for testing, click to download): [4class8000iteration](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/4class8000iteration.tar.gz) 157 | 158 | Average Loss: N/A 159 | 160 | - **PlasticNet YOLO 2 Classes Image Augmentation 416x416** 161 | 162 | Conv File Name (for transfer learning, click to download): [yolo2classes-imguag.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo2classes-imguag.conv.81.tar.gz) 163 | 164 | Weights File Name (for testing, click to download): [yolo2classes-imguagweights](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo2classes-imguagweights.tar.gz) 165 | 166 | Average Loss: N/A 167 | 168 | - **PlasticNet YOLO 2 Classes v1 416x416** 169 | 170 | Conv File Name (for transfer learning, click to download): [yolo2classes3000iterations.conv.81](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo2classes3000iterations.conv.81.tar.gz) 171 | 172 | Weights File Name (for testing, click to download): [yolo2classes3000iterations](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolo2classes3000iterations.tar.gz) 173 | 174 | Average Loss: 6.67 175 | 176 | - [**Darknet YOLO 12 Classes Pre-Trained Weights**](https://plasticnet-models.s3.us.cloud-object-storage.appdomain.cloud/yolov4.conv.137.tar.gz) 177 | 178 | **_This pre-trained weights is if you want to train a new PlasticNet model from scratch using the provided Darknet generic pretrained weights from [https://github.com/mattokc35/darknet](https://github.com/mattokc35/darknet), not using a pre-trained PlasticNet model!_** 179 | 180 | Model File Name: yolov4.conv.137 181 | 182 | Average Loss: N/A 183 | -------------------------------------------------------------------------------- /PlasticNet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | # Copyright 2021 IBM Corporation 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | from cmd import Cmd 17 | import sys 18 | import os 19 | import subprocess 20 | global currentModel 21 | global currentPath 22 | class MyPrompt(Cmd): 23 | 24 | 25 | def __init__(self): 26 | """ 27 | Constructor Function 28 | """ 29 | super(MyPrompt, self).__init__() 30 | self.currentModel = "default" 31 | 32 | # If CLI is installed, gets the path to the CLI. Otherwise, runs the path of this direct file as it was invoked directly 33 | if (os.environ.get("PLASTICNET_PATH") != None): 34 | self.currentPath = os.environ.get('PLASTICNET_PATH') 35 | else: 36 | self.currentPath = os.path.dirname(os.path.abspath(__file__)) 37 | 38 | def do_darknet_train(self, args): 39 | """ 40 | Trains a model using a pretrained YOLOv4 Model 41 | Parameters: 42 | model_name: Name of the model that was downloaded 43 | Return: 44 | None 45 | """ 46 | args = args.split() 47 | if (len(args) >= 1): 48 | self.currentModel = args[0] 49 | print("TRAINING YOLOV4 MODEL: " + str(self.currentModel)) 50 | os.chdir("darknet") 51 | os.system('./darknet detector train data/obj.data cfg/yolo-obj.cfg ' + str(self.currentModel)) 52 | os.chdir("../") 53 | def do_darknet_export(self, args): 54 | """ 55 | Freezes the initial layers of a model so that the final layers can be retrained 56 | Parameters: 57 | model_name: name you want the conv file to be called 58 | Return: 59 | None 60 | """ 61 | args = args.split() 62 | if (len(args) >= 1): 63 | print("FREEZING LAYERS OF YOLO MODEL: ") 64 | self.currentModel = args[0] 65 | print(self.currentModel) 66 | os.chdir("darknet") 67 | os.system('./darknet partial cfg/yolo-obj.cfg ' + 'backup/yolo-obj_final.weights' + " " + str(self.currentModel) +'.conv.81 81') 68 | os.chdir("../") 69 | else: 70 | print("Please specify the name you want the conv file to be called as an argument.") 71 | def help_darknet_export(self): 72 | """ 73 | Help Function 74 | """ 75 | print("syntax: darknet_export [string to name file]") 76 | print("Exports a model for transfer learning") 77 | def help_darknet_train(self): 78 | """ 79 | Help Function 80 | """ 81 | print("syntax: darknet_train [model_name]") 82 | print("trains model using pre-trained base model specified by user") 83 | def do_darknet_test(self, args): 84 | """ 85 | Tests a model using a pretrained YOLOv4 Model on either webcam, video, or images 86 | Parameters: 87 | video_name, score (minimum confidence threshold) 88 | """ 89 | os.chdir("darknet") 90 | args = args.split() 91 | if (len(args) >= 2): 92 | mode = args[0] 93 | if(mode == 'w'): 94 | score = args[1] 95 | print("Testing YOLOv4 Darknet on webcam...") 96 | os.system("./darknet detector demo data/obj.data cfg/yolo-obj.cfg " + str(self.currentModel) + ' -thresh ' + str(score) + " -c 0" ) 97 | #webcam 98 | if(mode == 'v'): 99 | if (len(args == 3)): 100 | video_name = args[1] 101 | score = args[2] 102 | else: 103 | print("ERROR: Invalid syntax. Use help darknet_test for help.") 104 | os.chdir("../") 105 | return 106 | print("Testing YOLOv4 Darknet on video...") 107 | os.system('./darknet detector demo data/obj.data cfg/yolo-obj.cfg ' + str(self.currentModel) + ' ./data/' + str(video_name) + ' -thresh ' + str(score)) 108 | #video 109 | if(mode == 'i'): 110 | score = args[1] 111 | print("Testing YOLOv4 Darknet on images...") 112 | os.system('./darknet detector test data/obj.data cfg/yolo-obj.cfg ' + str(self.currentModel) + ' -thresh ' + str(score)) 113 | else: 114 | print("ERROR: Invalid syntax. Use help darknet_test for help.") 115 | os.chdir("../") 116 | return 117 | os.chdir("../") 118 | 119 | def help_darknet_test(self): 120 | """ 121 | Help Function 122 | """ 123 | print("syntax: darknet_test [mode (w|v|i)] [video_name (only if \'v\' mode)] [score]") 124 | print("tests YOLOv4 Darknet model on either webcam, video, or images with a specified minimum confidence threshold (score)") 125 | 126 | 127 | def do_set_model(self, args): 128 | """ 129 | Function sets current model global variable that can be used for preparing, training, exporting, and testing. 130 | """ 131 | print("SETTING MODEL") 132 | args = args.split() 133 | if (len(args) == 1): 134 | self.currentModel = args[0] 135 | else: 136 | print("ERROR: Invalid syntax. Use help set_model for help.") 137 | return 138 | print('current model set to: ' + str(self.currentModel)) 139 | def help_set_model(self): 140 | """ 141 | Help Function 142 | """ 143 | print('--sets current model global variable that can be used for preparing, training, exporting, and testing') 144 | print('syntax: set_model [model name]') 145 | 146 | 147 | def do_prepare_training(self, args): 148 | """ 149 | Function generates all the necessary files for training to begin. After this is complete, run train_model. 150 | """ 151 | os.chdir(self.currentPath) 152 | args = args.split() 153 | model = "" 154 | num_classes = -1 155 | if (len(args) == 3): 156 | #yolo or tf 157 | mode = args[0] 158 | model = args[1] 159 | num_classes = args[2] 160 | print("PREPARING TRAINING FOR MODEL: " + str(model)) 161 | elif (len(args) == 2): 162 | mode = args[0] 163 | model = args[1] 164 | if ('.weights' in model): 165 | model = model[:-8] + '.conv.81' 166 | print("PREPARING TRAINING FOR MODEL: " + str(model)) 167 | else: 168 | print("ERROR: Invalid syntax. Use help prepare_training for help.") 169 | return 170 | if mode == 't': 171 | return_code = os.system("python prepare_training.py --arch_type " + str(mode) + " --model_name " + str(model) + " --num_classes " + str(num_classes)) 172 | elif mode == 'y': 173 | return_code = os.system("python prepare_training.py --arch_type " + str(mode) + " --yolo_weights " + str(model) + " --num_classes " + str(num_classes)) 174 | weightsString = model.replace('weights', '.weights') 175 | os.system("cp out/" + str(weightsString) + " " + str(self.currentPath) + "/darknet") 176 | os.system("rm out/" + str(weightsString)) 177 | 178 | #split train and test data for YOLOv4 179 | os.chdir("darknet") 180 | print(str(os.getcwd())) 181 | os.system("python traintestsplit.py") 182 | os.chdir("../") 183 | 184 | 185 | 186 | 187 | if return_code == 0: # successful 188 | self.currentModel = args[1] 189 | print(self.currentModel) 190 | else: 191 | print("ERROR: Prepare training has failed. Check the name of the model you've specified.") 192 | def help_prepare_training(self): 193 | """ 194 | Help Function 195 | """ 196 | print("syntax: prepare_training [model_type (YOLOv4 (y) or Tensorflow (t))] [model_name (tf) or yolo_weights (y)] [OPTIONAL: num_classes]\n") 197 | print("generates all the necessary files for training to begin (for Tensorflow). For YOLOv4, downloads the specified pre-trained weights from PlasticNet Model Zoo. After this is complete, run train_model\n") 198 | print("\n") 199 | print("List of available models:\n") 200 | print("TensorFlow:") 201 | print("faster_rcnn_640x640_9classes_v1\nfaster_rcnn_640x640_7classes_v2\nfaster_rcnn_640x640_7classes_v1\nssd_resnet_640x640_7classes_v1") 202 | print("YOLO:") 203 | print("yolo9classv2iter27000weights\nyolo9class480x480weights\nyolo7class576x576weights\n7classyoloV2weights\nyolo4classesimgaugweights"+ 204 | "\nyolo4classesupdatedweights\n4class8000iteration\nyolo2classes-imguag\nyolo2classes3000iterations\nyolov4weights") 205 | 206 | 207 | def do_train_model(self, args): 208 | """ 209 | Trains model using pre-trained base model specified by user 210 | """ 211 | os.chdir(self.currentPath) 212 | print("TRAINING MODEL: " + str(self.currentModel)) 213 | os.system("python train_model.py --pipeline_config_path=training/pipeline.config/ " + 214 | "--model_dir=out/" + str(self.currentModel)) 215 | def help_train_model(self): 216 | """ 217 | Help Function 218 | """ 219 | print ("syntax: train_model (no arguments needed, will use global variable \'current_model\'") 220 | print ("trains model using pre-trained base model specified by user") 221 | 222 | 223 | def do_export_model(self, args): 224 | """ 225 | Function exports a model after training. Specify the same model name as you used for training 226 | """ 227 | os.chdir(self.currentPath) 228 | args = args.split() 229 | if (len(args) == 1): 230 | model = args[0] 231 | else: 232 | print("ERROR: Invalid syntax. Use help export_model for help.") 233 | return 234 | print("EXPORTING MODEL: " + str(model)) 235 | os.chdir(self.currentPath) 236 | os.chdir('models/research/object_detection') 237 | os.system("python exporter_main_v2.py --input_type image_tensor --pipeline_config_path" + 238 | " ../../../training/pipeline.config --trained_checkpoint_dir ../../../out/" + str(model) + " " + 239 | "--output_directory ../../../out/" + str(model) + "/exported_model") 240 | os.chdir('../../../') 241 | def help_export_model(self): 242 | """ 243 | Help Function 244 | """ 245 | print("syntax: export_model [model_name]") 246 | print("Exports a model after training. Specify the same model name as you used for training") 247 | def do_test_model(self, args): 248 | """ 249 | Function tests model on either webcam (w), video (v), or a directory of images (i). 250 | """ 251 | os.chdir(self.currentPath) 252 | args = args.split() 253 | if (len(args) >= 1): 254 | mode = args[0] 255 | if(mode == 'w'): 256 | score = args[1] 257 | print('Testing on WEBCAM: ' + str(self.currentModel)) 258 | current_path = self.currentPath 259 | os.chdir(current_path + '/test_model') 260 | current_path = os.getcwd() 261 | os.system("python " + "/test_model/" + "testTensorWebcam.py --trained_model " + str(self.currentModel) + " --score " + str(score)) 262 | os.chdir('../') 263 | current_path = os.getcwd() 264 | print(current_path) 265 | elif(mode == 'v'): 266 | if (len(args) == 3): 267 | videoName = args[1] 268 | score = args[2] 269 | print('Testing on VIDEO: ' + str(self.currentModel)) 270 | current_path = self.currentPath 271 | os.chdir(current_path + '/test_model') 272 | current_path = os.getcwd() 273 | os.system("python " + "/test_model/" + "testTensorVideo.py --video_name " + str(videoName) + " --trained_model " + str(self.currentModel) + " --score " + str(score)) 274 | os.chdir('../') 275 | current_path = os.getcwd() 276 | print(current_path) 277 | else: 278 | print("ERROR: Invalid syntax. Use help test_model for help.") 279 | return 280 | elif(mode == 'i'): 281 | score = args[1] 282 | print('Testing on IMAGES: ' + str(self.currentModel)) 283 | current_path = self.currentPath 284 | os.chdir(current_path + '/test_model') 285 | current_path = os.getcwd() 286 | os.system("python " + "/test_model/" + "testTensorflow.py --trained_model" + str(self.currentModel) + " --score " + str(score)) 287 | os.chdir('../') 288 | current_path = os.getcwd() 289 | print(current_path) 290 | else: 291 | print("ERROR: Invalid mode. Use help test_model for help.") 292 | else: 293 | print("ERROR: Invalid syntax. Use help test_model for help.") 294 | return 295 | def help_test_model(self): 296 | """ 297 | Help Function 298 | """ 299 | print ("syntax: test_model [w (webcam) | v (video) | i (image directory)] [video name (only if you select to test on video] [score (minimum confidence threshold value)]") 300 | print ("tests model on either webcam (w), video (v), or a directory of images (i)") 301 | 302 | 303 | def do_generatetfrecord(self, args): 304 | """ 305 | Function generates both train.record and test.record for your tensorflow model, make sure your train_labels.csv is saved in images, and your images directories for train and test are saved as \'train\' and \'test\' in the images directory. 306 | """ 307 | current_path = str(self.currentPath) 308 | os.chdir(current_path + '/prepare_records/') 309 | # Partition the dataset 310 | os.system("python partition_dataset.py") 311 | # Create test_labels.csv and train_labels.csv 312 | os.system("python xml_to_csv.py") 313 | #generate for train data 314 | os.system("python generate_tfrecord.py --csv_input=images/train_labels.csv --output_path=train.record --image_dir=images/train") 315 | #generate for test data 316 | os.system("python generate_tfrecord.py --csv_input=images/test_labels.csv --output_path=test.record --image_dir=images/test") 317 | #subprocess.call(['python generate_tfrecord.py', csvPath, outputPath, imagePath]) 318 | os.chdir('../') 319 | def help_generatetfrecord(self): 320 | """ 321 | Help Function 322 | """ 323 | print("syntax: generatetfrecord") 324 | print("--generates both train.record and test.record for your tensorflow model," + "Make sure all images are stored in prepare_records/images") 325 | 326 | 327 | def do_quit(self, args): 328 | """Quits the program.""" 329 | print("Quitting.") 330 | raise SystemExit 331 | 332 | 333 | if __name__ == '__main__': 334 | prompt = MyPrompt() 335 | prompt.prompt = '***PLASTIC-NET-IBM*** ' 336 | prompt.cmdloop('WELCOME TO PLASTIC-NET! (developed by IBM Space Tech, 2021) \n Type \'help\' to see a list of documented commands!') 337 | --------------------------------------------------------------------------------