widzenie-komputerowe-projekt/rybki.ipynb at master

2023-02-02 01:22:08 +01:00

19 KiB

Raw Permalink Blame History

import cv2
import matplotlib.pyplot as plt
import keras
import numpy as np
import threading
from tqdm import tqdm
import torch
from ultralytics import YOLO
import tensorflow as tf

def object_movie_detection(movie_name, background_substract_method, substractor_threshold, mask_parameters, movie_resize, detection_area,object_area_to_detect):
    if background_substract_method == 'KNN':
        object_detector = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=substractor_threshold)
    else :
        object_detector = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=substractor_threshold)
    
    cap = cv2.VideoCapture(movie_name)
    #cap.set(cv2.CAP_PROP_FPS, 60)
    frames = []
    pbar = tqdm(total = cap.get(cv2.CAP_PROP_FRAME_COUNT))
    while True:
        _, frame = cap.read()
        if(frame is not None):
            pbar.update(1)
            detection_field = frame[detection_area[0][0]: detection_area[0][1],detection_area[1][0]:detection_area[1][1]]
            mask = object_detector.apply(detection_field)
            _, mask = cv2.threshold(mask,mask_parameters[0],mask_parameters[1], cv2.THRESH_BINARY)
            conturs, _ =cv2.findContours(mask, cv2.RETR_LIST , cv2.CHAIN_APPROX_TC89_L1)

            images = []
            counter = 0
            for cnt in conturs:
                area = cv2.contourArea(cnt)
                if area > object_area_to_detect:
                    counter += 1
                    x,y,w,h = cv2.boundingRect(cnt)
                    rectangle = cv2.rectangle(detection_field,(x,y),(x+w,y+h),(0,255,0),3)
            rectangle = cv2.rectangle(frame,(330,100),(1900,900),(255,0,0),3)
            cv2.putText(rectangle,f'DETECTION AREA    DETECTED OBJECTS:  {counter}',(330,80),cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 4)
            frames.append(frame)
            frame = cv2.resize(frame, (movie_resize[0], movie_resize[1])) 
            #cv2.imshow("roi", frame)

            key = cv2.waitKey(30)
            if key == 27:
                cv2.destroyAllWindows()
                pbar.close()
                return frames
        else:
            cv2.destroyAllWindows()
            pbar.close()
            return frames

