95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
import argparse
|
|
|
|
class Cascade():
|
|
|
|
def __init__(self,data_file):
|
|
"""[summary]
|
|
|
|
Args:
|
|
data_file (str): path to the cascade file
|
|
"""
|
|
|
|
self.cc = cv2.CascadeClassifier()
|
|
self.cc.load(data_file)
|
|
self.objects = []
|
|
self.scale = 1.1
|
|
self.neighbor = 0
|
|
|
|
def set_parameters(self, scale=1.1, neighbor=0):
|
|
"""
|
|
Set detection parameters
|
|
Args:
|
|
scale (float, optional): Scale factor. Defaults to 1.1.
|
|
neighbor (int, optional): Number of neighbors for detections. Defaults to 0.
|
|
"""
|
|
self.scale = scale
|
|
self.neighbor = neighbor
|
|
|
|
def get_detections(self, img):
|
|
"""
|
|
makes cascade detections
|
|
|
|
Args:
|
|
img (image): input - RGB image
|
|
returns a list of (x,y,w,h) for objects detected
|
|
"""
|
|
img_g = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
|
|
img_g = cv2.equalizeHist(img_g)
|
|
|
|
detections = self.cc.detectMultiScale(img_g,self.scale,self.neighbor, flags=cv2.CASCADE_DO_CANNY_PRUNING)
|
|
|
|
return detections
|
|
|
|
def display(self, img):
|
|
"""
|
|
display on an image the detections
|
|
|
|
Args:
|
|
img (image): color RGB image
|
|
"""
|
|
detections = self.get_detections(img)
|
|
|
|
for (x,y,w,h) in detections:
|
|
cv2.rectangle(img,(x,y),(x+w,y+h),[0,255,0],2)
|
|
|
|
return img
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# input arguments
|
|
parser = argparse.ArgumentParser(description='Cascade tester.\n Defaults: -cam 0 -n 0 -s 1.1')
|
|
parser.add_argument('-cam', type=int, help='Camera ID', default=0)
|
|
parser.add_argument('-n', type=int, help='Number of neighbors for detections', default=0)
|
|
parser.add_argument('-s', type=float, help='Scale factor', default=1.1)
|
|
|
|
args = parser.parse_args()
|
|
|
|
path, filename = os.path.split(os.path.realpath(__file__))
|
|
|
|
f = path + '/data/cascade.xml'
|
|
cap = cv2.VideoCapture(args.cam)
|
|
|
|
cd = Cascade(f)
|
|
cd.set_parameters(args.s,args.n)
|
|
|
|
if not cap.isOpened:
|
|
exit(0)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
|
|
if frame is None:
|
|
break
|
|
|
|
if cv2.waitKey(10) == 27:
|
|
break
|
|
|
|
frame = cd.display(frame)
|
|
cv2.imshow('Capture', frame)
|