Detectar y segmentar: dónde está cada cosa

import urllib.request
from collections import Counter
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

COMMONS = "https://upload.wikimedia.org/wikipedia/commons/"
CALLE = "Avenida_Faria_Lima%2C_cruce_de_peatones.jpg"
VIDEO = ("Congested_traffic_on_the_Dan_Ryan_Expy_"
         "%2810x_timelapse%29_-_April_2026.webm")


def bajar(nombre, url):
    """Baja un archivo a «datos» si todavía no está y devuelve su ruta."""
    ruta = Path("datos") / nombre
    if not ruta.exists():
        ruta.parent.mkdir(parents=True, exist_ok=True)
        pedido = urllib.request.Request(
            url, headers={"User-Agent": "ruta-obia/1.0 (iabolivia.org)"})
        with urllib.request.urlopen(pedido) as respuesta:
            ruta.write_bytes(respuesta.read())
    return ruta


ruta_calle = bajar("imagenes-ruta/t27-calle.jpg",
                   f"{COMMONS}thumb/9/97/{CALLE}/960px-{CALLE}")
ruta_video = bajar("video-trafico/dan-ryan-480p.webm",
                   f"{COMMONS}transcoded/6/67/{VIDEO}/{VIDEO}.480p.vp9.webm")

foto = Image.open(ruta_calle).convert("RGB")
print("Tamaño de la foto (ancho, alto):", foto.size)

def xyxy_a_xywh(c):
    x1, y1, x2, y2 = c
    return np.array([x1, y1, x2 - x1, y2 - y1])


def xywh_a_xyxy(c):
    x, y, w, h = c
    return np.array([x, y, x + w, y + h])


# Una caja alrededor del auto plateado de la izquierda
caja = np.array([251, 536, 348, 629])
print("xyxy:", caja)
print("xywh:", xyxy_a_xywh(caja))
print("y de vuelta:", xywh_a_xyxy(xyxy_a_xywh(caja)))

def area(c):
    # Si la caja está «al revés» (no hay intersección), el área es 0
    return max(0, c[2] - c[0]) * max(0, c[3] - c[1])


def iou(a, b):
    """IoU de dos cajas en formato xyxy."""
    interseccion = area([max(a[0], b[0]), max(a[1], b[1]),
                         min(a[2], b[2]), min(a[3], b[3])])
    union = area(a) + area(b) - interseccion
    return interseccion / union


a = [0, 0, 10, 10]
b = [5, 0, 15, 10]     # la misma caja corrida 5 a la derecha
c = [20, 20, 30, 30]   # lejos de las dos
print(f"IoU(a, a) = {iou(a, a):.3f}")
print(f"IoU(a, b) = {iou(a, b):.3f}")
print(f"IoU(a, c) = {iou(a, c):.3f}")

from torchvision.ops import box_iou


def iou_matriz(A, B):
    """IoU de cada caja de A (n x 4) con cada caja de B (m x 4)."""
    x1 = np.maximum(A[:, None, 0], B[None, :, 0])
    y1 = np.maximum(A[:, None, 1], B[None, :, 1])
    x2 = np.minimum(A[:, None, 2], B[None, :, 2])
    y2 = np.minimum(A[:, None, 3], B[None, :, 3])
    interseccion = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
    area_a = (A[:, 2] - A[:, 0]) * (A[:, 3] - A[:, 1])
    area_b = (B[:, 2] - B[:, 0]) * (B[:, 3] - B[:, 1])
    return interseccion / (area_a[:, None] + area_b[None, :] - interseccion)


cajas = np.array([a, b, c], dtype=float)
print(np.round(iou_matriz(cajas, cajas), 3))
print("¿Igual que torchvision?", np.allclose(
    iou_matriz(cajas, cajas), box_iou(torch.tensor(cajas),
                                      torch.tensor(cajas)).numpy()))

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
for eje, t in zip(ejes.flat, [0.3, 0.5, 0.7, 0.9]):
    lado = np.sqrt(200 * t / (1 + t))     # lado de la intersección
    d = 10 - lado
    fija, corrida = [0, 0, 10, 10], [d, d, 10 + d, 10 + d]
    eje.add_patch(plt.Rectangle((d, d), lado, lado, color="C2", alpha=0.5))
    for c in [fija, corrida]:
        eje.add_patch(plt.Rectangle((c[0], c[1]), 10, 10, fill=False,
                                    edgecolor="C0", linewidth=2))
    eje.set_title(f"IoU = {iou(fija, corrida):.1f}")
    eje.set_xlim(-1, 16)
    eje.set_ylim(16, -1)
    eje.set_aspect("equal")
    eje.axis("off")
