Intermediate Project: images only with NumPy

"""Leitura e gravação de PNG com a biblioteca padrão e NumPy.

Grava imagens de 8 bits em tons de cinza ou RGB, sem filtro por linha (filtro 0). Lê apenas
esse mesmo formato, ou seja, os arquivos gravados por este módulo.
"""

import struct
import zlib
from pathlib import Path

import numpy as np

from imagens.tipos import Bytes

ASSINATURA = b"\x89PNG\r\n\x1a\n"


def _bloco(tipo: bytes, dados: bytes) -> bytes:
    return (
        struct.pack(">I", len(dados)) + tipo + dados + struct.pack(">I", zlib.crc32(tipo + dados))
    )


def salvar_png(caminho: Path, img: Bytes) -> None:
    if img.dtype != np.uint8:
        raise ValueError("a imagem precisa ser uint8")
    if img.ndim == 2:
        tipo_cor = 0
    elif img.ndim == 3 and img.shape[2] == 3:
        tipo_cor = 2
    else:
        raise ValueError("use uma imagem 2D (cinza) ou (altura, largura, 3) (RGB)")
    altura, largura = img.shape[:2]
    # Cada linha começa com um byte de filtro (0 = sem filtro).
    linhas = np.concatenate(
        [np.zeros((altura, 1), dtype=np.uint8), img.reshape(altura, -1)], axis=1
    )
    cabecalho = struct.pack(">IIBBBBB", largura, altura, 8, tipo_cor, 0, 0, 0)
    conteudo = (
        ASSINATURA
        + _bloco(b"IHDR", cabecalho)
        + _bloco(b"IDAT", zlib.compress(linhas.tobytes(), 9))
        + _bloco(b"IEND", b"")
    )
    caminho.write_bytes(conteudo)


def ler_png(caminho: Path) -> Bytes:
    dados = caminho.read_bytes()
    if not dados.startswith(ASSINATURA):
        raise ValueError("não é um arquivo PNG")
    posicao = len(ASSINATURA)
    largura = altura = tipo_cor = 0
    comprimidos = b""
    while posicao < len(dados):
        (tamanho,) = struct.unpack(">I", dados[posicao : posicao + 4])
        tipo = dados[posicao + 4 : posicao + 8]
        corpo = dados[posicao + 8 : posicao + 8 + tamanho]
        if tipo == b"IHDR":
            largura, altura, profundidade, tipo_cor, *_ = struct.unpack(">IIBBBBB", corpo)
            if profundidade != 8 or tipo_cor not in (0, 2):
                raise ValueError("só leio PNG de 8 bits em cinza ou RGB")
        elif tipo == b"IDAT":
            comprimidos += corpo
        posicao += 12 + tamanho
    canais = 1 if tipo_cor == 0 else 3
    bruto = np.frombuffer(zlib.decompress(comprimidos), dtype=np.uint8)
    linhas = bruto.reshape(altura, 1 + largura * canais)
    if linhas[:, 0].any():
        raise ValueError("só leio PNG gravado sem filtro por linha (como o deste módulo)")
    pixels = linhas[:, 1:]
    return pixels.reshape(altura, largura) if canais == 1 else pixels.reshape(altura, largura, 3)

import numpy as np

from imagens.tipos import Bytes


def criar_cena(tamanho: int = 128, semente: int = 7) -> Bytes:
    """Imagem RGB de teste: fundo em gradiente, um círculo claro, um retângulo escuro e ruído."""
    rng = np.random.default_rng(semente)
    y, x = np.mgrid[0:tamanho, 0:tamanho]
    cena = np.zeros((tamanho, tamanho, 3), dtype=np.float64)
    cena[..., 0] = 60 + 80 * x / (tamanho - 1)
    cena[..., 1] = 60 + 80 * y / (tamanho - 1)
    cena[..., 2] = 90

    centro = tamanho * 0.35
    circulo = (x - centro) ** 2 + (y - centro) ** 2 <= (tamanho * 0.2) ** 2
    cena[circulo] = (230, 200, 60)
    retangulo = (x > tamanho * 0.55) & (x < tamanho * 0.9)
    retangulo &= (y > tamanho * 0.5) & (y < tamanho * 0.85)
    cena[retangulo] = (40, 60, 200)

    cena += rng.normal(0, 6, cena.shape)
    return np.clip(cena, 0, 255).round().astype(np.uint8)

import numpy as np
import numpy.typing as npt

from imagens.tipos import Bytes

PESOS_CINZA = np.array([0.299, 0.587, 0.114])


def para_cinza(rgb: Bytes) -> Bytes:
    """Média ponderada dos canais, com os pesos da percepção humana (verde pesa mais)."""
    cinza: Bytes = np.clip((rgb[..., :3] @ PESOS_CINZA).round(), 0, 255).astype(np.uint8)
    return cinza


