widzenie-komputerowe-projekt/rybki.ipynb

import cv2
import matplotlib.pyplot as plt
import keras
import numpy as np
import threading

model = keras.models.load_model('./model')

cap = cv2.VideoCapture("rybki.mp4")
cap.set(cv2.CAP_PROP_FPS, 60)

class_names=sorted(['Fish', "Jellyfish", 'Penguin', 'Puffin', 'Shark', 'Starfish', 'Stingray'])
object_detector = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=50)


width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter('track_fish.avi', cv2.VideoWriter_fourcc(*'MJPG'), fps, (width, height))

while True:
    ret, frame = cap.read()
    if(frame is not None):
        roi = frame[100: 900,330:1900]
        mask = object_detector.apply(roi)
        _, mask = cv2.threshold(mask,254,255, cv2.THRESH_BINARY)
        conturs, _ =cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        images = []

        for cnt in conturs:
            area = cv2.contourArea(cnt)
            if area > 200:
                #cv2.drawContours(roi,[cnt],-1,(0,255,0),2)
                x,y,w,h = cv2.boundingRect(cnt)
                rectangle = cv2.rectangle(roi,(x,y),(x+w,y+h),(0,255,0),3)
                image_to_predict = roi[y:y+h,x:x+w]
                image_to_predict = cv2.resize(image_to_predict,(227,227))
                images.append((x,y,rectangle,np.expand_dims(image_to_predict,axis=0)))
                
                # pred = model.predict(np.expand_dims(image_to_predict, axis=0))
                # label = class_names[np.argmax(pred)]
        if images:
            pred = model.predict(np.vstack([image[3] for image in images]))
            labels = [class_names[np.argmax(pre)] for pre in pred]
            for i,image in enumerate(images):
                cv2.putText(image[2], labels[i], (image[0], image[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)
        roi = cv2.resize(roi, (960, 540)) 
        cv2.imshow("roi", roi)

        key = cv2.waitKey(30)
        if key == 27:
            break

        #out.write(frame)
    else:
        break


#out.release()
cap.release()
cv2.destroyAllWindows()