appUtils.py

#Importamos pytorch
import torch
#Transformadores de imagenes para preprocesar la entrada
from torchvision import transforms
#Cargadores de Datos que introducen los datos al modelo
from torch.utils.data import DataLoader
#OpenCV para recortar la imagen al final
import cv2 as cv

#Numpy para trabajar con arrays y matplotlib para ver las imagenes
import numpy as np
import matplotlib.pyplot as plt

#Para cargar las imagenes
from PIL import Image

#Bibliotecas de google
from libs.vit_seg_modeling import VisionTransformer as ViT_seg
from libs.vit_seg_modeling import CONFIGS as CONFIGS_ViT_seg

import os

import warnings
warnings.filterwarnings("ignore")

#############################################################################
#                                                                           #
#############################################################################

def get_conf():

    dictionary = {
        "INPUT_SIZE" : 256,
        "TRANSFORM" : transforms.Compose([transforms.Resize((256, 256)),transforms.ToTensor()])
    }

    return dictionary

class VisionTransformerModel(torch.nn.Module):
    def __init__(self, configs):
        super(VisionTransformerModel, self).__init__()
        
        self.model = ViT_seg(configs, img_size=get_conf()["INPUT_SIZE"], num_classes=1)
    
    def forward(self, x):
        img_segs = self.model(x)
        
        return img_segs

class TestDataset(torch.utils.data.Dataset):
    def __init__(self, images, transform):

        self.images = images
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = Image.open(self.images[idx]).convert("RGB")
        image = self.transform(image)
        return image
        
def tensor_to_image_3(tensor):
    img_channels = []
    for channel in tensor:
        tensor_np = channel.numpy()
        min_val = np.min(tensor_np)
        max_val = np.max(tensor_np)
        tensor_normalizado = (tensor_np - min_val) / (max_val - min_val)
        arraynp = (tensor_normalizado * 255).astype(np.uint8)
        img_channels.append(arraynp)

    imagen_rgb = np.stack(img_channels, axis=-1)
    imagen_rgb = imagen_rgb.astype(np.uint8)

    return Image.fromarray(imagen_rgb)

def tensor_to_image_1(tensor):
    tensor_np = tensor.numpy()
    min_val = np.min(tensor_np)
    max_val = np.max(tensor_np)
    tensor_normalizado = (tensor_np - min_val) / (max_val - min_val)
    arraynp = (tensor_normalizado * 255).astype(np.uint8)
    imagen_rgb = np.stack((arraynp,)*3, axis=-1)
    return Image.fromarray(imagen_rgb)

def cut_object(img, pred, threshold=60, background=[255, 255, 255]):
    color_inferior = np.array([0, 0, 0])
    color_superior = np.array([threshold, threshold, threshold])

    mascara = cv.inRange(pred, color_inferior, color_superior)
    nuevo_color = background

    img[mascara == 0] = nuevo_color
    return img

def highlight_object(img, pred):
    resta = cv.subtract(img, pred)
    return resta