handai-detector/examples/ssd_obj_detector.py

"""USAGE
python cv_detector.py \
    --prototxt mobilenet_ssd/MobileNetSSD_deploy.prototxt \
    --model mobilenet_ssd/MobileNetSSD_deploy.caffemodel --video ~/Desktop/5min.mp4
"""
# import the necessary packages
from imutils.video import FPS
import numpy as np
import argparse
import imutils
import dlib
import cv2
from PIL import Image

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-p", "--prototxt", required=True,
    help="path to Caffe 'deploy' prototxt file")
ap.add_argument("-m", "--model", required=True,
    help="path to Caffe pre-trained model")
ap.add_argument("-v", "--video", required=True,
    help="path to input video file")
ap.add_argument("-o", "--output", type=str,
    help="path to optional output video file")
ap.add_argument("-c", "--confidence", type=float, default=0.2,
    help="minimum probability to filter weak detections")
args = vars(ap.parse_args())


# initialize the list of class labels MobileNet SSD was trained to
# detect
CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
    "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
    "dog", "horse", "motorbike", "person", "pottedplant", "sheep",
    "sofa", "train", "tvmonitor"]

# load our serialized model from disk
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])


# initialize the video stream and output video writer
print("[INFO] starting video stream...")
vs = cv2.VideoCapture(args["video"])
_fps = vs.get(cv2.CAP_PROP_FPS)

writer = None


# initialize the list of object trackers and corresponding class
# labels
# trackers = []
labels = []
# start the frames per second throughput estimator
fps = FPS().start()
frame_count = 0

# loop over frames from the video file stream
while True:
    # grab the next frame from the video file
    (grabbed, frame) = vs.read()
    frame_count += 1
    _duration = frame_count / _fps

    # check to see if we have reached the end of the video file
    if frame is None:
        break


    """
    ENTRANCE_1: from hospital
    """

    # 45. 325 == 164, 509



    cropped_frame = frame[325:509, 45:164]
    # cropped_frame.save("test.jpg")
    cv2.imwrite('test.jpg', cropped_frame)
    frame = cropped_frame


    # resize the frame for faster processing and then convert the
    # frame from BGR to RGB ordering (dlib needs RGB ordering)
    # frame = imutils.resize(frame, width=600)
    # rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # grab the frame dimensions and convert the frame to a blob
    (h, w) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(frame, 0.007843, (w, h), 127.5)

    # pass the blob through the network and obtain the detections
    # and predictions
    net.setInput(blob)
    detections = net.forward()


    # loop over the detections
    for i in np.arange(0, detections.shape[2]):
        # extract the confidence (i.e., probability) associated
        # with the prediction
        confidence = detections[0, 0, i, 2]


        # filter out weak detections by requiring a minimum
        # confidence
        if confidence > args["confidence"]:
            # extract the index of the class label from the
            # detections list
            idx = int(detections[0, 0, i, 1])
            label = CLASSES[idx]

            DROP = ["diningtable", "chair", "aeroplane"]
            if label in DROP:
                continue
            # if the class label is not a person, ignore it
            # if CLASSES[idx] != "person":
            #     continue

            # compute the (x, y)-coordinates of the bounding box
            # for the object
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")


            print(f"[{_duration:0.02f}] label: {label} (x,y) = ({startX}, {startY})")

            # construct a dlib rectangle object from the bounding
            # box coordinates and start the correlation tracker
            # t = dlib.correlation_tracker()
            # rect = dlib.rectangle(startX, startY, endX, endY)
            # t.start_track(rgb, rect)

            # update our set of trackers and corresponding class
            # labels
            labels.append(label)

            # grab the corresponding class label for the detection
            # and draw the bounding box
            cv2.rectangle(frame, (startX, startY), (endX, endY),
                (0, 255, 0), 2)
            cv2.putText(frame, label, (startX, startY - 15),
                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 2)

    # show the output frame
    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF

    # if the `q` key was pressed, break from the loop
    if key == ord("q"):
        break

    # update the FPS counter
    # fps.update()



# stop the timer and display FPS information
fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))


# do a bit of cleanup
cv2.destroyAllWindows()
vs.release()