def esticar_contraste(img: Bytes) -> Bytes:
    """Leva o menor valor a 0 e o maior a 255, mantendo a proporção entre os demais."""
    menor, maior = int(img.min()), int(img.max())
    if maior == menor:
        return img.copy()
    esticada = (img.astype(np.float64) - menor) * (255 / (maior - menor))
    return np.clip(esticada.round(), 0, 255).astype(np.uint8)


def histograma(img: Bytes) -> npt.NDArray[np.intp]:
    return np.bincount(img.ravel(), minlength=256)


def equalizar(img: Bytes) -> Bytes:
    """Redistribui os níveis para o histograma ficar o mais uniforme possível."""
    cdf = histograma(img).cumsum()
    cdf_min = int(cdf[np.flatnonzero(histograma(img))[0]])
    if cdf[-1] == cdf_min:
        return img.copy()
    tabela = np.clip(np.round((cdf - cdf_min) / (cdf[-1] - cdf_min) * 255), 0, 255).astype(np.uint8)
    equalizada: Bytes = tabela[img]
    return equalizada


def limiar_otsu(img: Bytes) -> int:
    """Limiar que separa a imagem em dois grupos com a maior variância entre eles (Otsu)."""
    prob = histograma(img) / img.size
    niveis = np.arange(256)
    omega = np.cumsum(prob)
    mu = np.cumsum(prob * niveis)
    with np.errstate(divide="ignore", invalid="ignore"):
        variancia_entre = (mu[-1] * omega - mu) ** 2 / (omega * (1 - omega))
    return int(np.argmax(np.nan_to_num(variancia_entre)))


def binarizar(img: Bytes, limiar: int) -> Bytes:
    return np.where(img > limiar, 255, 0).astype(np.uint8)

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

from imagens.tipos import Bytes, Reais


def aplicar_kernel(img: Reais, kernel: Reais) -> Reais:
    """Correlação 2D: cada pixel vira a soma ponderada da sua vizinhança.

    As bordas são tratadas repetindo o pixel mais próximo (``mode="edge"``). As janelas são uma
    visão (sem cópia) da imagem, e o ``einsum`` multiplica cada uma pelo kernel de uma vez.
    """
    altura_k, largura_k = kernel.shape
    if altura_k % 2 == 0 or largura_k % 2 == 0:
        raise ValueError("o kernel precisa ter dimensões ímpares")
    margem = ((altura_k // 2, altura_k // 2), (largura_k // 2, largura_k // 2))
    janelas = sliding_window_view(np.pad(img, margem, mode="edge"), (altura_k, largura_k))
    resultado: Reais = np.einsum("ijkl,kl->ij", janelas, kernel)
    return resultado


def kernel_caixa(k: int) -> Reais:
    return np.full((k, k), 1.0 / (k * k))


def kernel_gaussiano(k: int, sigma: float) -> Reais:
    eixo = np.arange(k) - k // 2
    g = np.exp(-(eixo**2) / (2 * sigma**2))
    kernel: Reais = np.outer(g, g)
    return kernel / kernel.sum()


SOBEL_X = np.array([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]])


def bordas_sobel(img: Bytes) -> Reais:
    """Magnitude do gradiente: grande onde o brilho muda de repente (as bordas)."""
    f = img.astype(np.float64)
    gx = aplicar_kernel(f, SOBEL_X)
    gy = aplicar_kernel(f, SOBEL_X.T)
    return np.hypot(gx, gy)


def para_uint8(img: Reais) -> Bytes:
    """Reescala qualquer array decimal para 0..255, para poder gravar como imagem."""
    menor, maior = float(img.min()), float(img.max())
    if maior == menor:
        return np.zeros(img.shape, dtype=np.uint8)
    return np.clip((img - menor) / (maior - menor) * 255, 0, 255).round().astype(np.uint8)

import argparse
from pathlib import Path

import numpy as np

from imagens.filtros import aplicar_kernel, bordas_sobel, kernel_gaussiano, para_uint8
from imagens.pngio import ler_png, salvar_png
from imagens.sintetica import criar_cena
from imagens.tons import binarizar, equalizar, esticar_contraste, limiar_otsu, para_cinza


def criar_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="imagens", description="Processamento de imagens")
    sub = parser.add_subparsers(dest="comando", required=True)
    gerar = sub.add_parser("gerar", help="grava uma imagem de teste")
    gerar.add_argument("pasta", type=Path)
    gerar.add_argument("--tamanho", type=int, default=128)
    processar = sub.add_parser("processar", help="aplica uma sequência de operações")
    processar.add_argument("entrada", type=Path)
    processar.add_argument("pasta", type=Path)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = criar_parser().parse_args(argv)
    args.pasta.mkdir(parents=True, exist_ok=True)
    if args.comando == "gerar":
        destino = args.pasta / "cena.png"
        salvar_png(destino, criar_cena(args.tamanho))
        print(f"imagem gravada em {destino}")
        return 0

    try:
        rgb = ler_png(args.entrada)
    except (OSError, ValueError) as erro:
        print(f"Erro: {erro}")
        return 1
    cinza = para_cinza(rgb) if rgb.ndim == 3 else rgb
    suave = aplicar_kernel(cinza.astype(float), kernel_gaussiano(5, 1.2))
    desfocada = np.clip(suave.round(), 0, 255).astype(np.uint8)
    limiar = limiar_otsu(cinza)
    etapas = {
        "cinza": cinza,
        "contraste": esticar_contraste(cinza),
        "equalizada": equalizar(cinza),
        "desfoque": desfocada,
        "bordas": para_uint8(bordas_sobel(desfocada)),
        "binaria": binarizar(cinza, limiar),
    }
    for nome, imagem in etapas.items():
        salvar_png(args.pasta / f"{nome}.png", imagem)
        minimo, maximo, media = int(imagem.min()), int(imagem.max()), imagem.mean()
        print(f"{nome:<11} min={minimo:>3} max={maximo:>3} média={media:6.1f}")
    original = rgb if rgb.ndim == 3 else np.stack([rgb] * 3, axis=-1)
    nomes = ("cinza", "equalizada", "desfoque", "bordas", "binaria")
    tiras = [np.stack([etapas[nome]] * 3, axis=-1) for nome in nomes]
    salvar_png(args.pasta / "montagem.png", np.hstack([original, *tiras]))
    print("montagem.png: original, " + ", ".join(nomes))
    print(f"limiar de Otsu: {limiar}")
    return 0