plt.show()

from torchvision.models import detection
from torchvision.transforms.functional import to_tensor

pesos_frcnn = detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT
frcnn = detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=pesos_frcnn)
frcnn = frcnn.to(device).eval()
pesos_ssd = detection.SSDLite320_MobileNet_V3_Large_Weights.DEFAULT
ssd = detection.ssdlite320_mobilenet_v3_large(weights=pesos_ssd)
ssd = ssd.to(device).eval()
# Los nombres de las clases de COCO, en inglés
CLASES = pesos_frcnn.meta["categories"]


def detectar(modelo, foto, umbral=0.5):
    """Cajas (xyxy, en píxeles), clases y puntajes por encima del umbral."""
    # Los modelos de torchvision reciben una lista de imágenes 0-1
    x = to_tensor(foto).to(device)
    with torch.no_grad():
        salida = modelo([x])[0]
    deja = salida["scores"] >= umbral
    cajas = salida["boxes"][deja].cpu().numpy()
    clases = [CLASES[i] for i in salida["labels"][deja].tolist()]
    return cajas, clases, salida["scores"][deja].cpu().numpy()


resultados = {}
for nombre, modelo in [("SSDlite", ssd), ("Faster R-CNN", frcnn)]:
    resultados[nombre] = detectar(modelo, foto)
    millones = sum(p.numel() for p in modelo.parameters()) / 1e6
    cajas, clases, puntajes = resultados[nombre]
    print(f"{nombre} ({millones:.1f} M parámetros): {len(cajas)} objetos",
          dict(Counter(clases)))

from transformers import AutoImageProcessor, AutoModelForObjectDetection

DETR = "facebook/detr-resnet-50"
procesador = AutoImageProcessor.from_pretrained(DETR)
detr = AutoModelForObjectDetection.from_pretrained(DETR).to(device).eval()


def detectar_detr(foto, umbral=0.7):
    """Lo mismo que detectar(), pero con DETR."""
    x = procesador(images=foto, return_tensors="pt").to(device)
    with torch.no_grad():
        salida = detr(**x)
    r = procesador.post_process_object_detection(
        salida, threshold=umbral, target_sizes=[foto.size[::-1]])[0]
    clases = [detr.config.id2label[i] for i in r["labels"].tolist()]
    return r["boxes"].cpu().numpy(), clases, r["scores"].cpu().numpy()


resultados["DETR"] = detectar_detr(foto)
millones = sum(p.numel() for p in detr.parameters()) / 1e6
cajas, clases, puntajes = resultados["DETR"]
print(f"DETR ({millones:.1f} M parámetros): {len(cajas)} objetos")
for clase, n in Counter(clases).most_common():
    print(f"  {clase}: {n}")

def dibujar_cajas(eje, cajas, color="C0", ancho=1.5):
    for x1, y1, x2, y2 in cajas:
        # Un borde blanco debajo para que la caja se vea en cualquier fondo
        for c, a in [("white", ancho + 1.5), (color, ancho)]:
            eje.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1,
                                        fill=False, edgecolor=c,
                                        linewidth=a))


fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
ejes[0, 0].imshow(foto)
ejes[0, 0].set_title("Foto")
for eje, nombre in zip(ejes.flat[1:], resultados):
    eje.imshow(foto)
    dibujar_cajas(eje, resultados[nombre][0], ancho=1)
    eje.set_title(f"{nombre}: {len(resultados[nombre][0])}")
for eje in ejes.flat:
    eje.axis("off")
plt.show()

frcnn.roi_heads.nms_thresh = 1.0
crudas, clases_crudas, puntajes_crudos = detectar(frcnn, foto)
frcnn.roi_heads.nms_thresh = 0.5          # el valor que trae
print(len(crudas), "cajas con puntaje de 0,5 o más antes de NMS")

def nms(cajas, puntajes, umbral_iou=0.5):
    """Índices de las cajas que sobreviven, de mayor a menor puntaje."""
    orden = list(np.argsort(-puntajes))     # de mayor a menor
    quedan = []
    while orden:
        mejor = orden.pop(0)
        quedan.append(mejor)
        # Sobreviven solo las que no se parecen demasiado a la mejor
        orden = [i for i in orden
                 if iou(cajas[mejor], cajas[i]) <= umbral_iou]
    return quedan


