smarthome-presence-detect/camera/person_detection.py

#!/usr/bin/env python

import argparse
import time
from statistics import median
import imutils
from imutils.video import VideoStream
import cv2
import numpy as np
import paho.mqtt.client as mqtt
from video_stream import imagezmq

VISUAL_DEBUG = True
BROKER = "141.75.33.126"
PORT = 1883

def getArgs():
    """ Arguments """
    ap = argparse.ArgumentParser()
    ap.add_argument("-v", "--video", help="path to the video file")
    ap.add_argument("-a", "--min-area", type=int, default=500, help="minimum area size")
    return vars(ap.parse_args())


def main():
    try:
        mqtt_client = mqtt.Client("pi-camera")
        mqtt_client.connect(BROKER, PORT)
    except:
        print("Connection to MQTT-Broker failed.")
        return 1

    try:
        args = getArgs()
        timer = Timer()

        # if the video argument is None, then the code will read from webcam (work in progress)
        if args.get("video", None) is None:
            #vs = VideoStream(src=0).start()
            image_hub = imagezmq.ImageHub()
            
            time.sleep(2.0)
        # otherwise, we are reading from a video file
        else:
            vs = cv2.VideoCapture(args["video"])

        cv2.namedWindow('Video stream', cv2.WINDOW_NORMAL)
        detector = DetectionFromFrame(args["min_area"], 0.8)
        while True:
            people_count = 0
            timer.start_frame_timer()
            if args.get("video", None) is None:
                rpi_name, detector.currentFrame = image_hub.recv_image()
                image_hub.send_reply(b'OK')
            else:
                detector.currentFrame = vs.read()
                detector.currentFrame = detector.currentFrame if args.get("video", None) is None else detector.currentFrame[1]
            # if the frame can not be grabbed, then we have reached the end of the video
            if detector.currentFrame is None:
                break

            # resize the frame to 500
            detector.currentFrame = imutils.resize(detector.currentFrame, width=500)
            detector.framecounter += 1
            if detector.framecounter > 1:
                
                cnts = detector.prepareFrame()

                for c in cnts:
                    bound_rect = cv2.boundingRect(c)
                    #(x, y, w, h) = cv2.boundingRect(c)
                    #initBB2 =(x,y,w,h)

                    prott1 = r'ML-Models/MobileNetSSD_deploy.prototxt'
                    prott2 = r'ML-Models/MobileNetSSD_deploy.caffemodel'
                    net = cv2.dnn.readNetFromCaffe(prott1, prott2)

                    #trackbox = detector.currentFrame[y:y+h, x:x+w]boundRect[1]
                    trackbox = detector.currentFrame[bound_rect[1]:bound_rect[1]+bound_rect[3], 
                                                    bound_rect[0]:bound_rect[0]+bound_rect[2]]
                    trackbox = cv2.resize(trackbox, (224, 224))
                    #cv2.imshow('image',trackbox)
                    blob = cv2.dnn.blobFromImage(cv2.resize(trackbox, (300, 300)),0.007843, (300, 300), 127.5)
                    net.setInput(blob)
                    detections = net.forward()
                    
                    for i in np.arange(0, detections.shape[2]):
                        people_count += detector.detectConfidentiallyPeople(i, detections, bound_rect)

                    cv2.rectangle(detector.currentFrame, (bound_rect[0], bound_rect[1]), 
                                    (bound_rect[0] + bound_rect[2], bound_rect[1] + bound_rect[3]), (255, 255, 0), 1)
                
            # show the frame and record if the user presses a key
            cv2.imshow("Video stream", detector.currentFrame)
            key = cv2.waitKey(1) & 0xFF

            # send number of people detected via mqtt
            mqtt_client.publish("/gso/bb/104/Camera", str(people_count))

            # if the `q` key is pressed, break from the lop
            if key == ord("q"):
                break
            if key == ord("d"):
                detector.firstFrame = None
            #detector.lastFrame = detector.currentFrame

            timer.print_frame_time()


        # finally, stop the camera/stream and close any open windows
        if args.get("video", None) is not None:
            vs.stop() if args.get("video", None) is None else vs.release()

        cv2.destroyAllWindows()
    finally:
        if args.get("video", None) is None:
            image_hub.send_reply(b'OK')

