Haziel1357's picture
Scripts, CSV de EMMa y guia de peso
a294091 verified
Raw History Blame Contribute Delete
13.5 kB
"""Une las fuentes descargadas en un solo conjunto de datos.
uv run python -m entrenamiento.preparar
1. Lee cada fuente descargada y traduce sus etiquetas a material y defectos.
2. Quita imágenes repetidas (mismo contenido aunque vengan de fuentes distintas).
3. Limita cuántas imágenes aporta cada clase de cada fuente y cada material, para que
ninguna domine el entrenamiento.
4. Reparte en entrenamiento / validación / prueba (80/10/10) por "grupo": las variantes
aumentadas de una misma foto (Roboflow las nombra foto_jpg.rf.<hash>.jpg) quedan siempre
en el mismo lado, para que la evaluación no sea tramposa.
5. Guarda una copia reducida (512 px) de cada imagen en datos/procesados/imagenes/ y el
listado en datos/procesados/manifiesto.jsonl.
"""
import hashlib
import json
import random
import re
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass, field
from pathlib import Path
from PIL import Image, ImageOps
from app.ml.estado.defectos import ORDEN_DEFECTOS
from app.ml.materiales import ORDEN_MATERIALES
from entrenamiento.descargar import MARCA
from entrenamiento.fuentes import FUENTES, Fuente
from entrenamiento.mapeo import LIMPIO, defectos_de_clase
from entrenamiento.rutas import (
CRUDOS,
IMAGENES,
LADO_GUARDADO,
MANIFIESTO,
asegurar_carpetas,
)
EXTENSIONES = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
MAX_POR_MATERIAL = 4000
SEMILLA = 2026
Image.MAX_IMAGE_PIXELS = 60_000_000
@dataclass
class Registro:
fuente: str
ruta: str # archivo original
material: str | None
# Solo los defectos conocidos: 1 = presente, 0 = ausente. Los que faltan, desconocidos.
defectos: dict[str, int] = field(default_factory=dict)
grupo: str = ""
id: str = ""
split: str = ""
# Recuadro (x1, y1, x2, y2) si el registro es un objeto recortado de una foto mayor.
recorte: list[int] | None = None
@property
def clave_cupo(self) -> str:
positivos = sorted(d for d, v in self.defectos.items() if v)
return f"{self.material}|{'+'.join(positivos) or ('limpio' if self.defectos else '-')}"
def _norm(texto: str) -> str:
return re.sub(r"[\s_\-]+", " ", texto.lower()).strip()
def _grupo(fuente: str, ruta: Path) -> str:
# Roboflow: "IMG_123_jpg.rf.abc123.jpg" → todas las variantes comparten "IMG_123_jpg".
return f"{fuente}/{ruta.name.split('.rf.')[0]}"
def _imagenes(carpeta: Path):
for ruta in carpeta.rglob("*"):
if (
ruta.suffix.lower() in EXTENSIONES
and "__MACOSX" not in ruta.parts
and not ruta.name.startswith("._")
):
yield ruta
# ------------------------------------------------------------------------------ lectores
def _leer_carpetas_material(fuente: Fuente, carpeta: Path) -> list[Registro]:
clases = {_norm(k): v for k, v in fuente.clases.items()}
ausentes = dict.fromkeys(fuente.defectos_ausentes, 0)
registros = []
for ruta in _imagenes(carpeta):
clase = next((_norm(p) for p in reversed(ruta.parts[:-1]) if _norm(p) in clases), None)
if clase is None or clases[clase] is None:
continue
registros.append(
Registro(
fuente.nombre, str(ruta), clases[clase], dict(ausentes), _grupo(fuente.nombre, ruta)
)
)
return registros
def _defectos_conocidos(nombres_clase: list[str]) -> set[str]:
conocidos: set[str] = set()
for nombre in nombres_clase:
conocidos |= defectos_de_clase(nombre) - {LIMPIO}
return conocidos
def _etiquetar(clases_imagen: set[str], conocidos: set[str]) -> dict[str, int]:
presentes: set[str] = set()
for nombre in clases_imagen:
presentes |= defectos_de_clase(nombre) - {LIMPIO}
return {d: int(d in presentes) for d in conocidos}
def _leer_coco(fuente: Fuente, carpeta: Path) -> list[Registro]:
anotaciones = list(carpeta.rglob("_annotations.coco.json"))
categorias_todas: list[str] = []
datos = []
for archivo in anotaciones:
coco = json.loads(archivo.read_text(encoding="utf-8"))
categorias = {c["id"]: c["name"] for c in coco.get("categories", [])}
categorias_todas += list(categorias.values())
datos.append((archivo.parent, coco, categorias))
conocidos = _defectos_conocidos(categorias_todas)
registros = []
material = fuente.material if fuente.usar_para_material else None
for carpeta_split, coco, categorias in datos:
por_imagen: dict[int, set[str]] = defaultdict(set)
for anotacion in coco.get("annotations", []):
por_imagen[anotacion["image_id"]].add(categorias.get(anotacion["category_id"], ""))
for imagen in coco.get("images", []):
ruta = carpeta_split / imagen["file_name"]
if ruta.is_file():
registros.append(
Registro(
fuente.nombre,
str(ruta),
material,
_etiquetar(por_imagen.get(imagen["id"], set()), conocidos),
_grupo(fuente.nombre, ruta),
)
)
return registros
MAX_RECORTES_POR_IMAGEN = 2
LADO_MINIMO_RECORTE = 48
MARGEN_RECORTE = 0.1
def _leer_recortes(fuente: Fuente, carpeta: Path) -> list[Registro]:
"""Un registro por objeto etiquetado (los más grandes de cada foto), con su material."""
patrones = [(re.compile(p, re.I), m) for p, m in fuente.recortes.items()]
ausentes = dict.fromkeys(fuente.defectos_ausentes, 0)
registros = []
for archivo in carpeta.rglob("_annotations.coco.json"):
coco = json.loads(archivo.read_text(encoding="utf-8"))
material_de = {}
for c in coco.get("categories", []):
material_de[c["id"]] = next((m for p, m in patrones if p.search(c["name"])), None)
imagenes = {i["id"]: i for i in coco.get("images", [])}
por_imagen: dict[int, list[tuple[float, list[int], str]]] = defaultdict(list)
for a in coco.get("annotations", []):
material = material_de.get(a["category_id"])
imagen = imagenes.get(a["image_id"])
if material is None or imagen is None:
continue
x, y, w, h = a["bbox"]
if min(w, h) < LADO_MINIMO_RECORTE:
continue
mx, my = w * MARGEN_RECORTE, h * MARGEN_RECORTE
caja = [
max(0, int(x - mx)),
max(0, int(y - my)),
min(imagen["width"], int(x + w + mx)),
min(imagen["height"], int(y + h + my)),
]
por_imagen[a["image_id"]].append((w * h, caja, material))
for id_imagen, objetos in por_imagen.items():
ruta = archivo.parent / imagenes[id_imagen]["file_name"]
if not ruta.is_file():
continue
for _, caja, material in sorted(objetos, reverse=True)[:MAX_RECORTES_POR_IMAGEN]:
registros.append(
Registro(
fuente.nombre,
str(ruta),
material,
dict(ausentes),
_grupo(fuente.nombre, ruta),
recorte=caja,
)
)
return registros
def _leer_clasificacion_rf(fuente: Fuente, carpeta: Path) -> list[Registro]:
rutas = [r for r in _imagenes(carpeta)]
conocidos = _defectos_conocidos(sorted({r.parent.name for r in rutas}))
material = fuente.material if fuente.usar_para_material else None
registros = []
for ruta in rutas:
clase = ruta.parent.name
if not defectos_de_clase(clase) and conocidos:
continue # clase que no dice nada del estado
registros.append(
Registro(
fuente.nombre,
str(ruta),
material,
_etiquetar({clase}, conocidos),
_grupo(fuente.nombre, ruta),
)
)
return registros
def leer_fuente(fuente: Fuente) -> list[Registro]:
carpeta = CRUDOS / fuente.nombre
if not fuente.entrenar or not (carpeta / MARCA).is_file():
return []
if fuente.recortes:
return _leer_recortes(fuente, carpeta)
if fuente.clases:
return _leer_carpetas_material(fuente, carpeta)
if any(carpeta.rglob("_annotations.coco.json")):
return _leer_coco(fuente, carpeta)
return _leer_clasificacion_rf(fuente, carpeta)
# ---------------------------------------------------------------------- limpieza y reparto
def _limitar(
registros: list[Registro], fuentes: dict[str, Fuente], rng: random.Random
) -> list[Registro]:
# 1) cupo por (fuente, material, tipo de defecto)
por_clave: dict[tuple[str, str], list[Registro]] = defaultdict(list)
for r in registros:
por_clave[(r.fuente, r.clave_cupo)].append(r)
elegidos: list[Registro] = []
for (fuente, _), lista in por_clave.items():
rng.shuffle(lista)
elegidos += lista[: fuentes[fuente].max_por_clase]
# 2) cupo por material: se priorizan las imágenes que también enseñan defectos
por_material: dict[str | None, list[Registro]] = defaultdict(list)
for r in elegidos:
por_material[r.material].append(r)
final: list[Registro] = []
for material, lista in por_material.items():
if material is None:
final += lista
continue
con_defectos = [r for r in lista if any(r.defectos.values()) or "limpio" in r.clave_cupo]
resto = [r for r in lista if r not in con_defectos]
rng.shuffle(con_defectos)
rng.shuffle(resto)
final += (con_defectos + resto)[:MAX_POR_MATERIAL]
return final
def _split(grupo: str) -> str:
valor = int(hashlib.md5(grupo.encode()).hexdigest(), 16) % 100
return "entrenamiento" if valor < 80 else "validacion" if valor < 90 else "prueba"
def _guardar_copia(registro: Registro) -> tuple[Registro, str | None]:
"""Guarda la copia reducida. Devuelve (registro, huella del contenido) o None si falla."""
try:
with Image.open(registro.ruta) as original:
if registro.recorte:
# Las cajas de COCO se refieren a los píxeles tal como están guardados.
imagen = original.convert("RGB").crop(tuple(registro.recorte))
else:
imagen = ImageOps.exif_transpose(original).convert("RGB")
if min(imagen.size) < 32:
return registro, None
imagen.thumbnail((LADO_GUARDADO, LADO_GUARDADO), Image.Resampling.LANCZOS)
huella = hashlib.md5(imagen.resize((16, 16)).convert("L").tobytes()).hexdigest()
imagen.save(IMAGENES / f"{registro.id}.jpg", quality=92)
return registro, huella
except Exception:
return registro, None
def main() -> None:
asegurar_carpetas()
rng = random.Random(SEMILLA)
fuentes = {f.nombre: f for f in FUENTES}
registros: list[Registro] = []
for fuente in FUENTES:
leidos = leer_fuente(fuente)
if leidos:
print(f"- {fuente.nombre}: {len(leidos)} imágenes")
registros += leidos
if not registros:
raise SystemExit(
"No hay fuentes descargadas. Ejecuta antes: python -m entrenamiento.descargar"
)
registros = _limitar(registros, fuentes, rng)
for r in registros:
r.id = hashlib.md5(f"{r.ruta}|{r.recorte}".encode()).hexdigest()[:16]
r.split = _split(r.grupo)
print(f"Guardando {len(registros)} copias reducidas en {IMAGENES} ...")
vistos: set[str] = set()
validos: list[Registro] = []
with ThreadPoolExecutor(max_workers=6) as pool:
for registro, huella in pool.map(_guardar_copia, registros):
if huella is None:
continue
if huella in vistos: # misma imagen en dos fuentes o repetida
(IMAGENES / f"{registro.id}.jpg").unlink(missing_ok=True)
continue
vistos.add(huella)
validos.append(registro)
with MANIFIESTO.open("w", encoding="utf-8") as salida:
for r in sorted(validos, key=lambda r: r.id):
salida.write(json.dumps(asdict(r), ensure_ascii=False) + "\n")
print(f"\nManifiesto: {MANIFIESTO} ({len(validos)} imágenes)")
print("\nMaterial por reparto:")
tabla = Counter((r.material, r.split) for r in validos)
for material in (*ORDEN_MATERIALES, None):
print(
f" {str(material):10} "
+ " ".join(
f"{s[:5]}={tabla[(material, s)]:5}"
for s in ("entrenamiento", "validacion", "prueba")
)
)
print("\nDefectos (positivos / conocidos):")
for d in ORDEN_DEFECTOS:
conocidos = sum(1 for r in validos if d in r.defectos)
positivos = sum(r.defectos.get(d, 0) for r in validos)
print(f" {d:12} {positivos:5} / {conocidos:5}")
if __name__ == "__main__":
main()