helper.py

import time
import os
import uuid
import datetime
import pytz
import yaml
from yaml.loader import SafeLoader

import pandas as pd
import streamlit as st
import streamlit_authenticator as stauth # type: ignore
import cv2
from pytube import YouTube
from ultralytics import YOLO

import settings
import database

@st.cache_resource  # 👈 Add the caching decorator
def load_model(model_path):
    """
    Loads a YOLO object detection model from the specified model_path.

    Parameters:
        model_path (str): The path to the YOLO model file.

    Returns:
        A YOLO object detection model.
    """
    model = YOLO(model_path)
    return model


def save_detection_job(res, uploaded_image, model_path, model_type, confidence):
    """
    Fonction de sauvegarde des fichiers et métadonnées d'une détection.

    Args : 
    ----------
    - res : Object "result" d'une détection 
    - uploaded_image : Image envoyée par l'utilisateur (objet Image PIL)
    - model_path : Chemin vers les poids personnalisés
    - model_type : Type du modèle (Détection ou Segmentation)
    - confidence : Seuil de confiance (réglé par l'utilisateur)

    Retourne : 
    ----------
    - job_id : l'identifiant de la tâche de détection
    """
    with st.status("Sauvegarde de la détection...", expanded=True) as status:

        # Création des dossiers pour sauvegarder les images et les labels
        if not os.path.exists('detections/imgs-original'):
            os.makedirs(os.path.join("detections", "imgs-original"))
        if not os.path.exists('detections/imgs-detected'):
            os.makedirs(os.path.join("detections", "imgs-detected"))
        if not os.path.exists('detections/labels'):
            os.makedirs(os.path.join("detections", "labels"))

        # Date et heure de la détection
        detection_timezone = pytz.timezone('Europe/Paris')
        detection_datetime = datetime.datetime.now(detection_timezone) 
        # String de date et heure
        UID = detection_datetime.strftime("%Y-%m-%d-%H%M%S")

        # Création d'identifiants uniques pour chaque table
        job_id = uuid.uuid4()
        og_id = uuid.uuid4()
        dt_id = uuid.uuid4()
        label_id = uuid.uuid4()

        # Table "app_imgs_original"
        st.write("Sauvegarde de l'image originale.")
        og_img_dict = {}
        og_img_dict['og_id'] = og_id
        og_img_dict['og_filename'] = f"{UID}-original.jpg"
        og_img_dict['og_filepath'] = f"detections/imgs-original/{og_img_dict['og_filename']}"
        og_img_dict['og_height'] = res[0].orig_shape[0]
        og_img_dict['og_width'] = res[0].orig_shape[1]
        # Sauvegarde dans la table
        og_img_df = pd.DataFrame(og_img_dict, index=[0])
        database.insert_dataframe_to_table(og_img_df, "app_imgs_original", "og_id", if_exists = 'append')
        # Sauvegarder le fichier image
        uploaded_image.save(og_img_dict['og_filepath'])

        # Table "app_imgs_detected"
        st.write("Sauvegarde de l'image des détections.")
        detected_img_dict = {}
        detected_img_dict['dt_id'] = dt_id
        detected_img_dict['dt_filename'] = f"{UID}-detected.jpg"
        detected_img_dict['dt_filepath'] = f"detections/imgs-detected/{detected_img_dict['dt_filename']}"
        # Sauvegarde dans la table
        detected_img_df = pd.DataFrame(detected_img_dict, index=[0])
        database.insert_dataframe_to_table(detected_img_df, "app_imgs_detected", "dt_id", if_exists = 'append')
        # Sauvegarder le fichier de l'image détectée
        res[0].save(filename=detected_img_dict['dt_filepath'])

        # Table "app_detection_labels"
        st.write("Sauvegarde des lables au format '.txt'.")
        label_dict = {}
        label_dict['label_id'] = label_id
        label_dict['label_filename'] = f"{UID}.txt"
        label_dict['label_filepath'] = f"detections/labels/{label_dict['label_filename']}"
        # Sauvegarde dans la table
        label_df = pd.DataFrame(label_dict, index=[0])
        database.insert_dataframe_to_table(label_df, "app_detection_labels", "label_id", if_exists = 'append')
        # Enregistre les prédictions dans un fichier txt.
        res[0].save_txt(label_dict['label_filepath'])

        # Table "app_detection_jobs"
        st.write("Sauvegarde de la tâche de détection.")
        # Métadonnées de la tâche de détection
        job_dict = {}
        job_dict['job_id'] = job_id
        job_dict['job_speed'] = float(sum(res[0].speed.values())) # vitesse de détection en ms
        job_dict['job_confidence'] = confidence
        job_dict['job_task'] = model_type
        job_dict['job_count'] = len(res[0].boxes)
        job_dict['job_model_filename'] = os.path.basename(str(model_path))
        job_dict['job_model_weights_path'] = str(model_path)
        # Clés étrangères
        job_dict['job_og_id'] = og_id
        job_dict['job_dt_id'] = dt_id
        job_dict['job_label_id'] = label_id
        job_dict['job_user_id'] = st.session_state["username"]
        # Création d'un DataFrame
        job_df = pd.DataFrame(job_dict, index=[0])
        job_df['job_created_at'] = pd.Timestamp.now(tz="Europe/Paris")
        # Sauvegarde du DataFrame dans la table "app_detection_jobs"
        database.insert_dataframe_to_table(job_df, "app_detection_jobs", "job_id", if_exists = 'append')

        # Table "app_detection_boxes"
        st.write("Sauvegarde des boîtes de détection.")
        # Nom des classes détectées par le modèle
        classes_dict = res[0].names
        # Boîtes de détection
        boxes = res[0].boxes
        boxes_dict_list = []
        for box in boxes:
            box_cpu = box.cpu()
            box_numpy = box_cpu.numpy()
            box_xywhn = box_numpy.xywhn.tolist()
            box_dict = {}
            box_dict['box_id'] = uuid.uuid4()
            box_dict['box_class_id'] = int(box_numpy.cls.tolist()[0])
            box_dict['box_class_name'] =  classes_dict[int(box_numpy.cls.tolist()[0])]
            box_dict['box_x_center'] = box_xywhn[0][0]
            box_dict['box_y_center'] = box_xywhn[0][1]
            box_dict['box_width'] = box_xywhn[0][2]
            box_dict['box_height'] = box_xywhn[0][3]
            box_dict['box_conf'] = round(box_numpy.conf.tolist()[0], 4)
            box_dict['box_label_id'] = label_id
            box_dict['box_job_id'] = job_id
            boxes_dict_list.append(box_dict)
        # Sauvegarder les boîtes de détection
        boxes_df = pd.DataFrame(boxes_dict_list)
        database.insert_dataframe_to_table(boxes_df, "app_detection_boxes", "box_id", if_exists = 'append')

        # Fin de la tâche de détection
        status.update(label="Tâche de détection terminée !", state="complete", expanded=False)
        print(f"Tâche de détection terminé ! Job ID : '{job_id}'.")
        return job_id