from pathlib import Path

import numpy as np
import pytest

from imagens.pngio import ASSINATURA, ler_png, salvar_png


def test_ida_e_volta_em_cinza_e_em_rgb(tmp_path: Path) -> None:
    rng = np.random.default_rng(0)
    cinza = rng.integers(0, 256, (7, 9), dtype=np.uint8)
    rgb = rng.integers(0, 256, (5, 6, 3), dtype=np.uint8)
    salvar_png(tmp_path / "c.png", cinza)
    salvar_png(tmp_path / "r.png", rgb)
    assert np.array_equal(ler_png(tmp_path / "c.png"), cinza)
    assert np.array_equal(ler_png(tmp_path / "r.png"), rgb)


def test_o_arquivo_tem_a_assinatura_do_png(tmp_path: Path) -> None:
    salvar_png(tmp_path / "x.png", np.zeros((2, 2), dtype=np.uint8))
    assert (tmp_path / "x.png").read_bytes().startswith(ASSINATURA)


def test_entradas_invalidas(tmp_path: Path) -> None:
    with pytest.raises(ValueError, match="uint8"):
        salvar_png(tmp_path / "a.png", np.zeros((2, 2)))
    with pytest.raises(ValueError, match="RGB"):
        salvar_png(tmp_path / "b.png", np.zeros((2, 2, 4), dtype=np.uint8))
    (tmp_path / "texto.png").write_text("não sou uma imagem", encoding="utf-8")
    with pytest.raises(ValueError, match="PNG"):
        ler_png(tmp_path / "texto.png")

import numpy as np
import pytest

from imagens.filtros import (
    aplicar_kernel,
    bordas_sobel,
    kernel_caixa,
    kernel_gaussiano,
    para_uint8,
)


def aplicar_por_laco(img: np.ndarray, kernel: np.ndarray) -> np.ndarray:  # type: ignore[type-arg]
    """Versão lenta e obviamente correta, usada para validar a vetorizada."""
    k = kernel.shape[0] // 2
    padded = np.pad(img, k, mode="edge")
    saida = np.zeros_like(img)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            saida[i, j] = (padded[i : i + 2 * k + 1, j : j + 2 * k + 1] * kernel).sum()
    return saida


def test_confere_com_a_versao_por_laco() -> None:
    rng = np.random.default_rng(0)
    img = rng.random((6, 7))
    kernel = rng.random((3, 3))
    assert np.allclose(aplicar_kernel(img, kernel), aplicar_por_laco(img, kernel))


def test_kernel_de_tamanho_par_e_recusado() -> None:
    with pytest.raises(ValueError, match="ímpares"):
        aplicar_kernel(np.zeros((4, 4)), np.ones((2, 2)))


