import streamlit as st
from ultralytics import YOLO
import cv2
from PIL import Image
from io import BytesIO
import cairosvg

# Read the SVG logo file
with open("marca-cor-fundo-escuro.svg", "r") as f:
    logo_svg = f.read()

# Convert SVG to PNG using cairosvg
logo_png_bytes = cairosvg.svg2png(bytestring=logo_svg.encode())
logo_png = Image.open(BytesIO(logo_png_bytes))

# Display the logo
st.image(logo_png, width=200)

# Load the YOLOv8 model
model = YOLO('yolov8n.pt')

# Set up the Streamlit app
#st.title('YOLOv8 Video/Webcam Inference')

# Add file uploader or webcam option
video_file = st.file_uploader("Upload a video", type=["mp4", "avi"])
use_webcam = st.checkbox("Use webcam")

# Placeholder for inference results
inference_placeholder = st.empty()

# Video/webcam inference loop
if video_file is not None or use_webcam:
    if video_file is not None:
        video = cv2.VideoCapture(video_file.name)
    else:
        video = cv2.VideoCapture(1)  # Use webcam

    while True:
        ret, frame = video.read()
        if not ret:
            break

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Run YOLOv8 inference on the frame
        results = model(frame)

        # Display the inference results
        inference_placeholder.image(results[0].plot(), use_container_width=True)

    video.release()
else:
    st.warning("Please upload a video file or select the webcam option.")