class DetectionFromFrame:
    def __init__(self, min_size, confidence):
        self.min_size = min_size
        self.confidence_level = confidence

        self.firstFrame = None
        self.currentFrame = None
        
        self.initBB2 = None
        self.fps = None
        self.differ = None
        self.now = ''
        self.framecounter = 0
        self.people_count_total = 0


    def prepareFrame(self):
        gray = cv2.cvtColor(self.currentFrame, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (21, 21), 0)

        # if the first frame is None, initialize it
        if self.firstFrame is None:
            self.firstFrame = gray
            return []

        # compute the absolute difference between the current frame and first frame
        frameDelta = cv2.absdiff(self.firstFrame, gray)
        thresh = cv2.threshold(frameDelta, 25, 255, cv2.THRESH_BINARY)[1]
        
        #debug
        """if VISUAL_DEBUG: 
            cv2.imshow("debug image", thresh)
            cv2.waitKey(0)
            cv2.destroyWindow("debug image")
            #cv2.destroyWindow("threshhold image")"""
        
        # dilate the thresholded image to fill in holes
        thresh = cv2.dilate(thresh, None, iterations=2)

        # find contours on thresholded image
        thresh = np.uint8(thresh)
        cnts, _  = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        return cnts

    def detectConfidentiallyPeople(self, i, detections, bound_rect):
        #CLASSES = ["person"]

        detected_color = (0, 255, 0)
        #COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))

        confidence = detections[0, 0, i, 2]

        if confidence > self.confidence_level:
            # draw a rectangle in green over the detected area
            cv2.rectangle(self.currentFrame, (bound_rect[0], bound_rect[1]),
                                (bound_rect[0] + bound_rect[2], bound_rect[1] + bound_rect[3]), detected_color, 3)
            
            label = "{:.2f}%".format(confidence * 100)
            cv2.putText(self.currentFrame, label, (bound_rect[0], bound_rect[1]-5), cv2.FONT_HERSHEY_SIMPLEX, 0.3, detected_color, 1)

            return 1
        else:
            return 0

class Timer:
    def __init__(self):
        self.frame_timer = None
        self.contour_timer = None
        self.detection_timer = None

        self.contour_time = []
        self.detection_time = []

    def start_frame_timer(self):
        self.frame_timer = time.time()

    def get_frame_time(self):
        return time.time() - self.frame_timer

    def start_contour_timer(self):
        self.contour_timer = time.time()

    def stop_contour_timer(self):
        self.contour_time.append(time.time() - self.contour_timer)

    def start_detection_timer(self):
        self.detection_timer = time.time()

    def stop_detection_timer(self):
        self.detection_time.append(time.time() - self.detection_timer)

    def print_frame_time(self):
        print("Time for Frame: {:.2f}.".format(self.get_frame_time()))
        

    def print_other_times(self):
        average_contour = 0 if not self.contour_time else sum(self.contour_time)/float(len(self.contour_time))
        average_detection = 0 if not self.detection_time else sum(self.detection_time)/float(len(self.detection_time))

        median_contour = 0 if not self.contour_time else median(self.contour_time)
        median_detection = 0 if not self.detection_time else median(self.detection_time)

        total_contour = sum(self.contour_time)
        total_detection = sum(self.detection_time)

        print("Contour Total: {:.2f}. Contour Median: {:.2f}. Contour Average: {:.2f}.".format(
            total_contour, median_contour, average_contour))
        print("Detection Total: {:.2f}. Detection Median: {:.2f}. Detection Average: {:.2f}. ".format(
            total_detection, median_detection, average_detection))

        self.contour_time = []
        self.detection_time = []


if __name__ == "__main__":
    main()