bug fix no viz

Author: gustavz

config.sample.yml

@@ -18,8 +18,8 @@ WIDTH: 600                     # OpenCV only supports 4:3 formats others will be
 HEIGHT: 600                    # 600x600 leads to 640x480
 MAX_FRAMES: 5000               # only used if visualize==False
 FPS_INTERVAL: 5                # Interval [s] to print fps of the last interval in console
-DET_INTERVAL: 500              # intervall [frames] to print detections to console
-DET_TH: 0.5                    # detection threshold for det_intervall
+PRINT_INTERVAL: 500            # intervall [frames] to print detections to console
+PRINT_TH: 0.5                  # detection threshold for det_intervall
 ## speed hack
 SPLIT_MODEL: True              # Splits Model into a GPU and CPU session (currently only works for ssd_mobilenets)
 SSD_SHAPE: 300                 # used for the split model algorithm (currently only supports ssd networks trained on 300x300 and 600x600 input)

rod/config.py

@@ -41,8 +41,8 @@ class Config(object):
     HEIGHT = cfg['HEIGHT']                  # 600x600 leads to 640x480
     MAX_FRAMES = cfg['MAX_FRAMES']          # only used if visualize==False
     FPS_INTERVAL = cfg['FPS_INTERVAL']      # Interval [s] to print fps of the last interval in console
-    DET_INTERVAL = cfg['DET_INTERVAL']      # intervall [frames] to print detections to console
-    DET_TH = cfg['DET_TH']                  # detection threshold for det_intervall
+    PRINT_INTERVAL = cfg['PRINT_INTERVAL']  # intervall [frames] to print detections to console
+    PRINT_TH = cfg['PRINT_TH']              # detection threshold for det_intervall
     ## speed hack
     SPLIT_MODEL = cfg['SPLIT_MODEL']        # Splits Model into a GPU and CPU session (currently only works for ssd_mobilenets)
     SSD_SHAPE = cfg['SSD_SHAPE']            # used for the split model algorithm (currently only supports ssd networks trained on 300x300 and 600x600 input)

rod/vis_utils.py

@@ -315,12 +315,56 @@ def draw_text_on_image(image, string, position, color = (77, 255, 9)):
         cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2)
 
 
