es.davy.ai

Preguntas y respuestas de programación confiables

¿Tienes una pregunta?

Si tienes alguna pregunta, puedes hacerla a continuación o ingresar lo que estás buscando.

PyTorch: guardar solo imágenes positivas donde se detecten peatones.

Estoy un poco atrapado en un código.

Quiero guardar SOLAMENTE IMÁGENES POSITIVAS cuando se detecta uno o más peatones. Cuando no se detecta nada, no hacer nada.

Comencé leyendo: https://github.com/ultralytics/yolov5/issues/36

Escribí esto:

import torch
import os

f = []
for dirpath, subdirs, files in os.walk('MyFolderWithPictures'):
    for x in files:
        if x.endswith(".jpg"):
            f.append(os.path.join(dirpath, x))

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

model.conf = 0.25  # Umbral de confianza de NMS
model.iou = 0.45  # Umbral de IoU para NMS
model.classes = 0   # Solo peatones
model.multi_label = False  # Múltiples etiquetas por caja en NMS
model.max_det = 1000  # Máximo número de detecciones por imagen

img = f  # lista de imágenes

results = model(img)

results.print()
results.save()

Pero esto imprime y guarda TODAS las imágenes (positivas y negativas).

Quiero guardar solo las imágenes con peatones.

¿Puede ayudarme? Gracias de antemano.

PD: La salida da:

imagen 1/13: 1080x1920 1 persona
imagen 2/13: 1080x1920 (no se encontraron detecciones)
imagen 3/13: 1080x1920 (no se encontraron detecciones)
imagen 4/13: 1080x1920 (no se encontraron detecciones)
imagen 5/13: 1080x1920 (no se encontraron detecciones)
imagen 6/13: 1080x1920 (no se encontraron detecciones)
imagen 7/13: 1080x1920 (no se encontraron detecciones)
imagen 8/13: 1080x1920 (no se encontraron detecciones)
imagen 9/13: 1080x1920 (no se encontraron detecciones)
imagen 10/13: 1080x1920 (no se encontraron detecciones)
imagen 11/13: 1080x1920 (no se encontraron detecciones)
imagen 12/13: 1080x1920 1 persona
imagen 13/13: 1080x1920 (no se encontraron detecciones)
Velocidad: 18.6ms de preprocesamiento, 119.8ms de inferencia, 1.9ms de NMS por imagen con forma (13, 3, 384, 640)
Guardadas 13 imágenes en runs\detect\exp

Agregar SOLUCIÓN:

for item in f:
    # Imágenes
    img = item  # o archivo, Ruta, PIL, OpenCV, numpy, lista
    # Inferencia
    results = model(img)
    # Resultados
    results.print()  # o .show(), .save(), .crop(), .pandas(), etc.
    if 0 in results.pandas().xyxy[0]['class']:
        results.save()
Tags:  , ,

Answer

  1. Avatar for davy.ai

    Para guardar solo imágenes con peatones detectados, puedes modificar el código de la siguiente manera:

    for item in f:
        # Imágenes
        img = item  # o archivo, Path, PIL, OpenCV, numpy, lista
        # Inferencia
        results = model(img)
        # Resultados
        results.print()  # o .show(), .save(), .crop(), .pandas(), etc.
        if 0 in results.pandas().xyxy[0]['class']:
            results.save()
    

    La condición if 0 in results.pandas().xyxy[0]['class'] verifica si se detecta un peatón en la imagen buscando la clase 0, que es el ID de clase para los peatones. Si se detecta un peatón, entonces results.save() guardará la imagen con los cuadros delimitadores de los peatones detectados.

Comments are closed.