projekt_PP/main.py

import cv2
import numpy as np
import vehicles
import time

counter_up=0    #licznik pojazdow jadacych w gore
counter_down=0  #licznik pojazdow jadacych w dol

while True:
    user_choice = input("Wybierz nagranie:\n1. video_1.mp4\n2. video_2.avi\n")
    if user_choice == "1":
        input_video = 'data/video_1.avi'
        vid_name = 'video_1'
        kernelSize = 6
        break
    elif user_choice == "2":
        input_video = 'data/video_2.mp4'
        vid_name = 'video_2'
        kernelSize = 3
        break

cap = cv2.VideoCapture(input_video)

height, width, frames_count, fps = cap.get(cv2.CAP_PROP_FRAME_HEIGHT), cap.get(
    cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_COUNT), cap.get(cv2.CAP_PROP_FPS),
height = int(height)
width = int(width)
frameArea = height*width

print("Wysokosc", height, "szerokosc", width, "liczba klatek", frames_count, "fps",fps)
font = cv2.FONT_HERSHEY_SIMPLEX
"""
#Polozenie linii "mety" (na srodku kadru)
lineUpper = int(4*(height/10))
lineLower = int(6*(height/10))

upLimit = int(2*(height/10))
downLimit = int(8*(height/10))
"""
#Polozenie linii "mety" (w dolnej czesci kadru)
lineUpper = int(7*(height/10))
lineLower = int(8*(height/10))

upLimit = int(4*(height/10))
downLimit = int(10*(height/10))

print("Zielona linia y:",str(lineLower))
print("Czerwona linia y:",str(lineUpper))
lineLower_color = (0,255,0)
lineUpper_color = (0,0,255)
point1 =  [0, lineLower]
point2 =  [width, lineLower]
points_lineLower = np.array([point1,point2], np.int32)
# Przekształcenie tablicy do 1x2, polecane w dokumentacji OpenCV ale chyba zbedne w tym przypadku
#points_lineLower = points_lineLower.reshape((-1,1,2))
point3 =  [0, lineUpper]
point4 =  [width, lineUpper]
points_LineUpper = np.array([point3,point4], np.int32)
#points_LineUpper = points_LineUpper.reshape((-1,1,2))
point5 =  [0, upLimit]
point6 =  [width, upLimit]
points_upLimit = np.array([point5,point6], np.int32)
#points_upLimit = points_upLimit.reshape((-1,1,2))
point7 =  [0, downLimit]
point8 =  [width, downLimit]
points_downLimit= np.array([point7,point8], np.int32)
#points_downLimit = points_downLimit.reshape((-1,1,2))

sub=cv2.createBackgroundSubtractorMOG2()    # utworz background subtractor

size = 2.0  #wielkosc obrazu
cars = []
maxAllowedAge = 3 #przez ile klatek "zgubiony" pojazd bedzie pozostawal na liscie do "sledzenia"
id = 1