-def visualize_objectdetection(image, boxes, classes, scores, masks, category_index,
-                fps=None, visualize=False, det_interval=5, det_th=0.5,
-                max_frames=500, cur_frame = None, model='realtime_object_detection'):
+def exit_visualization(millis):
     """
-    complicated visualization function for object detections
-    TODO: CLEAN UP
+    returns false if openCV exit is requested
+    """
+    if cv2.waitKey(millis) & 0xFF == ord('q'):
+        return False
+    return True
+
+def exit_print(cur_frame,max_frames):
+    """
+    returns false if max frames are reached
+    """
+    if cur_frame >= max_frames:
+        return False
+    return True
+
+def print_detection(boxes,scores,classes,category_index,cur_frame,max_frames=500,print_interval=100,print_th=0.5):
+    """
+    prints detection result above threshold to console
+    """
+    for box, score, _class in zip(boxes, scores, classes):
+        if cur_frame%print_interval==0 and score > print_th:
+            label = category_index[_class]['name']
+            print("label: {}\nscore: {}\nbox: {}".format(label, score, box))
+
+def draw_single_box_on_image(image,box,label):
+    """
+    draws single box and label on image
+    """
+    p1 = (box[1], box[0])
+    p2 = (box[3], box[2])
+    cv2.rectangle(image, p1, p2, (77,255,9), 2)
+    draw_text_on_image(image,label,(p1[0],p1[1]-10))
+
+
+def visualize_objectdetection(image,
+                            boxes,
+                            classes,
+                            scores,
+                            masks,
+                            category_index,
+                            cur_frame,
+                            max_frames=500,
+                            fps='N/A',
+                            print_interval=100,
+                            print_th=0.5,
+                            model_name='realtime_object_detection',
+                            visualize=True):
+    """
+    visualization function for object_detection
     """
     if visualize:
         visualize_boxes_and_labels_on_image(
@@ -332,44 +376,31 @@ def visualize_objectdetection(image, boxes, classes, scores, masks, category_ind
         instance_masks=masks,
         use_normalized_coordinates=True,
         line_thickness=2)
-        if fps:
-            draw_text_on_image(image,"fps: {}".format(fps), (10,30))
-        cv2.imshow(model, image)
-    elif not visualize and cur_frame:
-        # Exit after max frames if no visualization
-        for box, score, _class in zip(boxes, scores, classes):
-            if cur_frame%det_interval==0 and score > det_th:
-                label = category_index[_class]['name']
-                print("label: {}\nscore: {}\nbox: {}".format(label, score, box))
-    elif fps == "console":
-        for box, score, _class in zip(boxes, scores, classes):
-            if score > det_th:
-                label = category_index[_class]['name']
-                print("label: {}\nscore: {}\nbox: {}".format(label, score, box))
-    # Exit Option
-    if visualize:
-        if cv2.waitKey(1) & 0xFF == ord('q'):
-            return False
-    elif not visualize and fps:
-        if cur_frame >= max_frames:
-            return False
-    return True
-
-def draw_single_box_on_image(image,box,label):
-    p1 = (box[1], box[0])
-    p2 = (box[3], box[2])
-    cv2.rectangle(image, p1, p2, (77,255,9), 2)
-    draw_text_on_image(image,label,(p1[0],p1[1]-10))
-
+        draw_text_on_image(image,"fps: {}".format(fps), (10,30))
+        cv2.imshow(model_name, image)
+        vis = exit_visualization(1)
+    else:
+        print_detection(boxes,scores,classes,category_index,cur_frame,max_frames,print_interval,print_th)
+        vis = exit_print(cur_frame,max_frames)
+    return vis
 
-def visualize_deeplab(image,seg_map,model_name,fps,visualize=True):
+def visualize_deeplab(image,
+                    seg_map,
+                    cur_frame,
+                    max_frames=500,
+                    fps='N/A',
+                    print_interval=100,
+                    print_th=0.5,
+                    model_name='DeepLab',
+                    visualize=True):
+    """
+    visualization function for deeplab
+    """
     if visualize:
         draw_mask_on_image(image, seg_map)
         draw_text_on_image(image,"fps: {}".format(fps),(10,30))
         cv2.imshow(model_name,image)
-        if cv2.waitKey(1) & 0xFF == ord('q'):
-            return False
-        else:
-            return True
+        vis = exit_visualization(1)
     else:
-        return True
+        vis = exit_print(cur_frame,max_frames)
+    return vis

run_deeplab.py

@@ -46,8 +46,8 @@ def segmentation(model,config):
                         if config.VISUALIZE:
                             draw_single_box_on_image(frame,box,label)
 
-                vis = visualize_deeplab(frame,seg_map,config.OD_MODEL_NAME+config._DEV+config._OPT,
-                                        fps.fps_local(),config.VISUALIZE)
+                vis = visualize_deeplab(frame,seg_map,fps._glob_numFrames,config.MAX_FRAMES,fps.fps_local(),
+                                        config.PRINT_INTERVAL,config.PRINT_TH,config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                 if not vis:
                     break
                 fps.update()

run_objectdetection.py

@@ -120,9 +120,9 @@ def detection(model,config):
                     scores = np.squeeze(scores)
 
                     # Visualization
-                    vis = visualize_objectdetection(frame, boxes, classes, scores, masks, category_index, fps.fps_local(),
-                                        config.VISUALIZE, config.DET_INTERVAL, config.DET_TH, config.MAX_FRAMES,
-                                        fps._glob_numFrames, config.OD_MODEL_NAME)
+                    vis = visualize_objectdetection(frame,boxes,classes,scores,masks,category_index,fps._glob_numFrames,
+                                                    config.MAX_FRAMES,fps.fps_local(),config.PRINT_INTERVAL,config.PRINT_TH,
+                                                    config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                     if not vis:
                         break
 
@@ -146,9 +146,9 @@ def detection(model,config):
                     for idx,tracker in enumerate(trackers):
                         tracker_box = tracker.update(frame)
                         tracker_boxes[idx,:] = conv_track2detect(tracker_box, vs.real_width, vs.real_height)
-                    vis = visualize_objectdetection(frame, tracker_boxes, classes, scores, masks, category_index, fps.fps_local(),
-                                        config.VISUALIZE, config.DET_INTERVAL, config.DET_TH, config.MAX_FRAMES,
-                                        fps._glob_numFrames, config.OD_MODEL_NAME+config._DEV+config._OPT)
+                    vis = visualize_objectdetection(frame,tracker_boxes,classes,scores,masks,category_index,fps._glob_numFrames,
+                                                    config.MAX_FRAMES,fps.fps_local(),config.PRINT_INTERVAL,config.PRINT_TH,
+                                                    config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                     if not vis:
                         break
 
@@ -168,8 +168,11 @@ def detection(model,config):
         cpu_worker.stop()
 
 
-if __name__ == '__main__':
+def main():
     config = Config()
     model = Model('od',config.OD_MODEL_NAME,config.OD_MODEL_PATH,config.LABEL_PATH,
                 config.NUM_CLASSES,config.SPLIT_MODEL, config.SSD_SHAPE).prepare_od_model()
     detection(model, config)
+
+if __name__ == '__main__':
+    main()

test_deeplab.py

@@ -63,8 +63,8 @@ def segmentation(model,config):
                         if config.VISUALIZE:
                             draw_single_box_on_image(frame,box,label)
 
-                vis = visualize_deeplab(frame,seg_map,config.OD_MODEL_NAME+config._DEV+config._OPT,
-                                        timer.get_fps(),config.VISUALIZE)
+                vis = visualize_deeplab(frame,seg_map,timer.get_frame(),config.MAX_FRAMES,timer.get_fps(),
+                                        config.PRINT_INTERVAL,config.PRINT_TH,config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                 if not vis:
                     break
 	cv2.destroyAllWindows()

test_objectdetection.py

@@ -14,8 +14,8 @@
 from rod.helper import Timer, WebcamVideoStream, SessionWorker, TimeLiner, load_images
 from rod.model import Model
 from rod.config import Config
-from rod.utils import ops as utils_ops
 from rod.vis_utils import visualize_objectdetection
+from rod.tf_utils import reframe_box_masks_to_image_masks
 
 
 def detection(model,config):
@@ -38,7 +38,7 @@ def detection(model,config):
                 real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                 detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                 detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
-                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
+                detection_masks_reframed = reframe_box_masks_to_image_masks(
                         detection_masks, detection_boxes, config.HEIGHT, config.WIDTH)
                 detection_masks_reframed = tf.cast(tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                 # Follow the convention by adding back the batch dimension
@@ -120,9 +120,9 @@ def detection(model,config):
                 scores = np.squeeze(scores)
 
                 # Visualization
-                vis = visualize_objectdetection(frame, boxes, classes, scores, masks, category_index, timer.get_fps(),
-                                    config.VISUALIZE, config.DET_INTERVAL, config.DET_TH, config.MAX_FRAMES, None,
-                                    config.OD_MODEL_NAME+config._DEV+config._OPT)
+                vis = visualize_objectdetection(frame,boxes,classes,scores,masks,category_index,timer.get_frame(),
+                                                config.MAX_FRAMES,timer.get_fps(),config.PRINT_INTERVAL,config.PRINT_TH,
+                                                config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                 if not vis:
                     break