import streamlit as st
import cv2
import matplotlib.pyplot as plt
from inference import Stream
import supervision as sv
import os

import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Access the environment variable
api_key = os.getenv('ROBOFLOW_API_KEY')

# Now you can use api_key in your code



annotator = sv.BoxAnnotator()

def main():
    st.title("RTSP Stream with Hugging Face Streamlit")

    # Set up camera connection details
    ip_address = '65.108.95.124'  # IP address of the camera
    port = '554'  # Port number for the RTSP stream
    username = ''  # Username (if required by the camera)
    password = ''  # Password (if required by the camera)

    # Construct the RTSP stream URL
    rtsp_url = f"rtsp://{username}:{password}@{ip_address}:{port}/testia/live"

    # Display the RTSP stream using Hugging Face Streamlit
    stframe = st.empty()
    predictions_frame = st.empty()

    with Stream(
        source=rtsp_url,
        model="jakojdnfin/1",
        output_channel_order="BGR",
        use_main_thread=True,  # for OpenCV display
        on_prediction=lambda predictions, image: update_predictions(predictions, image, predictions_frame)
    ) as stream:
        while True:
            image = stream.read()
            if image is not None and not image.empty():  # Check if the image is not empty
                # Display the image using matplotlib in the Streamlit app
                fig, ax = plt.subplots()
                ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
                ax.axis('off')  # Hide axis
                stframe.pyplot(fig)

def update_predictions(predictions, image, predictions_frame):
    # Update predictions frame with annotated image
    annotated_image = annotator.annotate(scene=image, detections=sv.Detections.from_roboflow(predictions))
    predictions_frame.image(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB), use_column_width=True)


if __name__ == "__main__":
    main()