while(cap.isOpened()):
    ret, frame = cap.read()
    for i in cars:
        i.age_one()

    if ret == True: #przetwarzanie kadru

        #cv2.imshow("bez zmian", frame)  # wyswietlanie filmu wejsciowego bez zmian
        #gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        #cv2.imshow("czarnobialy", gray)  # wyswietlanie filmu w skali szarosci
        fgMask = sub.apply(frame)  # uzycie background subtraction
        #fgMask2 = sub.apply(frame)
        #cv2.imshow("fgMask", fgMask)
        # operacje morfologiczne. Wg tutoriala na docs.opencv.org. Opening, closing - usuwanie szumów
        #Binarization
        ret,imBinary = cv2.threshold(fgMask,200,255,cv2.THRESH_BINARY)
        #cv2.imshow("binarny", imBinary)

        kernelOpening = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernelSize, kernelSize))  # kernels to apply to the morphology
        kernelClosing = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
        #Opening is just another name of erosion followed by dilation.
        #It is useful in removing noise
        mask = cv2.morphologyEx(imBinary,cv2.MORPH_OPEN,kernelOpening)
        #cv2.imshow("opening", mask)

        #Closing is reverse of Opening, Dilation followed by Erosion.
        #It is useful in closing small holes inside the foreground objects, or small black points on the object.
        mask = cv2.morphologyEx(mask,cv2.MORPH_CLOSE,kernelClosing)
        #cv2.imshow("closing", mask)

        # kontury (obrysowywanie)
        contours,hierarchy=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)   # RETR_EXTERNAL - bierz pod uwage najbardziej 'zewnetrzne' kontury

        # kontury zajmujace zbyt duza lub mala czesc kadru beda ignorowane
        maxContourSize = frameArea / 4
        minContourSize = frameArea / 400

        for j in range(len(contours)):   # przejdz wszystkie kontury w kadrze
            area = cv2.contourArea(contours[j])
            #print("wielkosc konturu:",area)
            if minContourSize < area < maxContourSize:  # zbyt duze i male obiekty sa ignorowane
                ####        Sledzenie konturow      ######
                centroid = contours[j]
                moments = cv2.moments(centroid)
                cx = int(moments['m10'] / moments['m00'])
                cy = int(moments['m01'] / moments['m00'])

                # punkty graniczne (bounding points) konturu. x,y to wspolrzedne lewego gornego rogu, w,h to szerokosc i wysokosc prostokata
                x, y, w, h = cv2.boundingRect(centroid)

                new = True    #przyjmij ze kontur to nowy pojazd
                if cy in range(upLimit,downLimit):    # sprawdz czy jest wart sledzenia (czyli miedzy szarymi liniami limitujacymi)
                    for i in cars:
                        # jesli odnajdziesz sledzony kontur o bardzo zblizonych koordynatach,
                        #przymij ze to ten sam i zaktualizuj jego polozenie
                        if abs(x - i.getX()) <= w and abs(y - i.getY()) <= h:   #abs function calculates an absolute value of each matrix element.
                            new = False
                            i.updateCoords(cx, cy)
                            #sprawdz czy prekroczyl mete
                            if i.going_UP(lineLower,lineUpper) == True:
                                counter_up+=1
                                print("Pojazd o ID:",i.getId(),'przekroczyl linie mety w gore o czasie:', time.strftime("%c"))
                            elif i.going_DOWN(lineLower,lineUpper) == True:
                                counter_down+=1
                                print("Pojazd o ID:", i.getId(), 'przekroczyl linie mety w dol o czasie:', time.strftime("%c"))
                            break
                            #jesli kierunek ruchu jest znany i linia mety zostala przekroczona, przestan sledzic i usun z listy
                        if i.getState() == '1':
                            if i.getDir() == 'down'and i.getY() > downLimit:
                                i.setDone()
                            elif i.getDir() == 'up'and i.getY() < upLimit:
                                i.setDone()
                        if i.timedOut():
                            index=cars.index(i)
                            cars.pop(index)
                            del i

                    if new == True: #Jesli kontur nie byl dotychczas sledzony, dodaj nowy samochod do listy
                        newCar=vehicles.Car(id,cx,cy,maxAllowedAge)
                        cars.append(newCar)
                        id+=1

                    cv2.rectangle(frame,(x,y),(x+w,y+h),(255,0,0),2)

        #for i in cars:
            #cv2.putText(frame, str(i.getId()), (i.getX(), i.getY()), font, 0.3, i.getRGB(), 1, cv2.LINE_AA)    #wypisywanie id




        str_up='W GORE: '+str(counter_up)
        str_down='W DOL: '+str(counter_down)
        frame = cv2.polylines(frame,[points_lineLower],False,lineLower_color,thickness=2)
        frame = cv2.polylines(frame,[points_LineUpper],False,lineUpper_color,thickness=2)
        frame = cv2.polylines(frame,[points_upLimit],False,(255,255,255),thickness=1)
        frame = cv2.polylines(frame,[points_downLimit],False,(255,255,255),thickness=1)
        cv2.putText(frame, str_up, (10, 40), font, 0.4, (255, 255, 255), 2, cv2.LINE_AA)    #biale tlo zeby licznik byl wyrazniejszy
        cv2.putText(frame, str_up, (10, 40), font, 0.4, lineUpper_color, 1, cv2.LINE_AA)
        cv2.putText(frame, str_down, (10, 90), font, 0.4, (255, 255, 255), 2, cv2.LINE_AA)  #biale tlo zeby licznik byl wyrazniejszy
        cv2.putText(frame, str_down, (10, 90), font, 0.4, lineLower_color, 1, cv2.LINE_AA)
        frame = cv2.resize(frame, (0, 0), None, size, size)  # rozmiar obrazu x2
        cv2.imshow('Frame',frame)

        key = cv2.waitKey(60)
        if key == 27:  # Wyjdz po nacisnieciu escape
            break

    else:
        break
timestr = time.strftime("%Y_%m_%d-%H_%M_%S")
f = open("traffic logs/"+vid_name+" "+timestr+".txt","w+")
f.write("Natężenie ruchu:\r\nliczba samochodów jadących w górę: %d\rliczba samochodów jadących w dół: %d" %(counter_up, counter_down))
cap.release()
cv2.destroyAllWindows()