from torchvision.ops import nms as nms_torchvision

mias = nms(crudas, puntajes_crudos)
suyas = nms_torchvision(torch.tensor(crudas), torch.tensor(puntajes_crudos),
                        iou_threshold=0.5).tolist()
print("Quedan:", len(mias), "cajas:", [clases_crudas[i] for i in mias])
print("¿Igual que torchvision?", [int(i) for i in mias] == suyas)

fig, ejes = plt.subplots(2, 1, figsize=(4.8, 4.8))
for eje, cajas, titulo in [(ejes[0], crudas, "Antes de NMS: 15 cajas"),
                           (ejes[1], crudas[mias], "Después de NMS: 6")]:
    eje.imshow(foto)
    dibujar_cajas(eje, cajas, ancho=1)
    # Solo la parte de la foto donde están las cajas
    eje.set_xlim(220, 580)
    eje.set_ylim(790, 490)
    eje.set_title(titulo)
    eje.axis("off")
plt.show()

x = procesador(images=foto, return_tensors="pt").to(device)
with torch.no_grad():
    salida = detr(**x)
print("Imagen que entra:", tuple(x["pixel_values"].shape))
print("Clases por consulta:", tuple(salida.logits.shape))
print("Cajas por consulta:", tuple(salida.pred_boxes.shape))

# La última de las 92 columnas es «no hay objeto»
probabilidades = salida.logits.softmax(-1)[0]
nada = probabilidades[:, -1]
mejor_clase = probabilidades[:, :-1].max(dim=-1).values
print("Consultas que dicen «no hay objeto»:", int((nada >= 0.5).sum()))
print("Consultas con una clase de 0,7 o más:",
      int((mejor_clase >= 0.7).sum()))

from torchvision.ops import batched_nms

for nombre, (cajas, clases, puntajes) in [
        ("Faster R-CNN sin NMS", (crudas, clases_crudas, puntajes_crudos)),
        ("DETR", resultados["DETR"])]:
    # batched_nms necesita las clases como números
    codigos = torch.tensor([CLASES.index(c) for c in clases])
    quedan = batched_nms(torch.tensor(cajas), torch.tensor(puntajes),
                         codigos, iou_threshold=0.5)
    print(f"{nombre}: {len(cajas)} cajas, NMS deja {len(quedan)}")

cajas, clases, puntajes = detectar_detr(foto, umbral=0.3)
for umbral in [0.3, 0.5, 0.7, 0.9]:
    deja = puntajes >= umbral
    personas = sum(c == "person" for c, d in zip(clases, deja) if d)
    print(f"Umbral {umbral}: {deja.sum():2d} objetos, {personas} personas")

import cv2

VEHICULOS = {"car", "bus", "truck"}
video = cv2.VideoCapture(str(ruta_video))
fps = video.get(cv2.CAP_PROP_FPS)
tiempos, conteos = [], []
i = 0
while True:
    ok, cuadro = video.read()
    if not ok:                           # se acabó el video
        break
    if i % 10 == 0:
        rgb = cv2.cvtColor(cuadro, cv2.COLOR_BGR2RGB)
        cajas, clases, _ = detectar(frcnn, Image.fromarray(rgb))
        es_vehiculo = [c in VEHICULOS for c in clases]
        tiempos.append(i / fps)
        conteos.append(sum(es_vehiculo))
        if i == 180:                     # un cuadro para la figura
            ejemplo, cajas_ejemplo = rgb, cajas[es_vehiculo]
    i += 1
video.release()

print(f"{i} cuadros a {fps:.0f} por segundo ({i / fps:.1f} s); "
      f"analizamos {len(conteos)}")
print("Vehículos por cuadro: mínimo", min(conteos), "máximo", max(conteos),
      "mediana", np.median(conteos))

fig, ejes = plt.subplots(2, 1, figsize=(4.8, 4.8))
ejes[0].imshow(ejemplo)
dibujar_cajas(ejes[0], cajas_ejemplo, ancho=1)
ejes[0].set_title(f"Cuadro de los {180 / fps:.0f} s")
ejes[0].axis("off")
ejes[1].plot(tiempos, conteos, color="C0", marker="o", markersize=3)
ejes[1].plot(180 / fps, conteos[18], "o", color="C1", markersize=8)
ejes[1].set_xlabel("Tiempo (s)")
ejes[1].set_ylabel("Vehículos")
ejes[1].set_ylim(0, None)
plt.show()

