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 = """

Porto do Itaqui YOLO-World

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)