adding changes by LandyGuo

model.py

@@ -782,53 +782,61 @@ 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]