Spaces:

morinousagi
/

pytorch-vit-defect-detect

Sleeping

pytorch-vit-defect-detect / src /model.py

Upload model.py

ac8a72f verified 4 months ago

1.57 kB

import torch.nn as nn
from transformers import ViTModel

class SEMViTAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Load pre-trained ViT-Base as encoder
        self.encoder = ViTModel.from_pretrained(
            'google/vit-base-patch16-224',  # ViT expects 224p, 3 channels (RGB)
            add_pooling_layer=False 
        )
        
               
        # Custom convolutional decoder to reconstruct the image
        # Decoder: Mapping ViT features (768) back to image space
        # 512/16 = 32 patches. Latent grid is 32x32.
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(768, 256, kernel_size=4, stride=2, padding=1), # 64x64
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 128x128
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 256x256
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),   # 512x512
            nn.Sigmoid() # Output pixels in [0, 1]
        )

    def forward(self, x):
        # ViT requires position interpolation for resolutions != 224
        outputs = self.encoder(x, interpolate_pos_encoding=True)
        # Sequence to Grid: [Batch, 1025, 768] -> [Batch, 768, 32, 32]
        latent = outputs.last_hidden_state[:, 1:, :].transpose(1, 2).reshape(-1, 768, 32, 32)
        
        return self.decoder(latent)

	import torch.nn as nn
	from transformers import ViTModel

	class SEMViTAutoencoder(nn.Module):
	def __init__(self):
	super().__init__()

	# Load pre-trained ViT-Base as encoder
	self.encoder = ViTModel.from_pretrained(
	'google/vit-base-patch16-224', # ViT expects 224p, 3 channels (RGB)
	add_pooling_layer=False
	)


	# Custom convolutional decoder to reconstruct the image
	# Decoder: Mapping ViT features (768) back to image space
	# 512/16 = 32 patches. Latent grid is 32x32.
	self.decoder = nn.Sequential(
	nn.ConvTranspose2d(768, 256, kernel_size=4, stride=2, padding=1), # 64x64
	nn.BatchNorm2d(256),
	nn.ReLU(),
	nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 128x128
	nn.BatchNorm2d(128),
	nn.ReLU(),
	nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 256x256
	nn.BatchNorm2d(64),
	nn.ReLU(),
	nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1), # 512x512
	nn.Sigmoid() # Output pixels in [0, 1]
	)

	def forward(self, x):
	# ViT requires position interpolation for resolutions != 224
	outputs = self.encoder(x, interpolate_pos_encoding=True)
	# Sequence to Grid: [Batch, 1025, 768] -> [Batch, 768, 32, 32]
	latent = outputs.last_hidden_state[:, 1:, :].transpose(1, 2).reshape(-1, 768, 32, 32)

	return self.decoder(latent)