def test_desfoque_preserva_imagem_constante_e_reduz_ruido() -> None:
    assert np.allclose(aplicar_kernel(np.full((5, 5), 7.0), kernel_caixa(3)), 7.0)
    ruido = np.random.default_rng(1).normal(0, 10, (40, 40))
    assert aplicar_kernel(ruido, kernel_gaussiano(5, 1.2)).std() < 0.5 * ruido.std()


def test_kernel_gaussiano_soma_1_e_e_simetrico() -> None:
    k = kernel_gaussiano(5, 1.0)
    assert np.isclose(k.sum(), 1.0)
    assert np.allclose(k, k.T) and np.allclose(k, k[::-1, ::-1])


def test_sobel_acha_a_borda_vertical() -> None:
    img = np.zeros((9, 9), dtype=np.uint8)
    img[:, 5:] = 255
    mag = bordas_sobel(img)
    assert set(np.argmax(mag, axis=1)) <= {4, 5}
    assert np.allclose(mag[:, :3], 0) and np.allclose(mag[:, 7:], 0)


def test_para_uint8_reescala_para_0_255() -> None:
    r = para_uint8(np.array([[-5.0, 0.0], [5.0, 10.0]]))
    assert r.dtype == np.uint8 and r.min() == 0 and r.max() == 255
    assert para_uint8(np.ones((2, 2))).max() == 0

import numpy as np

from imagens.tons import (
    binarizar,
    equalizar,
    esticar_contraste,
    histograma,
    limiar_otsu,
    para_cinza,
)


def test_cinza_de_cores_conhecidas() -> None:
    vermelho = np.array([[[255, 0, 0]]], dtype=np.uint8)
    branco = np.array([[[255, 255, 255]]], dtype=np.uint8)
    assert para_cinza(vermelho)[0, 0] == 76
    assert para_cinza(branco)[0, 0] == 255


def test_esticar_contraste_leva_aos_extremos() -> None:
    img = np.array([[100, 120], [140, 150]], dtype=np.uint8)
    r = esticar_contraste(img)
    assert r.min() == 0 and r.max() == 255
    assert np.array_equal(esticar_contraste(np.full((2, 2), 9, dtype=np.uint8)), np.full((2, 2), 9))


def test_histograma_soma_o_numero_de_pixels() -> None:
    img = np.random.default_rng(0).integers(0, 256, (10, 10), dtype=np.uint8)
    assert histograma(img).sum() == img.size


def test_equalizar_aumenta_a_dispersao_de_uma_imagem_de_baixo_contraste() -> None:
    rng = np.random.default_rng(2)
    img = rng.integers(100, 140, (40, 40), dtype=np.uint8)
    r = equalizar(img)
    assert r.std() > 2 * img.std()
    assert r.min() == 0 and r.max() == 255


def test_otsu_separa_duas_populacoes() -> None:
    rng = np.random.default_rng(3)
    claro = np.clip(rng.normal(200, 10, 500), 0, 255)
    escuro = np.clip(rng.normal(50, 10, 500), 0, 255)
    img = np.concatenate([claro, escuro]).astype(np.uint8).reshape(25, 40)
    assert 80 < limiar_otsu(img) < 170
    assert set(np.unique(binarizar(img, limiar_otsu(img)))) == {0, 255}

from pathlib import Path

import pytest

from imagens.cli import main


def test_gerar_e_processar_de_ponta_a_ponta(
    tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
    assert main(["gerar", str(tmp_path), "--tamanho", "64"]) == 0
    assert main(["processar", str(tmp_path / "cena.png"), str(tmp_path / "saida")]) == 0
    saida = capsys.readouterr().out
    assert "limiar de Otsu" in saida
    nomes = {p.name for p in (tmp_path / "saida").glob("*.png")}
    assert nomes == {
        "cinza.png",
        "contraste.png",
        "equalizada.png",
        "desfoque.png",
        "bordas.png",
        "binaria.png",
        "montagem.png",
    }


def test_arquivo_que_nao_e_png_devolve_codigo_1(
    tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
    falso = tmp_path / "falso.png"
    falso.write_text("nada", encoding="utf-8")
    assert main(["processar", str(falso), str(tmp_path / "s")]) == 1
    assert "Erro" in capsys.readouterr().out

cd projetos/imagens
uv sync
uv run pytest
uv run imagens gerar saida
uv run imagens processar saida/cena.png saida

imagem gravada em saida/cena.png
cinza       min= 52 max=205 média= 106.8
contraste   min=  0 max=255 média=  91.3
equalizada  min=  0 max=255 média= 129.2
desfoque    min= 60 max=196 média= 106.8
bordas      min=  0 max=255 média=  14.8
binaria     min=  0 max=255 média=  32.1
montagem.png: original, cinza, equalizada, desfoque, bordas, binaria
limiar de Otsu: 143