Merge pull request #27 from JihongJu/rcnn_losses

Build R-CNN with ResNet

keras_rcnn/heads/__init__.py

@@ -0,0 +1 @@
+from .resnet import ResHead

keras_rcnn/heads/resnet.py

@@ -0,0 +1,40 @@
+import keras
+import keras_resnet.block.temporal
+
+
+def ResHead(classes, mask=False):
+    """Resnet heads as in Mask R-CNN."""
+    def f(x):
+        if keras.backend.image_data_format() == "channels_last":
+            channel_axis = 3
+        else:
+            channel_axis = 1
+
+        y = keras.layers.TimeDistributed(
+            keras.layers.Conv2D(1024, (1, 1)))(x)
+
+        # conv5 block as in Deep Residual Networks with first conv operates
+        # on a 7x7 RoI with stride 1 (instead of 14x14 / stride 2)
+        for i in range(3):
+            y = keras_resnet.block.temporal.bottleneck(
+                512, (1, 1), first=True)(y)
+
+        y = keras.layers.TimeDistributed(
+            keras.layers.BatchNormalization(axis=channel_axis))(y)
+        y = keras.layers.TimeDistributed(
+            keras.layers.Activation("relu"))(y)
+
+        # class and box branches
+        y = keras.layers.TimeDistributed(
+            keras.layers.AveragePooling2D((7, 7)))(y)
+
+        score = keras.layers.TimeDistributed(
+            keras.layers.Dense(classes, activation="softmax"))(y)
+
+        boxes = keras.layers.TimeDistributed(
+            keras.layers.Dense(4 * classes))(y)
+
+        # TODO{JihongJu} the mask branch
+
+        return [score, boxes]
+    return f

keras_rcnn/models.py

@@ -1,25 +1,67 @@
 # -*- coding: utf-8 -*-
 
 import keras
+import keras_resnet
 import keras_rcnn.layers
+import keras_rcnn.heads
 
 
 class RCNN(keras.models.Model):
-    def __init__(self, inputs, classes, regions_of_interest):
-        y = keras_rcnn.layers.ROI(14, regions_of_interest)(inputs)
+    """
+    Faster R-CNN model by S Ren et, al. (2015).
 
-        y = keras.layers.Dense(4096)(y)
-        y = keras.layers.TimeDistributed(y)
+    :param inputs: input tensor (e.g. an instance of `keras.layers.Input`)
+    :param encoder: (Convolutional) feature extractor,
+        e.g., `keras_resnet.ResNet50`
+    :param heads: R-CNN heads for object detection and/or segmentation
+        on the proposed regions
+    :param rois: integer, number of regions of interest per image
 
-        score = keras.layers.Dense(classes)(y)
-        score = keras.layers.Activation("softmax")(score)
-        score = keras.layers.TimeDistributed(score)
+    :return model: a functional model API for R-CNN.
+    """
 
-        boxes = keras.layers.Dense(4 * (classes - 1))(y)
-        boxes = keras.layers.Activation("linear")(boxes)
-        boxes = keras.layers.TimeDistributed(boxes)
+    def __init__(self, inputs, encoder, heads, rois):
+        # Extract features with the encoder
+        y = encoder(inputs)
+        features = y.layers[-2].output
+        # Propose regions given the features
+        rpn_classification = keras.layers.Conv2D(
+            9 * 1, (1, 1), activation="sigmoid")(features)
+        rpn_regression = keras.layers.Conv2D(9 * 4, (1, 1))(features)
 
-        super(RCNN, self).__init__(inputs, [score, boxes])
+        rpn_prediction = keras.layers.concatenate(
+            [rpn_classification, rpn_regression])
+
+        proposals = keras_rcnn.layers.object_detection.ObjectProposal(
+            rois)([rpn_regression, rpn_classification])
+        # Apply the heads on the proposed regions
+        slices = keras_rcnn.layers.ROI((7, 7))([inputs, proposals])
+
+        [score, boxes] = heads(slices)
+
+        super(RCNN, self).__init__(inputs,
+                                   [rpn_prediction, score, boxes])
+
+
+class ResNet50RCNN(RCNN):
+    """
+    Faster R-CNN model with ResNet50.
+
+    :param inputs: input tensor (e.g. an instance of `keras.layers.Input`)
+    :param classes: integer, number of classes
+    :param rois: integer, number of regions of interest per image
+
+    :return model: a functional model API for R-CNN.
+    """
+
+    def __init__(self, inputs, classes, rois=300):
+        # ResNet50 as encoder
+        encoder = keras_resnet.ResNet50
+
+        # ResHead with score and boxes
+        heads = keras_rcnn.heads.ResHead(classes)
+
+        super(ResNet50RCNN, self).__init__(inputs, encoder, heads, rois)
 
 
 class RPN(keras.models.Model):

setup.py

@@ -14,7 +14,8 @@
         ],
     },
     install_requires=[
-        "keras"
+        "keras",
+        "keras-resnet"
     ],
     license="MIT",
     name="keras-rcnn",

tests/backend/test_tensorflow_backend.py

@@ -1,4 +1,5 @@
 import numpy
+import keras.layers
 import keras.backend
 import keras_rcnn.backend
 
@@ -34,6 +35,17 @@ def test_inside_image():
     assert all_inside_anchors.shape == (84, 4)
 
 
+def test_propose():
+    rpn_scores = keras.backend.variable(numpy.random.random((1, 7, 7, 9)))
+    rpn_boxes = keras.backend.variable(numpy.random.random((1, 7, 7, 36)))
+    proposals = keras_rcnn.backend.propose(rpn_boxes, rpn_scores, 300)
+    assert keras.backend.eval(proposals).shape[0] == 1
+    assert keras.backend.eval(proposals).shape[-1] == 4
+
+    proposals = keras_rcnn.backend.propose(rpn_boxes, rpn_scores, 10)
+    assert keras.backend.eval(proposals).shape == (1, 10, 4)
+
+
 def test_overlap():
     x = numpy.asarray([
         [0, 10, 0, 10],

tests/test_models.py

@@ -0,0 +1,9 @@
+import keras.layers
+import keras_rcnn.models
+
+
+def test_resnet50_rcnn():
+    inputs = keras.layers.Input((224, 224, 3))
+    model = keras_rcnn.models.ResNet50RCNN(inputs, 21, 300)
+    model.compile(loss=["mse", "mse", "mse"],
+                  optimizer="adam")