def display_tracker_options():
    display_tracker = st.radio("Display Tracker", ('Yes', 'No'), key="tracker-display")
    is_display_tracker = True if display_tracker == 'Yes' else False
    if is_display_tracker:
        tracker_type = st.radio("Tracker", ("bytetrack.yaml", "botsort.yaml"), key="tracker-type")
        return is_display_tracker, tracker_type
    return is_display_tracker, None

def _display_detected_frames(conf, model, st_frame, image, is_display_tracking=None, tracker=None):
    """
    Display the detected objects on a video frame using the YOLOv8 model.

    Args:
    - conf (float): Confidence threshold for object detection.
    - model (YoloV8): A YOLOv8 object detection model.
    - st_frame (Streamlit object): A Streamlit object to display the detected video.
    - image (numpy array): A numpy array representing the video frame.
    - is_display_tracking (bool): A flag indicating whether to display object tracking (default=None).

    Returns:
    None
    """

    # Resize the image to a standard size
    image = cv2.resize(image, (720, int(720*(9/16))))

    # Display object tracking, if specified
    if is_display_tracking:
        res = model.track(image, conf=conf, persist=True, tracker=tracker)
    else:
        # Predict the objects in the image using the YOLOv8 model
        res = model.predict(image, conf=conf)

    # # Plot the detected objects on the video frame
    res_plotted = res[0].plot()
    st_frame.image(res_plotted,
                   caption='Vidéo détectée',
                   channels="BGR",
                   use_column_width=True
                   )

