Download entrenamiento/preparar.py from EagleDev-Innovation/Modelos-EagleDev-Innovation: direct link, hf CLI and curl.
- Browser
- Download file 13.5 kB
-
https://huggingface.co/EagleDev-Innovation/Modelos-EagleDev-Innovation/resolve/main/entrenamiento/preparar.py
- Command line
-
hf download hf://EagleDev-Innovation/Modelos-EagleDev-Innovation/entrenamiento/preparar.py
-
curl -L -o preparar.py https://huggingface.co/EagleDev-Innovation/Modelos-EagleDev-Innovation/resolve/main/entrenamiento/preparar.py
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 | |
| 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 | |
| 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() | |