Spaces:
Runtime error
Runtime error
File size: 2,632 Bytes
8e8f78f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import cv2
from ultralytics import YOLO
import numpy as np
import PIL
import streamlit as st
import io
class ObjectDetectionModel():
def __init__(self):
self.model = YOLO("models/yolov8n_openvino_model", task = "detect")
self.class_names = self.model.names
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 cls in result[0].boxes.cls:
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)
frame_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = camera.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'X264')
output_video_path = 'output_video.mp4'
out = cv2.VideoWriter(output_video_path,fourcc,fps, (frame_width,frame_height))
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()
processed_frames.append(img_plot)
st_frame.image(img_plot, channels = "BGR")
frame_count +=1
progress_bar.progress(frame_count/total_frames, text = None)
camera.release()
for frame in processed_frames:
out.write(frame)
out.release()
st_frame.empty()
progress_bar.empty()
return output_video_path
|