File size: 2,437 Bytes
9fc2e73 |
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 |
import streamlit as st
import torch
import cv2
import tempfile
import os
from PIL import Image
import numpy as np
# Muat model YOLOv5 yang sudah dilatih
@st.cache_resource
def load_model():
model = torch.hub.load('ultralytics/yolov11', 'custom', path='best.pt', force_reload=True)
return model
model = load_model()
# Fungsi untuk deteksi pada gambar
def detect_image(image):
results = model(image)
results.render() # Tambahkan bounding box ke gambar
annotated_image = results.imgs[0] # Ambil gambar yang sudah dianotasi
return annotated_image
# Fungsi untuk deteksi pada video
def detect_video(video_path, output_path):
cap = cv2.VideoCapture(video_path)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame)
results.render()
annotated_frame = np.squeeze(results.render())
out.write(annotated_frame)
cap.release()
out.release()
# Streamlit Interface
st.title("YOLO Object Detection")
# Opsi untuk memilih gambar atau video
option = st.radio("Pilih jenis input:", ("Gambar", "Video"))
if option == "Gambar":
uploaded_image = st.file_uploader("Unggah gambar", type=["jpg", "jpeg", "png"])
if uploaded_image is not None:
image = Image.open(uploaded_image)
st.image(image, caption="Gambar asli", use_column_width=True)
# Deteksi objek
annotated_image = detect_image(np.array(image))
st.image(annotated_image, caption="Hasil deteksi", use_column_width=True)
elif option == "Video":
uploaded_video = st.file_uploader("Unggah video", type=["mp4", "avi", "mov"])
if uploaded_video is not None:
# Simpan video sementara
temp_video_path = tempfile.NamedTemporaryFile(delete=False).name
with open(temp_video_path, "wb") as f:
f.write(uploaded_video.read())
# Proses video
output_video_path = "output_video.mp4"
st.text("Sedang memproses video...")
detect_video(temp_video_path, output_video_path)
st.text("Proses selesai!")
# Tampilkan hasil video
st.video(output_video_path)
|