Implemented RPN minibatch generation

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "C:\\Program Files\\Anaconda3\\python.exe"
+}

rpn_builder.py

@@ -8,10 +8,12 @@ def __init__(self, backbone, scales, ratios):
         # hard-coded parameters (for now)
         self.stride = 32
         self.base_anchor_size = 64
-        self.positive_iou_threshold = 0.7
+        self.positive_iou_threshold = 0.5
         self.negative_iou_threshold = 0.3
         self.batch_size = 256
         self.positives_ratio = 0.5
+        self.minibatch_positives_number = int(self.positives_ratio * self.batch_size)
+        self.minibatch_negatives_number = self.batch_size - self.minibatch_positives_number
         self.max_number_of_predictions = 400
 
         # parameters
@@ -40,11 +42,38 @@ def build_loss(self, ground_truths, predictions):
         # build minibatch
         # apply loss to minibatch
 
-    def generate_minibatch(self):
-        pass
-        anchors = generate_anchors(image.shape)
-        anchors_batch_indices, positive_anchors_indices, negative_anchors_indices = generate_minibatch_mask(anchors, ground_truths)
-
+    '''
+    Generates a random minibatch from positive and negative anchors
+        Args:
+        anchors: tensor of shape (1, height, width, num_anchors, 4)
+        ground_truths: tensor of shape (1, None, 4)
+    Returns:
+        positive_anchor_indices: tensor of shape (1, 3, num_positive_anchors)
+            Note: second dimension has indices for dimensions (height, width, num_anchors)
+        positive_gt_indices: tensor of shape (1, num_positive_anchors)
+        negative_anchor_indices: tensor of shape (1, 3, num_negative_anchors)
+            Note: second dimension has indices for dimensions (height, width, num_anchors)
+    '''
+    @tf.function
+    def generate_minibatch(self, anchors, ground_truths):
+        positive_anchor_indices, positive_gt_indices, negative_anchor_indices = self.assign_anchors_to_ground_truths(anchors, ground_truths)
+        n_positives = tf.minimum(tf.shape(positive_anchor_indices)[2], self.minibatch_positives_number)
+        n_negatives = tf.minimum(tf.shape(negative_anchor_indices)[2], self.batch_size - n_positives)
+
+        indices = tf.range(tf.shape(positive_anchor_indices)[2])
+        indices = tf.random.shuffle(indices)
+        indices = tf.slice(indices, [0], [n_positives])
+
+        positive_anchor_indices = tf.gather(positive_anchor_indices, indices,axis=2)
+        positive_gt_indices = tf.gather(positive_gt_indices, indices,axis=1)
+
+        indices = tf.range(tf.shape(negative_anchor_indices)[2])
+        indices = tf.random.shuffle(indices)
+        indices = tf.slice(indices, [0], [n_negatives])
+        negative_anchor_indices = tf.gather(negative_anchor_indices, indices,axis=2)
+
+        return positive_anchor_indices, positive_gt_indices, negative_anchor_indices
+
 
     '''
     Generates anchor templates, in XYXY format, centered at 0
@@ -76,23 +105,24 @@ def __get_anchor_templates(self):
     @tf.function
     def generate_anchors(self, feature_map):
         # TODO support minibatch by tiling anchors on first dimension
+        feature_map_shape = tf.shape(feature_map)
         assert feature_map.shape[0] == 1
-        vertical_stride = tf.range(0,feature_map.shape[1])
-        vertical_stride = tf.tile(vertical_stride,[feature_map.shape[2]])
-        vertical_stride = tf.reshape(vertical_stride, (feature_map.shape[2], feature_map.shape[1]))
+        vertical_stride = tf.range(0,feature_map_shape[1])
+        vertical_stride = tf.tile(vertical_stride,[feature_map_shape[2]])
+        vertical_stride = tf.reshape(vertical_stride, (feature_map_shape[2], feature_map_shape[1]))
         vertical_stride = tf.transpose(vertical_stride)
 
-        horizontal_stride = tf.range(0,feature_map.shape[2])
-        horizontal_stride = tf.tile(horizontal_stride, [feature_map.shape[1]])
-        horizontal_stride = tf.reshape(horizontal_stride, (feature_map.shape[1], feature_map.shape[2]))
+        horizontal_stride = tf.range(0,feature_map_shape[2])
+        horizontal_stride = tf.tile(horizontal_stride, [feature_map_shape[1]])
+        horizontal_stride = tf.reshape(horizontal_stride, (feature_map_shape[1], feature_map_shape[2]))
 
         centers_xyxy = tf.stack([horizontal_stride, vertical_stride, horizontal_stride, vertical_stride], axis=2)
 
         centers_xyxy = self.stride * centers_xyxy
         centers_xyxy = tf.cast(centers_xyxy,tf.float32)
 
         centers_xyxy = tf.tile(centers_xyxy,[1,1,self.anchor_templates.shape[0]])
-        centers_xyxy = tf.reshape(centers_xyxy, (feature_map.shape[1], feature_map.shape[2], self.anchor_templates.shape[0], 4))
+        centers_xyxy = tf.reshape(centers_xyxy, (feature_map_shape[1], feature_map_shape[2], self.anchor_templates.shape[0], 4))
         anchors = centers_xyxy + self.anchor_templates
         anchors = tf.expand_dims(anchors,axis=0)
         return anchors
@@ -108,9 +138,8 @@ def generate_anchors(self, feature_map):
         positive_gt_indices: tensor of shape (1, num_positive_anchors)
         negative_anchor_indices: tensor of shape (1, 3, num_negative_anchors)
             Note: second dimension has indices for dimensions (height, width, num_anchors)
-
     '''
