Spaces:

maripau22
/

Coloring2

Build error

Coloring2 / generator2.py

espero que ya color2

70764fd over 1 year ago

1.68 kB

import torch
import torch.nn as nn

class UNetGeneratorImproved(nn.Module):
    """
    UNet simplificado para tareas de colorización.
    Entrada: imagen en escala de grises (1 canal)
    Salida: imagen RGB colorizada (3 canales)
    """
    def __init__(self):
        super().__init__()

        # Codificador: reduce la resolución y captura características
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),  # 1 → 64 canales
            nn.ReLU(inplace=True),

            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # 64 → 128
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # 128 → 256
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )

        # Decodificador: recupera resolución y predice colores
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),  # 256 → 128
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # 128 → 64
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),  # 64 → 3 (RGB)
            nn.Tanh()  # Para que la salida esté en el rango [-1, 1] o lo normalices luego
        )

    def forward(self, x):
        """
        Propagación hacia adelante del modelo
        """
        x = self.encoder(x)  # Pasa por el codificador
        x = self.decoder(x)  # Reconstruye la imagen color
        return x

	import torch
	import torch.nn as nn

	class UNetGeneratorImproved(nn.Module):
	"""
	UNet simplificado para tareas de colorización.
	Entrada: imagen en escala de grises (1 canal)
	Salida: imagen RGB colorizada (3 canales)
	"""
	def __init__(self):
	super().__init__()

	# Codificador: reduce la resolución y captura características
	self.encoder = nn.Sequential(
	nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1), # 1 → 64 canales
	nn.ReLU(inplace=True),

	nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1), # 64 → 128
	nn.BatchNorm2d(128),
	nn.ReLU(inplace=True),

	nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1), # 128 → 256
	nn.BatchNorm2d(256),
	nn.ReLU(inplace=True),
	)

	# Decodificador: recupera resolución y predice colores
	self.decoder = nn.Sequential(
	nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 256 → 128
	nn.BatchNorm2d(128),
	nn.ReLU(inplace=True),

	nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 128 → 64
	nn.BatchNorm2d(64),
	nn.ReLU(inplace=True),

	nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1), # 64 → 3 (RGB)
	nn.Tanh() # Para que la salida esté en el rango [-1, 1] o lo normalices luego
	)

	def forward(self, x):
	"""
	Propagación hacia adelante del modelo
	"""
	x = self.encoder(x) # Pasa por el codificador
	x = self.decoder(x) # Reconstruye la imagen color
	return x