test / app.py
Juned24's picture
Update app.py
6a19f4c verified
Raw History Blame Contribute Delete
2.1 kB
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator
import streamlit as st
import cv2
from vidgear.gears import CamGear
# Membuat Streamlit UI
st.set_page_config(page_title="Ruang Belajar", layout="wide", initial_sidebar_state="expanded")
st.title('Ruang Belajar Deployment YoloV9')
# Load Model YOLOv9
model = YOLO('yolov9c.pt')
# Membuat Function untuk plot bounding boxes pada frames
def plot_boxes(frame, model):
results = model.predict(frame)
annotasi = Annotator(frame)
for r in results:
boxes = r.boxes
for box in boxes:
b = box.xyxy[0] # Box coordinat
c = box.cls # ID Class
annotasi.box_label(b, model.names[int(c)])
return annotasi.result()
# Membuat Function untuk process dan display video
def process_video(data, model, placeholder):
if data == 'Webcam':
camera = cv2.VideoCapture(0) # kode webcam
else: # YouTube link
camera = CamGear(source=data, stream_mode=True, logging=True).start()
while True:
if data == 'Webcam':
ret, frame = camera.read()
if not ret:
break
else: # YouTube link
frame = camera.read()
if frame is None:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = plot_boxes(frame, model)
placeholder.image(frame)
if data == 'Webcam':
camera.release()
# Sidebar
with st.sidebar:
video = st.radio('Pilih data video', ['Webcam', 'YouTube'])
link_youtube = ""
if video == 'YouTube':
link_youtube = st.text_input('Masukan Link Youtube', '')
# Memanggil pemrosesan data video di main area
process_placeholder = st.empty()
# Process video
with st.sidebar:
if st.button('Start'):
if video == 'YouTube' and link_youtube:
process_video(link_youtube, model, process_placeholder)
elif video == 'Webcam':
process_video('Webcam', model, process_placeholder)
with st.sidebar:
st.image("logo.png", use_column_width=True)