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))
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 [ ]: