smarthome-presence-detect/camera/counter_people.py

from imutils.object_detection import non_max_suppression
import numpy as np
import imutils
import cv2
import time
import argparse
import time
import base64

'''
Usage:

python peopleCounter.py -i PATH_TO_IMAGE  # Reads and detect people in a single local stored image
python peopleCounter.py -c  # Attempts to detect people using webcam
'''

HOGCV = cv2.HOGDescriptor()
HOGCV.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

def detector(image):
    '''
    @image is a numpy array
    '''

    clone = image.copy()

    (rects, weights) = HOGCV.detectMultiScale(image, winStride=(4, 4), padding=(8, 8), scale=1.05)

    # draw the original bounding boxes
    for (x, y, w, h) in rects:
        cv2.rectangle(clone, (x, y), (x + w, y + h), (0, 0, 255), 2)

    # Applies non-max supression from imutils package to kick-off overlapped
    # boxes
    rects = np.array([[x, y, x + w, y + h] for (x, y, w, h) in rects])
    result = non_max_suppression(rects, probs=None, overlapThresh=0.65)

    return result


def buildPayload(variable, value, context):
    return {variable: {"value": value, "context": context}}


def argsParser():
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image", default=None, help="path to image test file directory")
    ap.add_argument("-c", "--camera", default=False, help="Set as true if you wish to use the camera")
    ap.add_argument("-v", "--video", default=None, help="path to the video file")
    args = vars(ap.parse_args())

    return args


def localDetect(image_path):
    result = []
    image = cv2.imread(image_path)
    image = imutils.resize(image, width=min(400, image.shape[1]))
    clone = image.copy()
    if len(image) <= 0:
        print("[ERROR] could not read local image")
        return result
    print("[INFO] Detecting people")
    result = detector(image)

    """# shows the result
    for (xA, yA, xB, yB) in result:
        cv2.rectangle(image, (xA, yA), (xB, yB), (0, 255, 0), 2)

    cv2.imshow("result", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    cv2.imwrite("result.png", np.hstack((clone, image)))"""
    return result#(result, image)


def cameraDetect(video_path="", sample_time=5):

    if video_path:
        cap = cv2.VideoCapture(video_path)
    else:
        cap = cv2.VideoCapture(0)
    
    #init = time.time()

    while(True):
        # Capture frame-by-frame
        _, frame = cap.read()

        if frame is None:
            break

        frame = imutils.resize(frame, width=min(400, frame.shape[1]))
        result = detector(frame.copy())

        # shows the result
        for (xA, yA, xB, yB) in result:
            cv2.rectangle(frame, (xA, yA), (xB, yB), (0, 255, 0), 2)
        cv2.imshow('frame', frame)
        cv2.waitKey(0)

        #if time.time() - init >= sample_time:
        if len(result):
            print("{} people detected.".format(len(result)))
        #init = time.time()
        

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()


def convert_to_base64(image):
    image = imutils.resize(image, width=400)
    img_str = cv2.imencode('.png', image)[1].tostring()
    b64 = base64.b64encode(img_str)

    return b64.decode('utf-8')


def detectPeople(args):
    image_path = args["image"]
    video_path = args["video"]
    camera = True if str(args["camera"]) == 'true' else False

    # Routine to read local image
    if image_path != None and not camera and video_path == None:
        print("[INFO] Image path provided, attempting to read image")
        (result, image) = localDetect(image_path)
        print(str(len(result)) + " People detected.")

    if video_path != None and not camera:
        print("[INFO] reading video")
        cameraDetect(video_path)

    # Routine to read images from webcam
    if camera:
        print("[INFO] reading camera images")
        cameraDetect()


def main():
    args = argsParser()
    detectPeople(args)


if __name__ == '__main__':
    main()