lev1khachatryan / centroid-based_object_tracking Goto Github PK

# import the necessary packages
from pycoral.adapters import common
from pycoral.adapters import detect
from pycoral.utils.dataset import read_label_file
from pycoral.utils.edgetpu import make_interpreter
from imutils.video import VideoStream
from PIL import Image
from centroidtracker import CentroidTracker
import imutils
import time
import cv2
import numpy as np

#the arguments
label_input='/home/atomicpi/Desktop/smartcam/tflite/python/examples/detection/models/coco_labels.txt'
model_input='/home/atomicpi/Desktop/smartcam/tflite/python/examples/detection/models/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite'

confidence=0.2
tracked_labels = [0, 1, 2, 3, 5, 7] #person, bike, car, motorcycle, bus, truck

camera_streams=[]
camera_streams.append('rtsp://aaa:[email protected]/Streaming/channels/102') #frontdoor
camera_streams.append('rtsp://aaa:[email protected]/Streaming/channels/202') #driveway
camera_streams.append('rtsp://aaa:[email protected]/Streaming/channels/302') #rear
camera_streams.append('rtsp://aaa:[email protected]/Streaming/channels/702') #trees
camera_streams.append('rtsp://aaa:[email protected]/Streaming/channels/802') #street

red=(0,0,255) #BGR
green=(0,255,0)
blue=(255,0,0)

# initialize our centroid tracker and frame dimensions
ct = CentroidTracker()

# Initialize the TF interpreter
interpreter = make_interpreter(model_input)
interpreter.allocate_tensors()

#Define ROIs within each video stream, color rectangle based on ROI, notify if inside ROI
#SEE: https://stackoverflow.com/questions/36399381/whats-the-fastest-way-of-checking-if-a-point-is-inside-a-polygon-in-python

ROIcontours=[]
#FRONT DOOR
contour=np.array([
    [0,250],
    [300,70],
    [400,50],
    [704,30],
    [704,480],
    [0,480]],np.int32)
contour=contour.reshape((-1,1,2))
ROIcontours.append(contour)

#DRIVEWAY
contour=np.array([
    [0,355],
    [0,150],
    [110,75],
    [310,70],
    [400,80],
    [600,125],
    [610,355]],np.int32)
contour=contour.reshape((-1,1,2))
ROIcontours.append(contour)
#REAR
contour=np.array([
    [0,480],
    [0,0],
    [704,0],
    [704,480],
    ],np.int32)
contour=contour.reshape((-1,1,2))
ROIcontours.append(contour)

#TREES
contour=np.array([
    [0,480],
    [0,200],
    [260,70],
    [450,90],
    [704,90],
    [704,480]],np.int32)
contour=contour.reshape((-1,1,2))
ROIcontours.append(contour)

#WINDSOR
contour=np.array([
    [550,360],
    [550,200],
    [300,150],
    [550,65],
    [640,75],
    [640,360]],np.int32)
contour=contour.reshape((-1,1,2))
ROIcontours.append(contour)
draw_contours = True

#Open camera streams in new threads
print("[INFO] starting video streams...")
vs = []
for x, streams in enumerate(camera_streams):
    print(f'   Stream {x} started.')
    vs.append(VideoStream(src=camera_streams[x]).start())

#Get first frames
images=[]
for x,stream in enumerate(vs):
    #print(f'reading {x}')
    images.append(stream.read())


# initialize the labels dictionary
print("[INFO] parsing class labels...")
labels = {}
# loop over the class labels file
for row in open(label_input):
    # unpack the row and update the labels dictionary
    (classID, label) = row.strip().split(maxsplit=1)
    labels[int(classID)] = label.strip()

print('Tracking the following labels:')
for label in tracked_labels:
    print(labels.get(label))
# load the Google Coral object detection model
print("[INFO] loading Coral model...")
#model = DetectionEngine(model_input)

#Functions to concat the camera frames for display
def vconcat_resize_min(im_list, interpolation=cv2.INTER_CUBIC):
    w_min = 1280 #min(im.shape[1] for im in im_list)
    im_list_resize = [cv2.resize(im, (w_min, int(im.shape[0] * w_min / im.shape[1])), interpolation=interpolation)
                      for im in im_list]
    return cv2.vconcat(im_list_resize)

def hconcat_resize_min(im_list, interpolation=cv2.INTER_CUBIC):
    h_min = min(im.shape[0] for im in im_list)
    im_list_resize = [cv2.resize(im, (int(im.shape[1] * h_min / im.shape[0]), h_min), interpolation=interpolation)
                      for im in im_list]
    return cv2.hconcat(im_list_resize)

def concat_tile_resize(im_list_2d, interpolation=cv2.INTER_CUBIC): #https://note.nkmk.me/en/python-opencv-hconcat-vconcat-np-tile/
    im_list_v = [hconcat_resize_min(im_list_h, interpolation=interpolation) for im_list_h in im_list_2d]
    return vconcat_resize_min(im_list_v, interpolation=interpolation)

# loop over the frames from the video stream
while True:
    fps = time.time()
    # grab the frame from the threaded video stream and resize it
    for x,stream in enumerate(vs):
        images[x]=stream.read()

        # prepare the frame for object detection by converting (1) it
        # from BGR to RGB channel ordering and then (2) from a NumPy
        # array to PIL image format
        images[x] = cv2.cvtColor(images[x], cv2.COLOR_BGR2RGB)
        images[x] = Image.fromarray(images[x])
        scale = common.set_resized_input(interpreter, images[x].size, lambda size: images[x].resize(size, Image.ANTIALIAS))
        # make predictions on the input frame
        interpreter.invoke()
        results = detect.get_objects(interpreter, confidence, scale[1])
        
        #Convert back to CV2 array
        images[x] = np.array(images[x])
        images[x] = cv2.cvtColor(images[x], cv2.COLOR_BGR2RGB)
        if draw_contours:
            cv2.polylines(images[x],[ROIcontours[x]],isClosed=True,color=blue,thickness=1)
        # loop over the results
        rects = []
        for r in results:
            if r.id in tracked_labels:
                # extract the bounding box and box and predicted class label
                box = r.bbox#.flatten().astype("int")
                rects.append(box)
                #print(box)
                (startX, startY, endX, endY) = box
                box_center = ((endX-startX)/2+startX,(endY-startY)/2+startY)
                label = labels.get(r.id)
                
                #check if midpoint of bounding box is inside ROI
                if cv2.pointPolygonTest(ROIcontours[x],box_center,False) > 0:
                    cv2.rectangle(images[x], (startX, startY), (endX, endY), red, 2)
                    y = startY - 15 if startY - 15 > 15 else startY + 15
                    text = "{}: {:.1f}%".format(label, r.score * 100)
                    cv2.putText(images[x], text, (startX, y),cv2.FONT_HERSHEY_SIMPLEX, 0.5, red, 2)
                # draw the bounding box and label on the image
                else:
                    cv2.rectangle(images[x], (startX, startY), (endX, endY), green, 2)
                    y = startY - 15 if startY - 15 > 15 else startY + 15
                    text = "{}: {:.1f}%".format(label, r.score * 100)
                    cv2.putText(images[x], text, (startX, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, green, 2)
 
        # update our centroid tracker using the computed set of bounding
        # box rectangles
        objects = ct.update(rects)

        # loop over the tracked objects
        for (objectID, centroid) in objects.items():
            # draw both the ID of the object and the centroid of the
            # object on the output frame
            text = "ID {}".format(objectID)
            cv2.putText(images[x], text, (centroid[0] - 10, centroid[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            cv2.circle(images[x], (centroid[0], centroid[1]), 4, (0, 255, 0), -1)
 
    # show the output frame and wait for a key press
    im_tile_resize = concat_tile_resize([[images[0], images[1]],
                                        [images[2], images[3], images[4]]],interpolation=cv2.INTER_CUBIC)
    cv2.putText(im_tile_resize, str(int(1//(time.time()-fps))), (5,12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, red, 2)
    #cv2.putText(im_tile_resize, str(int(model.get_inference_time())), (1,22), cv2.FONT_HERSHEY_SIMPLEX, 0.5, red, 2)
    cv2.imshow("Frame", im_tile_resize)
    #fps.update()
    key = cv2.waitKey(1) & 0xFF
    # if the `q` key was pressed, break from the loop
    if key == ord("q"):
        break
    
# do a bit of cleanup
cv2.destroyAllWindows()
for x, streams in enumerate(camera_streams):
    print(f'stopping {x}')
    vs.append(VideoStream(src=camera_streams[x]).stop())