122 lines
2.8 KiB
Python

import cv2
import numpy as np
import matplotlib.pyplot as plt
from implement import MeanShift, CamShift, selectObject, calcHist
if True:
capture = cv2.VideoCapture(0)
flip = True
else:
capture = cv2.VideoCapture("colors_Trim.mp4")
flip = False
track_window = None
global xLZ
global yLZ
global xLast
global yLast
global frame
xLZ = 0
yLZ = 0
def clickOnImage(event, x, y, flags, param):
global frame, xLZ, yLZ, xLast, yLast
if event == cv2.EVENT_LBUTTONUP:
#print("EVENT_LBUTTONUP ({},{})".format(x,y))
xLZ = 0
yLZ = 0
if event == cv2.EVENT_LBUTTONDOWN:
#print("EVENT_LBUTTONDOWN ({},{})".format(x,y))
xLZ = x
yLZ = y
if event == cv2.EVENT_MOUSEMOVE:
#print("EVENT_MOUSEMOVE ({},{})".format(x,y))
xLast = x
yLast = y
cv2.namedWindow("Slika")
cv2.setMouseCallback("Slika", clickOnImage)
target_hist = None
while(1):
ret, frame = capture.read()
if ret == True:
if flip:
frame = cv2.flip(frame, 1)
frameCopy = frame.copy()
if xLZ != 0 or yLZ != 0:
cv2.rectangle(frameCopy,(xLZ,yLZ),(xLast,yLast),(0,255,0),2)
cv2.imshow('Slika',frameCopy)
if cv2.waitKey(100) & 0xFF == ord("q"):
miniFrame = frame[yLZ:yLast,xLZ:xLast]
miniFrame = cv2.cvtColor(miniFrame, cv2.COLOR_BGR2HSV)
target_hist = calcHist(miniFrame,2)
cv2.destroyAllWindows()
break
if target_hist is None:
print("No target selected")
exit(1)
track_window = selectObject(frame, target_hist)
cv2.rectangle(frame, (track_window[0], track_window[1]), (track_window[0]+track_window[2], track_window[1]+track_window[3]), (0, 255, 0), 2)
cv2.imshow("Frame", frame)
cv2.waitKey(0)
cv2.destroyAllWindows()
firstFrame = None
while True:
ret, frame = capture.read()
if not ret:
break
if flip:
frame = cv2.flip(frame, 1)
if firstFrame is None:
firstFrame = frame.copy()
frameCopy = frame.copy()
# prevFrame is not none
if firstFrame is not None:
diff = cv2.absdiff(firstFrame, frame)
gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (5, 5), 0)
tres = cv2.threshold(blur, 25, 255, cv2.THRESH_BINARY)[1]
ret, track_window = CamShift(tres, track_window, (3, 10, 1))
cv2.rectangle(frameCopy, (track_window[0], track_window[1]), (track_window[0]+track_window[2], track_window[1]+track_window[3]), (0, 255, 0), 2)
cv2.imshow("Frame", frameCopy)
cv2.imshow("ALG", tres)
if cv2.waitKey(10) & 0xFF == ord("q"):
capture.release()
cv2.destroyAllWindows()
break
cv2.destroyAllWindows()