-    # @tf.function
+    @tf.function
     def assign_anchors_to_ground_truths(self, anchors, ground_truths):
         anchors = tf.cast(anchors, tf.float32)
         ground_truths = tf.cast(ground_truths, tf.float32)

rpn_tests.py

@@ -35,7 +35,7 @@ def test_iou(self):
         assert ret[1,1] == 1
 
     def test_assign_anchors(self):
-        DEBUG = True
+        DEBUG = False
         if DEBUG:
             image = np.ones((500,500,3))
             box_size = 60
@@ -55,13 +55,52 @@ def test_assign_anchors(self):
             positive_anchors = np.squeeze(anchors)[positive_anchor_indices[0],positive_anchor_indices[1],positive_anchor_indices[2]]
             positive_ground_truths = bounding_boxes[positive_ground_truth_indices]
 
+            negative_anchor_indices = np.array(negative_anchor_indices).reshape((3,-1))
+            negative_anchors = np.squeeze(anchors)[negative_anchor_indices[0],negative_anchor_indices[1],negative_anchor_indices[2]]
+
             for anchor in positive_anchors:
-                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (0,0,255),1)
+                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (255,0,0),2)
             for anchor in positive_ground_truths:
                 cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (0,255,0),1)
+            for anchor in negative_anchors:
+                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (0,0,255),1)
+            cv2.imshow('anchors', image)
+            cv2.waitKey(0)
+
+
+    def test_get_minibatch(self):
+        DEBUG = True
+        if DEBUG:
+            image = np.ones((500,500,3))
+            box_size = 60
+            bounding_boxes = np.array([[100,100,100+box_size,100+box_size],[300,300,300+box_size,300+box_size]])
+            for box in bounding_boxes:
+                image[box[1]:box[3],box[0]:box[2]] = 0
 
+            backbone = None
+            scales = [0.5, 1, 2]
+            ratios = [0.5, 1, 2]
+            rpn = RegionProposalNetwork(backbone, scales, ratios)
+            image_feature_map = np.zeros((1, image.shape[0] // rpn.stride, image.shape[1] // rpn.stride, 2048))
+            anchors = rpn.generate_anchors(image_feature_map)
+            positive_anchor_indices, positive_ground_truth_indices, negative_anchor_indices = rpn.generate_minibatch(anchors, np.expand_dims(bounding_boxes,axis=0))
+            positive_anchor_indices = np.array(positive_anchor_indices).reshape((3,-1))
+            positive_ground_truth_indices = np.array(positive_ground_truth_indices).reshape((-1))
+            positive_anchors = np.squeeze(anchors)[positive_anchor_indices[0],positive_anchor_indices[1],positive_anchor_indices[2]]
+            positive_ground_truths = bounding_boxes[positive_ground_truth_indices]
+
+            negative_anchor_indices = np.array(negative_anchor_indices).reshape((3,-1))
+            negative_anchors = np.squeeze(anchors)[negative_anchor_indices[0],negative_anchor_indices[1],negative_anchor_indices[2]]
+
+            for anchor in positive_anchors:
+                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (255,0,0),1)
+            for anchor in positive_ground_truths:
+                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (0,255,0),1)
+            for anchor in negative_anchors:
+                cv2.rectangle(image, (int(anchor[0]), int(anchor[1])), (int(anchor[2]), int(anchor[3])), (0,0,255),1)
             cv2.imshow('anchors', image)
             cv2.waitKey(0)
+
     def test_rpn(self):
         DEBUG = False
         if DEBUG: