light_person_yolo_webcam_v2.py

#!/usr/bin/env python3

# USAGE
# python light_person_yolo_webcam_v2.py --yolo yolo-coco
# python light_person_yolo_webcam_v2.py --use-gpu 1 --yolo yolo-coco
# ex: set tabstop=8 softtabstop=0 expandtab shiftwidth=2 smarttab:

# import the necessary packages
from imutils.video import VideoStream
from imutils.video import FPS
import numpy as np
import argparse
import imutils
import time
import cv2
import os
import json

import socket
import sys
import threading

from pprint import pprint as pp

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
#ap.add_argument("-i", "--input", required=True,
#  help="path to input video")
#ap.add_argument("-o", "--output", required=True,
#  help="path to output video")
ap.add_argument("-y", "--yolo", required=True,
  help="base path to YOLO directory")
ap.add_argument("-c", "--confidence", type=float, default=0.5,
  help="minimum probability to filter weak detections")
ap.add_argument("-t", "--threshold", type=float, default=0.3,
  help="threshold when applyong non-maxima suppression")
ap.add_argument("-u", "--use-gpu", type=int, default=0,
	help="boolean indicating if CUDA GPU should be used")
ap.add_argument("-co", "--codec", type=int, default=0,
	help="chose codec for compression, 0 = MJEG, 1=H264")
args = vars(ap.parse_args())



################## Server Module ##################
connections = []

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

server_address = ('0.0.0.0', 10000)

sock.bind(server_address)

sock.listen(10)

def accept_conn():
  while True:
    connection, client_address = sock.accept()
    connections.append({'conn': connection, 'addr': client_address})

t = threading.Thread(target=accept_conn)

t.start()

###################################################


# load the COCO class labels our YOLO model was trained on
labelsPath = os.path.sep.join([args["yolo"], "coco.names"])
LABELS = open(labelsPath).read().strip().split("\n")

# initialize a list of colors to represent each possible class label
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3),
  dtype="uint8")

# derive the paths to the YOLO weights and model configuration
weightsPath = os.path.sep.join([args["yolo"], "yolov3.weights"])
configPath = os.path.sep.join([args["yolo"], "yolov3.cfg"])

# load our YOLO object detector trained on COCO dataset (80 classes)
# and determine only the *output* layer names that we need from YOLO
print("[INFO] loading YOLO from disk...")
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
print(args["use_gpu"])
if args["use_gpu"]:
	# set CUDA as the preferable backend and target
	print("[INFO] setting preferable backend and target to CUDA...")
	net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
	net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# initialize the video stream, pointer to output video file, and
# frame dimensions
#vs = cv2.VideoCapture(args["input"])
print("[INFO] starting video stream...")
vs = VideoStream(src=3).start()
time.sleep(2.0)
#writer = None
(W, H) = (None, None)

# try to determine the total number of frames in the video file
#try:
#  prop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() \
#    else cv2.CAP_PROP_FRAME_COUNT
#  total = int(vs.get(prop))
#  print("[INFO] {} total frames in video".format(total))

# an error occurred while trying to determine the total
# number of frames in the video file
#except:
#  print("[INFO] could not determine # of frames in video")
#  print("[INFO] no approx. completion time can be provided")
#  total = -1

# loop over frames from the video file stream
while True:
  # read the next frame from the file
  #(grabbed, frame) = vs.read()
  frame = vs.read()
  #frame = imutils.resize(frame, width=400)
  frame = imutils.resize(frame, width=1200)
  # if the frame was not grabbed, then we have reached the end
  # of the stream
  #if not grabbed:
  #  break

  # if the frame dimensions are empty, grab them
  if W is None or H is None:
    (H, W) = frame.shape[:2]

  # construct a blob from the input frame and then perform a forward
  # pass of the YOLO object detector, giving us our bounding boxes
  # and associated probabilities
  blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416),
    swapRB=True, crop=False)
  net.setInput(blob)
  start = time.time()
  layerOutputs = net.forward(ln)
  end = time.time()

  # initialize our lists of detected bounding boxes, confidences,
  # and class IDs, respectively
  boxes = []
  confidences = []
  classIDs = []
  send_data = []

  # loop over each of the layer outputs
  for output in layerOutputs:
    # loop over each of the detections
    for detection in output:
      # extract the class ID and confidence (i.e., probability)
      # of the current object detection
      scores = detection[5:]
      classID = np.argmax(scores)
      confidence = scores[classID]

      # filter out weak predictions by ensuring the detected
      # probability is greater than the minimum probability
      if confidence > args["confidence"]: # and classID in [0]:
        # scale the bounding box coordinates back relative to
        # the size of the image, keeping in mind that YOLO
        # actually returns the center (x, y)-coordinates of
        # the bounding box followed by the boxes' width and
        # height
        box = detection[0:4] * np.array([W, H, W, H])
        (centerX, centerY, width, height) = box.astype("int")

        # use the center (x, y)-coordinates to derive the top
        # and and left corner of the bounding box
        x = int(centerX - (width / 2))
        y = int(centerY - (height / 2))

        # update our list of bounding box coordinates,
        # confidences, and class IDs
        boxes.append([x, y, int(width), int(height)])
        confidences.append(float(confidence))
        classIDs.append(classID)

  # apply non-maxima suppression to suppress weak, overlapping
  # bounding boxes
  idxs = cv2.dnn.NMSBoxes(boxes, confidences, args["confidence"],
    args["threshold"])

  # ensure at least one detection exists

  if len(idxs) > 0:
    # loop over the indexes we are keeping
    for i in idxs.flatten():
      send_data.append({
          'x': int(boxes[i][0]),
          'y': int(boxes[i][1]),
          'width': int(boxes[i][2]),
          'height': int(boxes[i][3]),
          'confidences': float(confidences[i]),
          'class_id': int(classIDs[i]),
          'label': LABELS[classIDs[i]]
      })
      # extract the bounding box coordinates
      (x, y) = (boxes[i][0], boxes[i][1])
      (w, h) = (boxes[i][2], boxes[i][3])

      # draw a bounding box rectangle and label on the frame
      color = [int(c) for c in COLORS[classIDs[i]]]
      cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
      text = "{}: {:.4f}".format(LABELS[classIDs[i]],
        confidences[i])
      cv2.putText(frame, text, (x, y - 5),
        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

  pp(send_data)

  ################## Server Module ##################
  tmp = json.dumps(send_data).encode('utf-8')
  for i in range(len(connections)):
    try:
      connections[i]['conn'].sendall(tmp)
    except BrokenPipeError as e:
      connections.pop(i)
      break
    except ConnectionResetError as e:
      connections.pop(i)
      break

  ###################################################


  # check if the video writer is None
  #if writer is None:
    # initialize our video writer
  #  fourcc = cv2.VideoWriter_fourcc(*"MJPG")
  #  writer = cv2.VideoWriter(args["output"], fourcc, 30,
  #    (frame.shape[1], frame.shape[0]), True)

    # some information on processing single frame
  #  if total > 0:
  #    elap = (end - start)
  #    print("[INFO] single frame took {:.4f} seconds".format(elap))
  #    print("[INFO] estimated total time to finish: {:.4f}".format(
  #      elap * total))

  # write the output frame to disk
  #writer.write(frame)

    # 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
# release the file pointers
print("[INFO] cleaning up...")
#writer.release()
vs.stop()

cv2.destroyAllWindows()
vs.release()