utils/object_detection_brainai.py

import cv2
from ultralytics import YOLO
import numpy as np
import PIL
import streamlit as st
import io
import tempfile
import moviepy.editor as mpy
import random

class ObjectDetectionModel():
    def __init__(self):
        self.model = YOLO("models/yolov8n_openvino_model", task = "detect")
        self.class_names = self.model.names

        with open("utils/game_classes.txt", "r") as file:
            self.game_classes = [line.strip() for line in file if line.strip()]

    def process(self, img):
        if isinstance(img, np.ndarray):
            uploaded_img_cv = img
        else:
            uploaded_img = PIL.Image.open(img)
            uploaded_img_cv = np.array(uploaded_img)
            if uploaded_img_cv.shape[-1] == 4:
                uploaded_img_cv = cv2.cvtColor(uploaded_img_cv, cv2.COLOR_RGBA2RGB)
        
        result = self.model(uploaded_img_cv)
        img_plot = result[0].plot()

        detected_classes = set()
        for box in result[0].boxes:
            class_id = int(box.cls[0])
            class_name = self.class_names[class_id]
            detected_classes.add(class_name)
        detected_objects = f'Objects Detected: {", ".join(detected_classes) if detected_classes else "No objects detected"}'

        return img_plot, detected_objects
    
    def play_video(self, video_path):
        uploaded_video = io.BytesIO(video_path.read())
        temporary_location = "upload.mp4"
        
        with open(temporary_location, "wb") as temp_out:
            temp_out.write(uploaded_video.read())
        temp_out.close()
        
        camera = cv2.VideoCapture(temporary_location)
        fps = camera.get(cv2.CAP_PROP_FPS)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        processed_frames = []

        total_frames = int(camera.get(cv2.CAP_PROP_FRAME_COUNT))
        frame_count = 0
        progress_bar = st.progress(0)
        st_frame = st.empty()
        
        while(True):
            ret, frame = camera.read()
            if not ret:
                break
    
            result = self.model(frame, verbose=False)
            img_plot = result[0].plot()
            img_plot_rgb = cv2.cvtColor(img_plot, cv2.COLOR_BGR2RGB)
            processed_frames.append(img_plot_rgb)

            st_frame.image(img_plot, channels = "BGR")
            frame_count +=1
            progress_bar.progress(frame_count/total_frames, text = None)
            
        camera.release()
        clip = mpy.ImageSequenceClip(processed_frames, fps=fps)
        clip.write_videofile(temp_file.name)
    
        st_frame.empty()
        progress_bar.empty()
    
        return temp_file.name

    def call_class(self):
        random_class = random.choice(list(self.game_classes))
        
        return random_class