def save_video(video_name, frames, frame_sizes,fps):
    out = cv2.VideoWriter(f'{video_name}.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (frame_sizes[1], frame_sizes[0]))
    for frame in frames:
        out.write(frame)
    out.release()

movie_detection_1 = object_movie_detection(
    '../examples_movie/rybki.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|██████████| 660/660.0 [01:12<00:00,  9.14it/s]

save_video('saved_movies/example_detection_1',movie_detection_1,movie_detection_1[0].shape,60)

movie_detection_2 = object_movie_detection(
    '../examples_movie/rybki2.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|██████████| 1037/1037.0 [03:21<00:00,  5.15it/s]

save_video('saved_movies/example_detection_2',movie_detection_2,movie_detection_2[0].shape,60)

movie_detection_3 = object_movie_detection(
    '../examples_movie/rybki3.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|█████████▉| 962/963.0 [01:20<00:00, 11.98it/s]

save_video('saved_movies/example_detection_3',movie_detection_3,movie_detection_3[0].shape,60)

movie_detection_4 = object_movie_detection(
    '../examples_movie/rybki4.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|█████████▉| 857/858.0 [01:55<00:00,  7.43it/s]

save_video('saved_movies/example_detection_4',movie_detection_4,movie_detection_4[0].shape,60)

movie_detection_5 = object_movie_detection(
    '../examples_movie/rybki5.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|█████████▉| 779/780.0 [01:01<00:00, 12.64it/s]

save_video('saved_movies/example_detection_5',movie_detection_5,movie_detection_5[0].shape,60)

movie_detection_6 = object_movie_detection(
    '../examples_movie/jelly2.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|██████████| 388/388.0 [00:26<00:00, 14.55it/s]

save_video('saved_movies/example_detection_6',movie_detection_6,movie_detection_6[0].shape,60)

movie_detection_7 = object_movie_detection(
    '../examples_movie/jelly3.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

100%|██████████| 388/388.0 [00:37<00:00, 10.38it/s]

save_video('saved_movies/example_detection_7',movie_detection_7,movie_detection_7[0].shape,60)

frames = object_movie_detection(
    '../examples_movie/rybki.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(100,900),(330,1900)],
    object_area_to_detect=1000
)

def object_movie_yolo_detection(movie_name,model,detection_time):

    cap = cv2.VideoCapture(movie_name)
    cap.set(cv2.CAP_PROP_FPS, 60)

    frames = []
    pbar = tqdm(total = detection_time * 60)
    i = 0

    while True:
        i += 1
        ret, frame = cap.read()
        if(frame is not None):
            pbar.update(1)
            model.predict(frame,save_crop=True)
            frames.append(frame)
            key = cv2.waitKey(30)
            if key == 27:
                cv2.destroyAllWindows()
                return frames
            if i >= detection_time * 60:
                pbar.close
                cv2.destroyAllWindows()
                return frames
        else:
            cv2.destroyAllWindows()
            pbar.close
            return frames

model = YOLO("../yolo_model/best.pt")

jelly2 = object_movie_yolo_detection('../examples_movie/jelly2.mp4',model)
jelly3 = object_movie_yolo_detection('../examples_movie/jelly3.mp4',model)

save_video('saved_movies/yolo_jelly1', jelly2,jelly2[0].shape,60)

save_video('saved_movies/yolo_jelly2', jelly3,jelly3[0].shape,60)

yolo_detect_1 = object_movie_yolo_detection('../examples_movie/rybki3.mp4',model,1)

save_video('saved_movies/example_detection_8', yolo_detect_1,yolo_detect_1[0].shape,60)

def object_detection_with_neural_pred(
        movie_name, 
        background_substract_method, 
        substractor_threshold, 
        mask_parameters, 
        movie_resize, 
        detection_area,
        object_area_to_detect,
        model,
        class_names,
        time_to_capture):
    if background_substract_method == 'KNN':
        object_detector = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=substractor_threshold)
    else :
        object_detector = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=substractor_threshold)
    
    cap = cv2.VideoCapture(movie_name)
    cap.set(cv2.CAP_PROP_FPS, 60)

    pbar = tqdm(total = time_to_capture * 60)
    i = 0
    frames = []
    while True:
        ret, frame = cap.read()
        if(frame is not None):
            pbar.update(1)
            i += 1
            detection_field = frame[detection_area[0][0]: detection_area[0][1],detection_area[1][0]:detection_area[1][1]]
            mask = object_detector.apply(detection_field)
            _, mask = cv2.threshold(mask,mask_parameters[0],mask_parameters[1], cv2.THRESH_BINARY)
            conturs, _ =cv2.findContours(mask, cv2.RETR_LIST , cv2.CHAIN_APPROX_TC89_L1)

            images = []
            counter = 0
            for cnt in conturs:
                area = cv2.contourArea(cnt)
                if area > object_area_to_detect:
                    counter += 1
                    x,y,w,h = cv2.boundingRect(cnt)
                    rectangle = cv2.rectangle(detection_field,(x,y),(x+w,y+h),(0,255,0),3)
                    image_to_predict = detection_field[y:y+h,x:x+w]
                    image_to_predict = cv2.resize(image_to_predict,(224,224))
                    images.append((x,y,rectangle,image_to_predict))
            if images:
                for i,image in enumerate(images):
                    pred = model(x=tf.convert_to_tensor(image[3][None, :], dtype='float32'))
                    labels = [class_names[np.argmax(pre)] for pre in pred]
                    cv2.putText(image[2], labels[0], (image[0], image[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)
            rectangle = cv2.rectangle(frame,(detection_area[1][0],detection_area[0][0]),(detection_area[1][1],detection_area[0][1]),(255,0,0),3)
            cv2.putText(rectangle,f'DETECTION AREA    DETECTED OBJECTS:  {counter}',(330,80),cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 4)
            frames.append(frame)
            if i >= time_to_capture * 60:
                cv2.destroyAllWindows()
                pbar.close()
                return frames
            #frame = cv2.resize(frame, (movie_resize[0], movie_resize[1])) 
            #cv2.imshow("roi", frame)

            key = cv2.waitKey(30)
            if key == 27:
                cv2.destroyAllWindows()
                pbar.close()
                return frames
        else:
            cv2.destroyAllWindows()
            pbar.close()
            return frames

def wrap_frozen_graph(graph_def, inputs, outputs, print_graph=False):
    def _imports_graph_def():
        tf.compat.v1.import_graph_def(graph_def, name="")

    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
    import_graph = wrapped_import.graph

    if print_graph == True:
        print("-" * 50)
        print("Frozen model layers: ")
        layers = [op.name for op in import_graph.get_operations()]
        for layer in layers:
            print(layer)
        print("-" * 50)

    return wrapped_import.prune(
        tf.nest.map_structure(import_graph.as_graph_element, inputs),
        tf.nest.map_structure(import_graph.as_graph_element, outputs))

with tf.io.gfile.GFile("../frozen_models/frozen_graph_best_vgg.pb", "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    loaded = graph_def.ParseFromString(f.read())

frozen_func = wrap_frozen_graph(graph_def=graph_def,
                                inputs=["x:0"],
                                outputs=["Identity:0"],
                                print_graph=False)

class_names=sorted(['Fish', "Jellyfish", 'Lionfish', 'Shark', 'Stingray', 'Turtle'])

frames = object_detection_with_neural_pred(
    '../examples_movie/rybki2.mp4',
    background_substract_method='KNN',
    substractor_threshold=100,
    mask_parameters=[254,255],
    movie_resize=[960,540],
    detection_area=[(500,900),(800,1900)],
    object_area_to_detect=1000,
    model=frozen_func,
    class_names=class_names,
    time_to_capture = 5
)

1037it [27:07,  1.57s/it]

save_video('saved_movies/example_detection_9', frames,frames[0].shape,60)

19 KiB Raw Permalink Blame History

19 KiB

Raw Permalink Blame History