LADO = 48


def generar_formas(n, rng):
    """n imágenes de 48 x 48 con círculos y cuadrados, y sus máscaras."""
    yy, xx = np.mgrid[0:LADO, 0:LADO]
    imagenes = np.zeros((n, LADO, LADO), dtype=np.float32)
    mascaras = np.zeros((n, LADO, LADO), dtype=np.int64)
    for i in range(n):
        # Fondo oscuro con ruido
        imagen = rng.uniform(0, 0.4) + rng.normal(0, 0.08, (LADO, LADO))
        for _ in range(rng.integers(1, 4)):        # de 1 a 3 formas
            clase = rng.integers(1, 3)             # 1 círculo, 2 cuadrado
            r = rng.uniform(4, 9)
            cy, cx = rng.uniform(r, LADO - r, size=2)
            if clase == 1:
                dentro = (yy - cy) ** 2 + (xx - cx) ** 2 <= r ** 2
            else:
                dentro = ((abs(yy - cy) <= 0.85 * r)
                          & (abs(xx - cx) <= 0.85 * r))
            # Cada forma con su propio brillo, también con ruido
            imagen[dentro] = (rng.uniform(0.5, 1)
                              + rng.normal(0, 0.08, dentro.sum()))
            mascaras[i][dentro] = clase
        imagenes[i] = imagen
    return imagenes, mascaras


rng = np.random.default_rng(42)
X_ent, Y_ent = generar_formas(1000, rng)
X_val, Y_val = generar_formas(200, rng)
print("Imágenes:", X_ent.shape, " máscaras:", Y_ent.shape)
for clase, nombre in enumerate(["fondo", "círculo", "cuadrado"]):
    print(f"Píxeles de {nombre}: {(Y_ent == clase).mean():.1%}")

from matplotlib.colors import LinearSegmentedColormap, ListedColormap

blanco_tinta = LinearSegmentedColormap.from_list("bt", ["white", "C0"])
# Fondo, círculo y cuadrado, cada uno con su color
colores_mascara = ListedColormap(["C3", "C0", "C2"])

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
for fila, k in enumerate([0, 1]):
    ejes[fila, 0].imshow(X_ent[k], cmap=blanco_tinta, vmin=0, vmax=1)
    ejes[fila, 1].imshow(Y_ent[k], cmap=colores_mascara, vmin=0, vmax=2)
ejes[0, 0].set_title("Imagen")
ejes[0, 1].set_title("Máscara")
for eje in ejes.flat:
    eje.axis("off")
plt.show()

from torch import nn
from torch.utils.data import DataLoader, TensorDataset


def doble_conv(entrada, salida):
    """Dos convoluciones 3 x 3, cada una con BatchNorm y ReLU."""
    return nn.Sequential(
        nn.Conv2d(entrada, salida, 3, padding=1), nn.BatchNorm2d(salida),
        nn.ReLU(),
        nn.Conv2d(salida, salida, 3, padding=1), nn.BatchNorm2d(salida),
        nn.ReLU())


class UNet(nn.Module):
    def __init__(self, clases=3):
        super().__init__()
        self.bajar = nn.MaxPool2d(2)
        # Codificador: 48 x 48, 24 x 24 y 12 x 12
        self.cod1 = doble_conv(1, 8)
        self.cod2 = doble_conv(8, 16)
        self.centro = doble_conv(16, 32)
        # Decodificador: sube y junta lo del codificador
        self.subir2 = nn.ConvTranspose2d(32, 16, 2, stride=2)
        self.dec2 = doble_conv(16 + 16, 16)
        self.subir1 = nn.ConvTranspose2d(16, 8, 2, stride=2)
        self.dec1 = doble_conv(8 + 8, 8)
        self.salida = nn.Conv2d(8, clases, 1)

    def forward(self, x):
        c1 = self.cod1(x)                       # (8, 48, 48)
        c2 = self.cod2(self.bajar(c1))          # (16, 24, 24)
        centro = self.centro(self.bajar(c2))    # (32, 12, 12)
        # Conexiones de salto: se pegan los canales del codificador
        d2 = self.dec2(torch.cat([self.subir2(centro), c2], dim=1))
        d1 = self.dec1(torch.cat([self.subir1(d2), c1], dim=1))
        return self.salida(d1)                  # (3, 48, 48)


