In [522]:
import pickle
import cv2
from IPython.display import display, clear_output
from PIL import Image
from os.path import join, splitext, isfile
from collections import defaultdict
from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets

In [526]:
''' Note - Use this notebook after you have already processed an image or video with one or more detection 
algorithms (dlib, mtcnn, tinyface, frcnn). The original media is required in /media and associated bounding 
box pickle files (ex:"...dlib_bboxes.pickle") are required in /bboxes.
'''

# Enter file name found in /media for image or video that has been processed with one or more detection algorithms
media_filename = "Your_File_Here.jpg"
# dlib=yellow, mtcnn=blue, tinyface=red, frcnn=green
detection_methods = [("dlib",(255,255,0)), ("mtcnn",(0,0,255)), ("tinyface",(255,0,0)),("frcnn",(0,255,0))]

results = list()
# Load pickle output files from detection algorithms
for detection_method in detection_methods:
        bbox_filename = join("/bboxes",media_filename + "." + detection_method[0]+"_bboxes.pickle")
        if isfile(bbox_filename):
            results.append((pickle.load(open(bbox_filename,"rb"))[2],detection_method[1]))
            print("Loaded {0}".format(bbox_filename))

if len(results) == 0:
    print("No detection bounding box files found for {0}".format(media_filename))
    
# Build a dictionary to look up frames in videos and get detected bounding boxes
frame_to_bbox = defaultdict(list)
for result in results:
    frame_box_list, color = result
    for frame_box in frame_box_list:
        frame_num, boxes = frame_box
        frame_to_bbox[frame_num].append((boxes,color))


Loaded bboxes/afghanistan-afghan-people-war-2015-red-cross_0.jpg.dlib_bboxes.pickle
Loaded bboxes/afghanistan-afghan-people-war-2015-red-cross_0.jpg.mtcnn_bboxes.pickle
Loaded bboxes/afghanistan-afghan-people-war-2015-red-cross_0.jpg.tinyface_bboxes.pickle
Loaded bboxes/afghanistan-afghan-people-war-2015-red-cross_0.jpg.frcnn_bboxes.pickle

In [527]:
def show_detections(media_filename, frame_index, frame_to_bbox):
    
    ext = splitext(media_filename)[1]

    # videos
    if ext in ['.avi', '.mov', '.mp4']:            
        cap = cv2.VideoCapture(join("/media",media_filename))        
        this_frame = list(frame_to_bbox.keys())[frame_index]      
        print("Selected Frame: {0}".format(this_frame))
        cap.set(1, this_frame)
        ret, test_image = cap.read()
        test_image = cv2.cvtColor(test_image,cv2.COLOR_BGR2RGB)
        if ret:
            for bbox_list, color in frame_to_bbox[this_frame]:
                for box in bbox_list:
                    t,r,b,l = box
                    cv2.rectangle(test_image,(l,b),(r,t),color, 2)
        else:
            print("Couldn't retrieve frame number {0}".format(this_frame))
        cap.release()

    # images
    elif ext in ['.jpg', '.png', '.jpeg', '.bmp', '.gif']:
        # Load image and convert to RGB
        test_image = cv2.imread(join("/media",media_filename))
        test_image = cv2.cvtColor(test_image,cv2.COLOR_BGR2RGB)
        for result in results:
            frame_box_list, color = result 
            frame_num, boxes = frame_box_list[0]
            for box in boxes:
                t,r,b,l = box
                cv2.rectangle(test_image,(l,b),(r,t),color, 2)

    # unknown extension
    else:
        print("Unknown file extension {0}".format(ext))

    image_box = Image.fromarray(test_image)
    display(image_box)

In [528]:
interact(show_detections,
         media_filename=fixed(media_filename),
         frame_index=widgets.IntSlider(min=0,max=len(frame_to_bbox.keys())-1,step=1,value=0),
         frame_to_bbox=fixed(frame_to_bbox));



In [ ]: