app.py

from typing import List

import os
import cv2
import gradio as gr
import numpy as np
import supervision as sv
import torch
from tqdm import tqdm
import cv2
from translate import Translator

from inference.models.yolo_world.yolo_world import YOLOWorld


import datetime
import uuid
from typing import List


def generate_file_name(extension="mp4"):
    current_datetime = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    unique_id = uuid.uuid4()
    return f"{current_datetime}_{unique_id}.{extension}"


def list_files_older_than(directory: str, diff_minutes: int) -> List[str]:
    diff_seconds = diff_minutes * 60
    now = datetime.datetime.now()
    older_files: List[str] = []

    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path):
            file_mod_time = os.path.getmtime(file_path)
            file_mod_datetime = datetime.datetime.fromtimestamp(file_mod_time)
            time_diff = now - file_mod_datetime
            if time_diff.total_seconds() > diff_seconds:
                older_files.append(file_path)

    return older_files

def remove_files_older_than(directory: str, diff_minutes: int) -> None:
    older_files = list_files_older_than(directory, diff_minutes)
    file_count = len(older_files)

    for file_path in older_files:
        os.remove(file_path)

    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(
        f"[{now}] Removed {file_count} files older than {diff_minutes} minutes from "
        f"'{directory}' directory."
    )


def calculate_end_frame_index(source_video_path: str) -> int:
    video_info = sv.VideoInfo.from_video_path(source_video_path)
    return video_info.total_frames

def create_directory(directory_path: str) -> None:
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)


MARKDOWN = """
<h1>Porto do Itaqui YOLO-World </h1>


Este é protótipo em fase de execução que será apresentado ao porto do Itaqui com o objetivo de entregar alguma coisa.

"""

RESULTS = "results"

# IMAGE_EXAMPLES = [
#     ['https://media.roboflow.com/dog.jpeg', 'dog, eye, nose, tongue, car', 0.005, 0.1, True, False, False],
#     ['https://media.roboflow.com/albert-4x.png', 'hand, hair', 0.005, 0.1, True, False, False],
# ]
# VIDEO_EXAMPLES = [
#     ['https://media.roboflow.com/supervision/video-examples/croissant-1280x720.mp4', 'croissant', 0.01, 0.2, False, False, False],
#     ['https://media.roboflow.com/supervision/video-examples/suitcases-1280x720.mp4', 'suitcase', 0.1, 0.2, False, False, False],
#     ['https://media.roboflow.com/supervision/video-examples/tokyo-walk-1280x720.mp4', 'woman walking', 0.1, 0.2, False, False, False],
#     ['https://media.roboflow.com/supervision/video-examples/wooly-mammoth-1280x720.mp4', 'mammoth', 0.01, 0.2, False, False, False],
# ]

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# EFFICIENT_SAM_MODEL = load(device=DEVICE)
YOLO_WORLD_MODEL = YOLOWorld(model_id="yolo_world/l")

BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator(thickness=2)
MASK_ANNOTATOR = sv.MaskAnnotator()
LABEL_ANNOTATOR = sv.LabelAnnotator(text_thickness=2, text_scale=1, text_color=sv.Color.BLACK)

# creating video results directory
create_directory(directory_path=RESULTS)


def process_categories(categories: str) -> List[str]:
    # Traduzindo as palavras do português para o inglês
    translator = Translator(from_lang="pt", to_lang="en")
    translation = translator.translate(categories)

    return [category.strip() for category in translation.split(',')]


def annotate_image(
    input_image: np.ndarray,
    detections: sv.Detections,
    categories: List[str],
    with_confidence: bool = False,
) -> np.ndarray:
    labels = [
        (
            f"{categories[class_id]}: {confidence:.3f}"
            if with_confidence
            else f"{categories[class_id]}"
        )
        for class_id, confidence in
        zip(detections.class_id, detections.confidence)
    ]
    output_image = MASK_ANNOTATOR.annotate(input_image, detections)
    output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
    output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
    return output_image


def process_image(
    input_image: np.ndarray,
    categories: str,
    confidence_threshold: float = 0.3,
    iou_threshold: float = 0.5,
    # with_segmentation: bool = True,
    with_confidence: bool = False,
    with_class_agnostic_nms: bool = False,
) -> np.ndarray:
    # cleanup of old video files
    remove_files_older_than(RESULTS, 30)

    categories = process_categories(categories)
    YOLO_WORLD_MODEL.set_classes(categories)
    results = YOLO_WORLD_MODEL.infer(input_image, confidence=0.02)
    detections = sv.Detections.from_inference(results)
    detections = detections.with_nms(
        class_agnostic=with_class_agnostic_nms,
        threshold=iou_threshold
    )
    # if with_segmentation:
    #     detections.mask = inference_with_boxes(
    #         image=input_image,
    #         xyxy=detections.xyxy,
    #         model=EFFICIENT_SAM_MODEL,
    #         device=DEVICE
    #     )
    output_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
    output_image = annotate_image(
        input_image=output_image,
        detections=detections,
        categories=categories,
        with_confidence=with_confidence
    )
    return cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)


def process_video(
    input_video: str,
    categories: str,
    confidence_threshold: float = 0.3,
    iou_threshold: float = 0.5,
    # with_segmentation: bool = True,
    with_confidence: bool = False,
    with_class_agnostic_nms: bool = False,
    progress=gr.Progress(track_tqdm=True)
) -> str:
    # cleanup of old video files
    remove_files_older_than(RESULTS, 30)

    categories = process_categories(categories)
    YOLO_WORLD_MODEL.set_classes(categories)
    video_info = sv.VideoInfo.from_video_path(input_video)
    total = calculate_end_frame_index(input_video)
    frame_generator = sv.get_video_frames_generator(
        source_path=input_video,
        end=total
    )
    result_file_name = generate_file_name(extension="mp4")
    result_file_path = os.path.join(RESULTS, result_file_name)
    with sv.VideoSink(result_file_path, video_info=video_info) as sink:
        for _ in tqdm(range(total), desc="Processing video..."):
            frame = next(frame_generator)
            results = YOLO_WORLD_MODEL.infer(frame, confidence=confidence_threshold)
            detections = sv.Detections.from_inference(results)
            detections = detections.with_nms(
                class_agnostic=with_class_agnostic_nms,
                threshold=iou_threshold
            )
            # if with_segmentation:
            #     detections.mask = inference_with_boxes(
            #         image=frame,
            #         xyxy=detections.xyxy,
            #         model=EFFICIENT_SAM_MODEL,
            #         device=DEVICE
            # )
            frame = annotate_image(
                input_image=frame,
                detections=detections,
                categories=categories,
                with_confidence=with_confidence
            )
            sink.write_frame(frame)
    return result_file_path


confidence_threshold_component = gr.Slider(
    minimum=0,
    maximum=1.0,
    value=0.03,
    step=0.01,
    label="Limite de Confiança",
    info=(
        "O limite de confiança para o modelo YOLO-World. Reduza o limite para "
        "reduzir falsos negativos, aumentando a sensibilidade do modelo para detectar "
        "objetos procurados. Por outro lado, aumente o limite para minimizar falsos "
        "positivos, evitando que o modelo identifique objetos que não deveria."
    ))

iou_threshold_component = gr.Slider(
    minimum=0,
    maximum=1.0,
    value=0.1,
    step=0.01,
    label="Limite IoU",
    info=(
        "Limite de intersecção sobre união (Intersection over Union ou IoU) para supressão não máxima. "
        "Diminua o valor para diminuir a ocorrência de caixas delimitadoras sobrepostas, "
        "tornando o processo de detecção mais rigoroso. Por outro lado, aumente o valor "
        "para permitir mais caixas delimitadoras sobrepostas, acomodando uma gama mais ampla de "
        "detecções."
    ))

with_segmentation_component = gr.Checkbox(
    value=False,
    label="With Segmentation",
    info=(
        "Whether to run EfficientSAM for instance segmentation."
    )
)

with_confidence_component = gr.Checkbox(
    value=False,
    label="Mostrar confiança.",
    info=(
        "Mostrar ou não a confiança dos objetos detectados."
    )
)

with_class_agnostic_nms_component = gr.Checkbox(
    value=False,
    label="Use NMS (Non-Max Supression ou Supressão Não Máxima) independente da classe.",
    info=(
        "Suprima caixas delimitadoras sobrepostas em todas as classes."
    )
)


with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Accordion("Confiduração", open=False):
        confidence_threshold_component.render()
        iou_threshold_component.render()
        with gr.Row():
            # with_segmentation_component.render()
            with_confidence_component.render()
            with_class_agnostic_nms_component.render()
    with gr.Tab(label="Imagem"):
        with gr.Row():
            input_image_component = gr.Image(
                type='numpy',
                label='Imagem de entrada'
            )
            output_image_component = gr.Image(
                type='numpy',
                label='Imagem de saída'
            )
        with gr.Row():
            image_categories_text_component = gr.Textbox(
                label='Categorias',
                placeholder='Digite as categorias separadas por vírgula',
                scale=7
            )
            image_submit_button_component = gr.Button(
                value='Submeter',
                scale=1,
                variant='primary'
            )
        # gr.Examples(
        #     fn=process_image,
        #     examples=IMAGE_EXAMPLES,
        #     inputs=[
        #         input_image_component,
        #         image_categories_text_component,
        #         confidence_threshold_component,
        #         iou_threshold_component,
        #         with_segmentation_component,
        #         with_confidence_component,
        #         with_class_agnostic_nms_component
        #     ],
        #     outputs=output_image_component
        # )
    with gr.Tab(label="Video"):
        with gr.Row():
            input_video_component = gr.Video(
                label='Video de entrada'
            )
            output_video_component = gr.Video(
                label='Video de saída'
            )
        with gr.Row():
            video_categories_text_component = gr.Textbox(
                label='Categorias',
                placeholder='Digite as categorias separadas por vírgula',
                scale=7
            )
            video_submit_button_component = gr.Button(
                value='Submeter',
                scale=1,
                variant='primary'
            )
        # gr.Examples(
        #     fn=process_video,
        #     examples=VIDEO_EXAMPLES,
        #     inputs=[
        #         input_video_component,
        #         video_categories_text_component,
        #         confidence_threshold_component,
        #         iou_threshold_component,
        #         with_segmentation_component,
        #         with_confidence_component,
        #         with_class_agnostic_nms_component
        #     ],
        #     outputs=output_image_component
        # )

    image_submit_button_component.click(
        fn=process_image,
        inputs=[
            input_image_component,
            image_categories_text_component,
            confidence_threshold_component,
            iou_threshold_component,
            # with_segmentation_component,
            with_confidence_component,
            with_class_agnostic_nms_component
        ],
        outputs=output_image_component
    )
    video_submit_button_component.click(
        fn=process_video,
        inputs=[
            input_video_component,
            video_categories_text_component,
            confidence_threshold_component,
            iou_threshold_component,
            # with_segmentation_component,
            with_confidence_component,
            with_class_agnostic_nms_component
        ],
        outputs=output_video_component
    )

demo.launch(debug=False, show_error=True, max_threads=1)