def play_youtube_video(conf, model):
    """
    Plays a webcam stream. Detects Objects in real-time using the YOLOv8 object detection model.

    Parameters:
        conf: Confidence of YOLOv8 model.
        model: An instance of the `YOLOv8` class containing the YOLOv8 model.

    Returns:
        None

    Raises:
        None
    """
    source_youtube = st.sidebar.text_input("URL de la vidéo YouTube", key="source-youtube")

    is_display_tracker, tracker = display_tracker_options()

    if st.sidebar.button('Lancer la détection', use_container_width=True, key="start-detection-youtube"):
        try:
            yt = YouTube(source_youtube)
            stream = yt.streams.filter(file_extension="mp4", res=720).first()
            vid_cap = cv2.VideoCapture(stream.url)

            st_frame = st.empty()
            while (vid_cap.isOpened()):
                success, image = vid_cap.read()
                if success:
                    _display_detected_frames(conf,
                                             model,
                                             st_frame,
                                             image,
                                             is_display_tracker,
                                             tracker,
                                             )
                else:
                    vid_cap.release()
                    break
        except Exception as e:
            st.sidebar.error("Erreur de chargement de la vidéo : " + str(e))

def play_rtsp_stream(conf, model):
    """
    Plays an rtsp stream. Detects Objects in real-time using the YOLOv8 object detection model.

    Parameters:
        conf: Confidence of YOLOv8 model.
        model: An instance of the `YOLOv8` class containing the YOLOv8 model.

    Returns:
        None

    Raises:
        None
    """
    source_rtsp = st.sidebar.text_input("URL du flux RTSP :")
    st.sidebar.caption('Example URL: rtsp://admin:12345@192.168.1.210:554/Streaming/Channels/101')
    is_display_tracker, tracker = display_tracker_options()
    if st.sidebar.button('Lancer la détection', use_container_width=True, key="start-detection-rtsp"):
        try:
            vid_cap = cv2.VideoCapture(source_rtsp)
            st_frame = st.empty()
            while (vid_cap.isOpened()):
                success, image = vid_cap.read()
                if success:
                    _display_detected_frames(conf,
                                             model,
                                             st_frame,
                                             image,
                                             is_display_tracker,
                                             tracker
                                             )
                else:
                    vid_cap.release()
                    # vid_cap = cv2.VideoCapture(source_rtsp)
                    # time.sleep(0.1)
                    # continue
                    break
        except Exception as e:
            vid_cap.release()
            st.sidebar.error("Erreur lors du chargement du flux RTSP : " + str(e))

def play_webcam(conf, model):
    """
    Plays a webcam stream. Detects Objects in real-time using the YOLOv8 object detection model.

    Parameters:
        conf: Confidence of YOLOv8 model.
        model: An instance of the `YOLOv8` class containing the YOLOv8 model.

    Returns:
        None

    Raises:
        None
    """
    source_webcam = settings.WEBCAM_PATH
    is_display_tracker, tracker = display_tracker_options()
    if st.sidebar.button('Lancer la détection', use_container_width=True, key="start-detection-webcam"):
        try:
            vid_cap = cv2.VideoCapture(source_webcam)
            st_frame = st.empty()
            while (vid_cap.isOpened()):
                success, image = vid_cap.read()
                if success:
                    _display_detected_frames(conf,
                                             model,
                                             st_frame,
                                             image,
                                             is_display_tracker,
                                             tracker,
                                             )
                else:
                    vid_cap.release()
                    break
        except Exception as e:
            st.sidebar.error("Erreur de chargement de la vidéo : " + str(e))

def play_stored_video(conf, model):
    """
    Plays a stored video file. Tracks and detects objects in real-time using the YOLOv8 object detection model.

    Parameters:
        conf: Confidence of YOLOv8 model.
        model: An instance of the `YOLOv8` class containing the YOLOv8 model.

    Returns:
        None

    Raises:
        None
    """
    source_vid = st.sidebar.selectbox(
        "Choisissez une vidéo...", settings.VIDEOS_DICT.keys())

    is_display_tracker, tracker = display_tracker_options()

    with open(settings.VIDEOS_DICT.get(source_vid), 'rb') as video_file:
        video_bytes = video_file.read()
    if video_bytes:
        st.video(video_bytes)

    if st.sidebar.button('Lancer la détection', use_container_width=True, key="start-detection-button"):
        try:
            vid_cap = cv2.VideoCapture(
                str(settings.VIDEOS_DICT.get(source_vid)))
            st_frame = st.empty()
            while (vid_cap.isOpened()):
                success, image = vid_cap.read()
                if success:
                    _display_detected_frames(conf,
                                             model,
                                             st_frame,
                                             image,
                                             is_display_tracker,
                                             tracker
                                             )
                else:
                    vid_cap.release()
                    break
        except Exception as e:
            st.sidebar.error("Erreur de chargement de la vidéo : " + str(e))


# --- DELETE ---

def delete_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)

def delete_dir_files(directory_path):
    for filename in os.listdir(directory_path):
        if os.path.isfile(os.path.join(directory_path, filename)):
            os.remove(os.path.join(directory_path, filename))
    print(f"Suppression de tous les fichiers de '{directory_path}'.")

def clear_past_detections_files(dir_path):
    for root, sub_dir_names, files in os.walk(dir_path):
        for sub_dir_name in sub_dir_names:
            sub_dir_path = os.path.join(root, sub_dir_name)
            if os.path.isdir(sub_dir_path):
                for filename in os.listdir(sub_dir_path):
                    filename_path = os.path.join(sub_dir_path, filename)
                    if os.path.isfile(filename_path):
                        os.remove(filename_path)
                print(f"Suppression de tous les fichiers de '{sub_dir_path}'.")

def delete_detection(job_id):

    job_query = f"""
        SELECT job_id, job_og_id, og_filepath, job_dt_id, dt_filepath, job_label_id, label_filepath
        FROM app_detection_jobs
        JOIN app_imgs_original
        ON job_og_id = og_id
        JOIN app_imgs_detected
        ON job_dt_id = dt_id
        JOIN app_detection_labels
        ON job_label_id = label_id
        WHERE job_id = '{job_id}';
        """
    job_df = database.sql_query_to_dataframe(job_query)
    
    if not job_df.empty :
        job_dict = job_df.to_dict('records')[0]
        # Effacer les fichiers
        delete_file(job_dict['og_filepath'])
        delete_file(job_dict['dt_filepath'])
        delete_file(job_dict['label_filepath'])
        # Effacer les enregistrements dans les tables
        delete_query = f"""
        DELETE FROM app_detection_boxes WHERE box_job_id = '{job_id}';
        DELETE FROM app_detection_jobs WHERE job_id = '{job_id}';
        DELETE FROM app_imgs_original WHERE og_id = '{job_dict['job_og_id']}';
        DELETE FROM app_imgs_detected WHERE dt_id = '{job_dict['job_dt_id']}';
        DELETE FROM app_detection_labels WHERE label_id = '{job_dict['job_label_id']}';
        """
        database.simple_query_to_db(delete_query)
        print(f"Détection '{job_id}' effacée !")
        st.rerun()
    else:
        print(f"Row with job_id = '{job_id}' doesn't exist!")

def delete_all_detections():

    tables_list = [
        "app_detection_boxes",
        "app_detection_jobs",
        "app_imgs_original",
        "app_imgs_detected",
        "app_detection_labels"
    ]
    for table in tables_list:
        database.erase_table(table)
    clear_past_detections_files("./detections")


# --- DISPLAY ---
                
def display_detection_imgs(job_id):
    """
    Affichage des images d'une détection, à partir de l'identifiant de la tâche de détection (job_id).
    Affiche l'image originale et l'image détectée, dans deux colonnes.
    """

    # Requête de récupération des informations sur la tâche de détection à partir de son identifiant
    job_query = f"""
        SELECT job_id, og_filepath, dt_filepath
        FROM app_detection_jobs
        JOIN app_imgs_original
        ON job_og_id = og_id
        JOIN app_imgs_detected
        ON job_dt_id = dt_id
        WHERE job_id = '{job_id}';
        """
    # Résultats de la requête au format DataFrame
    job_df = database.sql_query_to_dataframe(job_query)
    
    if not job_df.empty :
        job_dict = job_df.to_dict('records')[0]

        col1, col2 = st.columns(2)
        with col1 : 
            if os.path.isfile(job_dict['og_filepath']):
                # st.subheader(job_dict['og_filename'])
                st.image(job_dict['og_filepath'])
        with col2 :
            if os.path.isfile(job_dict['dt_filepath']):
                # st.subheader(job_dict['dt_filename'])
                st.image(job_dict['dt_filepath'])

