DAY 77-100 DAYS MLCODE: Custom Object detection

Custom Object detection: In the past few blogs, we discussed object detection using ImageAI, TensorFlow and Yolo V3 using CV2, in this blog, we’ll try to use TensorFlow Object detection model to create Custom Object detection model.

Create Dataset

First, we have to annotate the image and for the learning purpose, I created a custom class and training data has the image of 50 dogs which I kept as class PT.

Create TFRecord

%cd ~/datalab

!python /content/models/research/object_detection/dataset_tools/create_pet_tf_record.py –label_map_path=label_map.pbtxt –data_dir=. –output_dir=. –num_shards=1

!mv pet_faces_train.record-00000-of-00001 tf_train.record

!mv pet_faces_val.record-00000-of-00001 tf_val.record

Download pre-trained model

We’ll use pre-trained Faster RCNN Inception V2 model which is trained on coco objects

MODEL = ‘faster_rcnn_inception_v2_coco_2018_01_28’
MODEL_FILE = MODEL + ‘.tar.gz’
DOWNLOAD_BASE = ‘http://download.tensorflow.org/models/object_detection/’
DEST_DIR = ‘pretrained_model’

if not (os.path.exists(MODEL_FILE)):
opener = urllib.URLopener()
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)

tar = tarfile.open(MODEL_FILE)
tar.extractall()
tar.close()

os.remove(MODEL_FILE)
if (os.path.exists(DEST_DIR)):
shutil.rmtree(DEST_DIR)
os.rename(MODEL, DEST_DIR)

Edit model config file

To you use a different trained model in step before, update accordingly filename var and re.sub functions in next cell.

with open(filename) as f:
s = f.read()
with open(filename, ‘w’) as f:
s = re.sub(‘PATH_TO_BE_CONFIGURED/model.ckpt’, ‘/content/datalab/pretrained_model/model.ckpt’, s)
s = re.sub(‘PATH_TO_BE_CONFIGURED/pet_faces_train.record-\?\?\?\?\?-of-00010’, ‘/content/datalab/tf_train.record’, s)
s = re.sub(‘PATH_TO_BE_CONFIGURED/pet_faces_val.record-\?\?\?\?\?-of-00010’, ‘/content/datalab/tf_val.record’, s)
s = re.sub(‘PATH_TO_BE_CONFIGURED/pet_label_map.pbtxt’, ‘/content/datalab/label_map.pbtxt’, s)
f.write(s)

Train model

%cd ~/datalab

!python ~/models/research/object_detection/model_main.py \
–pipeline_config_path=/content/models/research/object_detection/samples/configs/faster_rcnn_inception_v2_pets.config \
–model_dir=/content/datalab/trained \
–alsologtostderr \
–num_train_steps=3000 \
–num_eval_steps=500

Export trained model with highest step number in filename.

%cd ~/datalab

lst = os.listdir(‘trained’)
lf = filter(lambda k: ‘model.ckpt-‘ in k, lst)
last_model = sorted(lf)[-1].replace(‘.meta’, ”)

!python ~/models/research/object_detection/export_inference_graph.py \
–input_type=image_tensor \
–pipeline_config_path=/content/models/research/object_detection/samples/configs/faster_rcnn_inception_v2_pets.config \
–output_directory=fine_tuned_model \
–trained_checkpoint_prefix=trained/$last_model

You can find the entire code here.