adding changes by LandyGuo

model.py

@@ -768,7 +768,7 @@ def refine_detections_graph(rois, probs, deltas, window, config):
     # Shape: [boxes, (y1, x1, y2, x2)] in normalized coordinates
     refined_rois = apply_box_deltas_graph(
         rois, deltas_specific * config.BBOX_STD_DEV)
-    
+
     # Convert coordiates to image domain
     # TODO: better to keep them normalized until later
     height, width = config.IMAGE_SHAPE[:2]
@@ -782,72 +782,80 @@ def refine_detections_graph(rois, probs, deltas, window, config):
 
     # Filter out background boxes
     keep = tf.where(class_ids > 0)[:,0]
+
     # Filter out low confidence boxes
+    conf_keep = tf.where(class_scores >= config.DETECTION_MIN_CONFIDENCE)[:,0]
+
     if config.DETECTION_MIN_CONFIDENCE:
-        keep = tf.sets.set_intersection(
-                keep, tf.where(class_scores >= config.DETECTION_MIN_CONFIDENCE))[:,0]
+    keep = tf.sparse_tensor_to_dense(tf.sets.set_intersection(
+            tf.expand_dims(keep, 0), tf.expand_dims(conf_keep, 0)))[0]
 
     # Apply per-class NMS
     pre_nms_class_ids = tf.gather(class_ids, keep)
     pre_nms_scores = tf.gather(class_scores, keep)
     pre_nms_rois = tf.gather(refined_rois,   keep)
-    print('pre_nms_class_ids = {}'.format(pre_nms_class_ids.shape))
-    print('pre_nms_scores = {}'.format(pre_nms_scores.shape))
-    print('pre_nms_rois = {}'.format(pre_nms_rois.shape))
 
     uniq_pre_nms_class_ids = tf.unique(pre_nms_class_ids)[0]
 
+    # sort unique class ids
+    _,max_index = tf.nn.top_k(-uniq_pre_nms_class_ids, tf.size(uniq_pre_nms_class_ids))
+    uniq_pre_nms_class_ids = tf.gather(uniq_pre_nms_class_ids,max_index)
+
     nms_keep = []
-    def nms_keep_map(class_id):
-        #class_id = tf.expand_dims(class_id, -1)
-        print('pre_nms_class_ids.shape', pre_nms_class_ids.shape)
-        print('class_id', class_id.shape)
-        ixs = tf.where(pre_nms_class_ids == class_id)
+    def nms_keep_map(i, ret):
+
+        class_id = uniq_pre_nms_class_ids[i]
+        scale = tf.fill(tf.shape(pre_nms_class_ids), class_id)
+        ixs = tf.cast(tf.where(tf.equal(scale, pre_nms_class_ids))[:,0], tf.int32)
 
-        print('ixs {}'.format(ixs.shape))
         # Apply NMS
         class_keep = tf.image.non_max_suppression(
                 tf.to_float(tf.gather(pre_nms_rois,ixs)),
                 tf.gather(pre_nms_scores, ixs),
-                max_output_size=ixs.shape[1],
+                max_output_size=tf.shape(ixs)[0],
                 iou_threshold=config.DETECTION_NMS_THRESHOLD)
 
         # Map indicies
-        return tf.gather(keep, tf.gather(ixs, class_keep))
-    
-    print('uniq_pre_nms_class_ids: {}'.format(uniq_pre_nms_class_ids.shape))
-    nms_keep = tf.map_fn(nms_keep_map, uniq_pre_nms_class_ids)
-    nms_keep = tf.concat(nms_keep, axis=0) 
-    nms_keep = tf.unique(nms_keep)[0] 
-    nms_keep = tf.to_int64(nms_keep)
-    keep = tf.expand_dims(keep, 0)
-    nms_keep = tf.expand_dims(nms_keep, 0)
-    keep = tf.sparse_tensor_to_dense(tf.sets.set_intersection(keep, nms_keep))[:,0]
-    """#tf.to_int32(
-    #np.intersect1d(keep, nms_keep).astype(np.int32)
-    """
+        cur_keep_indexes = tf.gather(keep, tf.gather(ixs, class_keep))
+        return i+1, tf.concat([ret,cur_keep_indexes], axis=0)
+
+    nums_iters = tf.shape(uniq_pre_nms_class_ids)[0] # unique class ids
+    i = tf.constant(0)
+    ret = tf.ones([1], dtype=tf.int32)
+    c = lambda i, unique_pre_nms:tf.less(i, nums_iters)
+    b = nms_keep_map
+    r = tf.while_loop(c, b, [i, -ret],
+    	shape_invariants=[i.get_shape(), tf.TensorShape([None])])
+
+    nms_keep = r[1]
+
+    # remove initial_value background
+    nms_keep = tf.gather(nms_keep, tf.where(nms_keep >= 0)[:,0])
+
+    keep = tf.sparse_tensor_to_dense(tf.sets.set_intersection(tf.expand_dims(keep,0), tf.expand_dims(nms_keep,0)))[0]
+    keep = tf.cast(keep, tf.int32)
+
     # Keep top detections
     roi_count = tf.convert_to_tensor(config.DETECTION_MAX_INSTANCES)
-    print('class scores: {}'.format(class_scores.shape))
     class_scores_keep = tf.gather(class_scores, keep)
     num_keep = tf.minimum(tf.shape(class_scores_keep)[0], roi_count)
     top_ids = tf.nn.top_k(class_scores_keep, k=num_keep, sorted=True)[1]
-    
+
     #np.argsort(class_scores[keep])[::-1][:roi_count]
     print('keep before: {}'.format(keep.shape))
     keep = tf.gather(keep, top_ids)
     print('keep after: {}'.format(keep.shape))
- 
+
     refined_rois_keep = tf.gather(tf.to_float(refined_rois), keep)
-    class_ids_keep = tf.gather(tf.to_float(class_ids), keep)[..., tf.newaxis]  
-    class_scores_keep = tf.gather(class_scores, keep)[..., tf.newaxis] 
+    class_ids_keep = tf.gather(tf.to_float(class_ids), keep)[..., tf.newaxis]
+    class_scores_keep = tf.gather(class_scores, keep)[..., tf.newaxis]
     print('refined_rois_keep = ', refined_rois_keep.shape)
     print('class_ids_keep = ', class_ids_keep.shape)
     print('class_scores_keep = ', class_scores_keep.shape)
 
     # Arrange output as [N, (y1, x1, y2, x2, class_id, score)]
     # Coordinates are in image domain.
-    detections = tf.concat((refined_rois_keep, class_ids_keep, 
+    detections = tf.concat((refined_rois_keep, class_ids_keep,
                 class_scores_keep), axis=1)
     print('detections.shape = ', detections.shape)