def display_detection_details(job_id):
    """
    Affichage des détails d'une détection à partir de son identifiant de tâche de détection (job_id).
    """

    # Requête de récupération des informations sur la tâche de détection à partir de son identifiant
    job_query = f"""
        SELECT job_model_filename, job_confidence, job_speed, job_count, og_filename, dt_filename, job_created_at, user_name
        FROM app_detection_jobs
        JOIN app_imgs_original
        ON job_og_id = og_id
        JOIN app_imgs_detected
        ON job_dt_id = dt_id
        JOIN app_users
        ON job_user_id = user_id
        WHERE job_id = '{job_id}';
        """
    job_df = database.sql_query_to_dataframe(job_query)
    if not job_df.empty:
        job_dict = job_df.to_dict('records')[0]

        st.markdown(f"""
            - **Date de détection** : {job_dict['job_created_at'].strftime('%Y-%m-%d %X')}
            - **Utilisateur** : {job_dict['user_name']}     
            - **Modèle** : {job_dict['job_model_filename']}
            - **Seuil de confiance** : {job_dict['job_confidence']}
            - **Vitesse de détection** : {round(job_dict['job_speed'], 2)} ms
            - **Nom de l'image originale** : {job_dict['og_filename']}
            - **Nom de l'image détectée** : {job_dict['dt_filename']}
        """)
        if job_dict['job_count'] == 0:
            st.warning(f"🚨 Cette image ne contient aucune détection !")
        

def display_detection_boxes(job_id):
    """
    Affichage des boîtes de détection d'une détection, à partir de l'identifiant de tâche de détection. 
    """
    # Requête pour récupérer les boîtes de détection dans la base de données
    boxes_sql_query = f"""
        SELECT * FROM app_detection_boxes
        WHERE box_job_id = '{job_id}'
        """
    boxes_df = database.sql_query_to_dataframe(boxes_sql_query)
    if not boxes_df.empty:
        st.markdown("""##### 📦 Boîtes de détection""")
        st.markdown(f"""
            - **Nombre de boîtes de détection** : {len(boxes_df)}
            - **Nombre de classes différentes** : {boxes_df['box_class_id'].nunique()}
            - **Liste des classes détectées** : {boxes_df['box_class_name'].unique().tolist()}
        """)
        st.dataframe(boxes_df[['box_class_name', 'box_class_id', 'box_x_center', 'box_y_center', 'box_width', 'box_height', 'box_conf']])

def display_label_download_button(job_id):
    """
    Affichage du boutton de téléchargement du fichier label.txt, résultat lié à une détection.
    """
    # Requête de récupération des informations sur la tâche de détection à partir de son identifiant
    job_query = f"""
        SELECT label_filepath, label_filename
        FROM app_detection_jobs
        JOIN app_detection_labels
        ON job_label_id = label_id
        WHERE job_id = '{job_id}';
        """
    job_df = database.sql_query_to_dataframe(job_query)
    if not job_df.empty:
        job_dict = job_df.to_dict('records')[0]
        if os.path.exists(job_dict['label_filepath']):
            with open(job_dict['label_filepath'], "rb") as file:
                btn = st.download_button(label="📄 Télécharger le fichier label.txt", data=file, file_name=job_dict['label_filename'], mime="text/plain")


# --- Authentication ---
        
def authentification_main():
    
    # Get credentials in Database, table 'app_users'
    with open('auth.yaml') as file:
        config = yaml.load(file, Loader=SafeLoader)

    authenticator = stauth.Authenticate(
        config['credentials'],
        config['cookie']['name'],
        config['cookie']['key'],
        config['cookie']['expiry_days'],
        config['preauthorized']
    )

    authenticator.login()

    if st.session_state["authentication_status"]:
        st.sidebar.subheader(f'Bienvenue {st.session_state["name"]} !')
        authenticator.logout(key="logout-button", button_name="Se déconnecter", location="sidebar")
    elif st.session_state["authentication_status"] is False:
        st.error("Le nom d'utilisateur ou le mot de passe est incorrect.")
    elif st.session_state["authentication_status"] is None:
        st.warning("Veuillez entrer votre nom d'utilisateur et votre mot de passe.")

    # Si l'utilisateur n'est pas identifié, ne pas afficher la suite de la page.
    if not st.session_state["authentication_status"]: 
        st.stop()