torch.manual_seed(0)
unet = UNet().to(device)
x_val = torch.tensor(X_val).unsqueeze(1).to(device)   # (200, 1, 48, 48)
y_val = torch.tensor(Y_val).to(device)
print("Salida:", tuple(unet(x_val[:2]).shape))
print("Parámetros:", sum(p.numel() for p in unet.parameters()))

def predecir_mascaras(modelo, x):
    modelo.eval()
    with torch.no_grad():
        return modelo(x).argmax(dim=1)          # la clase con más puntaje


def iou_clases(pred, real, clases=3):
    """IoU de las máscaras de cada clase."""
    return [((pred == c) & (real == c)).sum().item()
            / ((pred == c) | (real == c)).sum().item()
            for c in range(clases)]


cargador = DataLoader(
    TensorDataset(torch.tensor(X_ent).unsqueeze(1), torch.tensor(Y_ent)),
    batch_size=32, shuffle=True, generator=torch.Generator().manual_seed(0))
optimizador = torch.optim.Adam(unet.parameters(), lr=3e-3)
perdida_pixel = nn.CrossEntropyLoss()

historia = {0: predecir_mascaras(unet, x_val[:1])[0].cpu()}
for epoca in range(1, 9):
    unet.train()
    suma = 0.0
    for x, y in cargador:
        x, y = x.to(device), y.to(device)
        optimizador.zero_grad()
        perdida = perdida_pixel(unet(x), y)
        perdida.backward()
        optimizador.step()
        suma += perdida.item() * len(x)
    pred = predecir_mascaras(unet, x_val)
    historia[epoca] = pred[0].cpu()
    fondo, circulo, cuadrado = iou_clases(pred, y_val)
    print(f"Época {epoca}: pérdida {suma / len(X_ent):.3f}, IoU círculo "
          f"{circulo:.2f}, cuadrado {cuadrado:.2f}, fondo {fondo:.2f}")

fig, ejes = plt.subplots(2, 3, figsize=(4.8, 3.6))
ejes[0, 0].imshow(X_val[0], cmap=blanco_tinta, vmin=0, vmax=1)
ejes[0, 0].set_title("Imagen", fontsize=12)
ejes[0, 1].imshow(Y_val[0], cmap=colores_mascara, vmin=0, vmax=2)
ejes[0, 1].set_title("Real", fontsize=12)
for eje, epoca in zip(ejes.flat[2:], [0, 1, 3, 8]):
    eje.imshow(historia[epoca], cmap=colores_mascara, vmin=0, vmax=2)
    eje.set_title(f"Época {epoca}", fontsize=12)
for eje in ejes.flat:
    eje.axis("off")
plt.show()

def dice_clase(pred, real, c):
    """Coeficiente de Dice de la clase c."""
    interseccion = ((pred == c) & (real == c)).sum().item()
    return 2 * interseccion / ((pred == c).sum().item()
                               + (real == c).sum().item())


pred = predecir_mascaras(unet, x_val)
print(f"Exactitud por píxel de la U-Net: "
      f"{(pred == y_val).float().mean():.3f}")
print(f"Exactitud de decir siempre «fondo»: {(y_val == 0).float().mean():.3f}")
for c, nombre in [(1, "círculo"), (2, "cuadrado")]:
    print(f"{nombre}: IoU {iou_clases(pred, y_val)[c]:.3f}, "
          f"Dice {dice_clase(pred, y_val, c):.3f}")

from torchvision import tv_tensors
from torchvision.transforms import v2

imagen = tv_tensors.Image(torch.tensor(X_val[:1]))
mascara = tv_tensors.Mask(torch.tensor(Y_val[0]))
caja = tv_tensors.BoundingBoxes([[10, 5, 20, 15]], format="XYXY",
                                canvas_size=(LADO, LADO))

voltear = v2.RandomHorizontalFlip(p=1)    # p=1: voltea siempre
imagen_v, mascara_v, caja_v = voltear(imagen, mascara, caja)
print("Caja a mano:", [LADO - 20, 5, LADO - 10, 15])
print("Caja con v2:", caja_v.tolist()[0])
print("¿La máscara se volteó igual que la imagen?",
      torch.equal(mascara_v, mascara.flip(-1)))