Spaces:
Sleeping
Sleeping
File size: 4,253 Bytes
d04393d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import os
import pandas as pd
from PIL import Image
from animaloc.vizual import draw_points, draw_text
from inference.utils_io import mkdir, save_csv
def draw_detections_on_image(
image_path: str,
detections_df: pd.DataFrame,
output_path: str = None
) -> Image.Image:
"""
Dibuja puntos visibles sobre una imagen y añade una leyenda
con el total de detecciones y el desglose por especie.
Parámetros
----------
image_path : str
Ruta de la imagen original.
detections_df : pd.DataFrame
DataFrame con las detecciones (columnas: x, y, species, scores, etc.).
output_path : str, opcional
Ruta donde se guardará la imagen anotada.
Retorna
-------
Image.Image
Imagen con los puntos y leyenda dibujados.
"""
img = Image.open(image_path)
img_cpy = img.copy()
# Extraer coordenadas (y, x)
pts = list(detections_df[["y", "x"]].to_records(index=False))
pts = [(y, x) for y, x in pts]
# Dibujar puntos sobre la imagen
output = draw_points(img_cpy, pts, color=(255, 0, 0), size=60)
# Construir texto de leyenda
species_counts = detections_df["species"].value_counts().to_dict()
total = sum(species_counts.values())
legend = f"Detecciones: {total} | " + ", ".join(
[f"{sp}: {n}" for sp, n in species_counts.items()]
)
# Posicionar texto en parte inferior
overlay_y = img_cpy.height - int(0.08 * img_cpy.height)
output = draw_text(
output,
text=legend,
position=(20, overlay_y),
font_size=int(0.04 * img_cpy.height),
)
# Guardar imagen si se especifica ruta de salida
if output_path:
mkdir(os.path.dirname(output_path))
output.save(output_path, quality=95)
return output
def compute_species_counts(detections_df: pd.DataFrame) -> dict:
"""
Calcula el número de detecciones por especie.
Retorna
-------
dict
Diccionario con las especies y sus conteos.
Retorna un diccionario vacío si no hay detecciones.
"""
if detections_df.empty:
return {}
return detections_df["species"].value_counts().to_dict()
def generate_thumbnails(
image_path: str,
detections_df: pd.DataFrame,
output_dir: str,
thumb_size: int = 256
) -> None:
"""
Genera miniaturas recortadas alrededor de cada detección,
con el nombre de la especie y su puntaje de confianza.
Parámetros
----------
image_path : str
Ruta de la imagen original.
detections_df : pd.DataFrame
DataFrame con las detecciones.
output_dir : str
Directorio donde se guardarán las miniaturas.
thumb_size : int
Tamaño (en píxeles) de cada miniatura cuadrada.
"""
mkdir(output_dir)
img = Image.open(image_path)
img_cpy = img.copy()
sp_score = list(detections_df[["species", "scores"]].to_records(index=False))
pts = list(detections_df[["y", "x"]].to_records(index=False))
for i, ((y, x), (sp, score)) in enumerate(zip(pts, sp_score)):
off = thumb_size // 2
coords = (x - off, y - off, x + off, y + off)
# Recortar miniatura
thumbnail = img_cpy.crop(coords)
# Dibujar texto con especie y score
score = round(score * 100, 1)
thumbnail = draw_text(
thumbnail,
f"{sp} | {score}%",
position=(10, 5),
font_size=int(0.08 * thumb_size),
)
filename = os.path.basename(image_path)[:-4] + f"_{i}.JPG"
thumbnail.save(os.path.join(output_dir, filename))
def save_detections(
detections_df: pd.DataFrame,
output_dir: str,
logger=None
) -> str:
"""
Guarda las detecciones en formato CSV dentro del directorio de salida.
Parámetros
----------
detections_df : pd.DataFrame
DataFrame con las detecciones.
output_dir : str
Directorio de salida.
logger : logging.Logger, opcional
Logger para registrar el proceso.
Retorna
-------
str
Ruta del archivo CSV guardado.
"""
mkdir(output_dir)
csv_path = os.path.join(output_dir, "detections.csv")
save_csv(detections_df, csv_path, logger)
return csv_path
|