mask detection bug fix

balaji1994 · Mar 22, 2018 · 490986d · 490986d
1 parent fd4a5c4
commit 490986d
Show file tree

Hide file tree

Showing 17 changed files with 7 additions and 16 deletions.
diff --git a/object_detection.py b/object_detection.py
@@ -281,8 +281,9 @@ def detection(detection_graph, category_index, score, expand):
                 fps.update()
 
     # End everything
-    gpu_worker.stop()
-    cpu_worker.stop()
+    if split_model:
+        gpu_worker.stop()
+        cpu_worker.stop()
     fps.stop()
     video_stream.stop()
     cv2.destroyAllWindows()

diff --git a/object_detection/__init__.pyc b/object_detection/__init__.pyc
diff --git a/object_detection/core/__init__.pyc b/object_detection/core/__init__.pyc
diff --git a/object_detection/core/box_list.pyc b/object_detection/core/box_list.pyc
diff --git a/object_detection/core/box_list_ops.pyc b/object_detection/core/box_list_ops.pyc
diff --git a/object_detection/core/standard_fields.pyc b/object_detection/core/standard_fields.pyc
diff --git a/object_detection/protos/__init__.pyc b/object_detection/protos/__init__.pyc
diff --git a/object_detection/protos/string_int_label_map_pb2.pyc b/object_detection/protos/string_int_label_map_pb2.pyc
diff --git a/object_detection/utils/__init__.pyc b/object_detection/utils/__init__.pyc
diff --git a/object_detection/utils/label_map_util.pyc b/object_detection/utils/label_map_util.pyc
diff --git a/object_detection/utils/ops.pyc b/object_detection/utils/ops.pyc
diff --git a/object_detection/utils/shape_utils.pyc b/object_detection/utils/shape_utils.pyc
diff --git a/object_detection/utils/static_shape.pyc b/object_detection/utils/static_shape.pyc
diff --git a/object_detection/utils/visualization_utils.pyc b/object_detection/utils/visualization_utils.pyc
diff --git a/object_detection_mask_test.py b/object_detection_mask_test.py
@@ -48,7 +48,6 @@
 num_classes         = cfg['num_classes']
 split_model         = cfg['split_model']
 log_device          = cfg['log_device']
-convert_rgb         = cfg['convert_rgb']
 ssd_shape           = cfg['ssd_shape']
 
 
@@ -220,21 +219,13 @@ def detection(detection_graph, category_index, score, expand):
             cur_frames = 0
             while video_stream.isActive():
                 image = video_stream.read()
-                fps.update()
-                # read video frame and expand dimensions
-                if convert_rgb:
-                    try:
-                        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-                        cvt = True
-                    except:
-                        print("Error converting BGR2RGB")
-                        cvt = False    
+                image_expanded = np.expand_dims(cv2.cvtColor(image, cv2.COLOR_BGR2RGB),0)
                 # detection
                 if split_model:
-                    (score, expand) = sess.run([score_out, expand_out], feed_dict={image_tensor: np.expand_dims(image, 0)})
+                    (score, expand) = sess.run([score_out, expand_out], feed_dict={image_tensor:image_expanded})
                     output_dict = sess.run(tensor_dict, feed_dict={score_in:score, expand_in: expand})
                 else:
-                    output_dict = sess.run(tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)})
+                    output_dict = sess.run(tensor_dict, feed_dict={image_tensor: image_expanded})
                 #num = int(output_dict['num_detections'][0])
                 classes = output_dict['detection_classes'][0].astype(np.uint8)
                 boxes = output_dict['detection_boxes'][0]
@@ -243,8 +234,6 @@ def detection(detection_graph, category_index, score, expand):
                     output_dict['detection_masks'] = output_dict['detection_masks'][0]
                 # Visualization of the results of a detection.
                 if visualize:
-                    if convert_rgb and cvt:
-                        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                     vis_util.visualize_boxes_and_labels_on_image_array(
                         image,
                         boxes,
@@ -270,6 +259,7 @@ def detection(detection_graph, category_index, score, expand):
                                 print("label: {}\nscore: {}\nbox: {}".format(label, score, box))
                     if cur_frames >= max_frames:
                         break
+                fps.update()
     # End everything
     fps.stop()
     video_stream.stop()     

diff --git a/stuff/__init__.pyc b/stuff/__init__.pyc
diff --git a/stuff/helper.pyc b/stuff/helper.pyc