from darkflow.net.build import TFNet
import numpy as np
import cv2
options = {"model": "cfg/yolo.cfg",
"load": "bin/yolo.weights",
"threshold": 0.1,
"cpu": 1.0}
tfnet = TFNet(options)
def boxing(original_img, predictions):
newImage = np.copy(original_img)
for result in predictions:
top_x = result['topleft']['x']
top_y = result['topleft']['y']
btm_x = result['bottomright']['x']
btm_y = result['bottomright']['y']
confidence = result['confidence']
label = result['label'] + " " + str(round(confidence, 3))
if confidence > 0.3:
newImage = cv2.rectangle(newImage, (top_x, top_y), (btm_x, btm_y), (255,0,0), 3)
newImage = cv2.putText(newImage, label, (top_x, top_y-5), cv2.FONT_HERSHEY_COMPLEX_SMALL , 0.8, (0, 230, 0), 1, cv2.LINE_AA)
return newImage
cap = cv2.VideoCapture('./sample_video/input.mp4')
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
out = cv2.VideoWriter('./sample_video/output.avi',fourcc, 20.0, (int(width), int(height)))
while(True):
ret, frame = cap.read()
if ret == True:
frame = np.asarray(frame)
results = tfnet.return_predict(frame)
new_frame = boxing(frame, results)
out.write(new_frame)
cv2.imshow('frame',new_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
cap.release()
out.release()
cv2.destroyAllWindows()