rpn_simple.py

import keras.backend
import keras.engine
import keras.layers
import numpy
import numpy as np
import tensorflow
import keras_resnet.models
import keras_resnet.blocks
import keras_rcnn.backend
import keras_rcnn.classifiers
import keras_rcnn.layers
import cv2

img_size = 1036

configuration = tensorflow.ConfigProto()

session = tensorflow.Session(config=configuration)

keras.backend.set_session(session)

image = keras.layers.Input((img_size, img_size, 1), name="image")

# Bounding boxes are the ground truth bounding boxes
bounding_boxes = keras.layers.Input((None, 4), name="bounding_boxes")

labels = keras.layers.Input((None, 2), name="labels")

# Metadata is a 3-tuple that contains the feature width, feature height, and scale
metadata = keras.layers.Input((3,), name="metadata")

options = {
    "activation": "relu",
    "kernel_size": (3, 3),
    "padding": "same"
}

features = keras.layers.Conv2D(64, name="convolution_1_1", **options)(image)
features = keras.layers.Conv2D(64, name="convolution_1_2", **options)(features)

features = keras.layers.MaxPooling2D(strides=(2, 2), name="max_pooling_1")(features)

features = keras.layers.Conv2D(128, name="convolution_2_1", **options)(features)
features = keras.layers.Conv2D(128, name="convolution_2_2", **options)(features)

features = keras.layers.MaxPooling2D(strides=(2, 2), name="max_pooling_2")(features)

features = keras.layers.Conv2D(256, name="convolution_3_1", **options)(features)
features = keras.layers.Conv2D(256, name="convolution_3_2", **options)(features)
features = keras.layers.Conv2D(256, name="convolution_3_3", **options)(features)

features = keras.layers.MaxPooling2D(strides=(2, 2), name="max_pooling_3")(features)

features = keras.layers.Conv2D(512, name="convolution_4_1", **options)(features)
features = keras.layers.Conv2D(512, name="convolution_4_2", **options)(features)
features = keras.layers.Conv2D(512, name="convolution_4_3", **options)(features)

features = keras.layers.MaxPooling2D(strides=(2, 2), name="max_pooling_4")(features)

features = keras.layers.Conv2D(512, name="convolution_5_1", **options)(features)
features = keras.layers.Conv2D(512, name="convolution_5_2", **options)(features)
features = keras.layers.Conv2D(512, name="convolution_5_3", **options)(features)

convolution_3x3 = keras.layers.Conv2D(512, name="convolution_3x3", **options)(features)

deltas = keras.layers.Conv2D(9 * 4, (1, 1), name="deltas")(convolution_3x3)
scores = keras.layers.Conv2D(9 * 2, (1, 1), name="scores")(convolution_3x3)

all_anchors, rpn_labels, bounding_box_targets = keras_rcnn.layers.AnchorTarget()([scores, bounding_boxes, metadata])

scores_reshaped = keras.layers.Reshape((-1, 2))(scores)
scores_reshaped = keras.layers.Activation('softmax')(scores_reshaped)

scores_loss = keras_rcnn.layers.losses.RPNClassificationLoss(9)([scores_reshaped, rpn_labels])

model = keras.models.Model([image, bounding_boxes, metadata, labels], [scores_loss, rpn_labels, all_anchors])

from keras.optimizers import Adam
opt = Adam(lr=1e-5)
model.compile(opt, None)

x_image = numpy.zeros((1, img_size, img_size, 1))

x_boxes = numpy.zeros((1, 1, 4))

x_boxes[0,0,0] = 64
x_boxes[0,0,1] = 64
x_boxes[0,0,2] = 64 + 128
x_boxes[0,0,3] = 64 + 128

for box_idx in range(x_boxes.shape[1]):
	x1, y1, x2, y2 = x_boxes[0,box_idx, :]
	x_image[0,int(y1):int(y2),int(x1):int(x2),:] = 1.0

x_metadata = numpy.expand_dims([img_size, img_size, 1], 0)

x_labels = numpy.zeros((1,1,2))
x_labels[0,0,0] = 1

model.fit([x_image, x_boxes, x_metadata, x_labels], epochs=200)

P = model.predict([x_image, x_boxes, x_metadata, x_labels])

img = np.zeros((img_size, img_size, 3)).astype(np.uint8)

for box_idx in range(x_boxes.shape[1]):
	x1, y1, x2, y2 = x_boxes[0,box_idx, :]
	img[int(y1):int(y2),int(x1):int(x2), :] = 255

for ix in range(P[0].shape[1]):
	if P[0][0, ix, 1] > 0.9:
		x1, y1, x2, y2 = P[2][0,ix,:]
		cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255))

img = cv2.resize(img, None, fx=0.5, fy=0.5)
cv2.imwrite('rpn.png', img)