helper.py

import torch
import torch.nn as nn
import numpy as np
import cv2
from torchvision import transforms,datasets
from PIL import Image
input_nc = 3
output_nc = 3
class SEBlock(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SEBlock, self).__init__()
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),  # Squeeze: output size (N, channel, 1, 1)
            nn.Conv2d(channel, channel // reduction, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channel // reduction, channel, 1),
            nn.Sigmoid()  # Excitation: channel weights between 0 and 1
        )
    
    def forward(self, x):
        weights = self.fc(x)
        return x * weights  # channel-wise multiplication

class ResnetBlock(nn.Module):
    def __init__(self, dim, reduction=16):
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim)
        self.se = SEBlock(dim, reduction)
    
    def build_conv_block(self, dim):
        conv_block = [
            nn.ReflectionPad2d(1),
            nn.Conv2d(dim, dim, kernel_size=3, padding=0),
            nn.InstanceNorm2d(dim),
            nn.ReLU(True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(dim, dim, kernel_size=3, padding=0),
            nn.InstanceNorm2d(dim)
        ]
        return nn.Sequential(*conv_block)
    
    def forward(self, x):
        out = self.conv_block(x)
        out = self.se(out)  # apply squeeze-and-excitation
        return x + out

class GeneratorResNet(nn.Module):
    def __init__(self, input_nc, output_nc, n_residual_blocks=9):
        super(GeneratorResNet, self).__init__()

        # Initial convolution block
        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, 64, 7),
                 nn.InstanceNorm2d(64),
                 nn.ReLU(inplace=True)]

        # Downsampling
        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                      nn.InstanceNorm2d(out_features),
                      nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features * 2

        # Residual blocks
        for _ in range(n_residual_blocks):
            model += [ResnetBlock(in_features)]

        # Upsampling
        out_features = in_features // 2
        for _ in range(2):
            model += [nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                      nn.InstanceNorm2d(out_features),
                      nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features // 2

        # Output layer
        model += [nn.ReflectionPad2d(3),
                  nn.Conv2d(64, output_nc, 7),
                  nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, x):
        return self.model(x)


netG_A2B = GeneratorResNet(input_nc, output_nc)
netG_B2A = GeneratorResNet(input_nc, output_nc)
# Load weights for netG_A2B
device = 'cpu'
netG_A2B.load_state_dict(torch.load('./netG_A2B_epoch130.pth',map_location=device))

# Load weights for netG_B2A
netG_B2A.load_state_dict(torch.load('./netG_B2A_epoch130.pth',map_location=device))

def generate_Y2O(uploaded_image):
    to_tensor = transforms.ToTensor()
    tensor = to_tensor(uploaded_image)
    old = netG_A2B(tensor)
    return (old.detach().permute(1, 2, 0).numpy()+1)/2


def generate_O2Y(uploaded_image):
    img = cv2.resize(uploaded_image, (256,256))
    to_tensor = transforms.ToTensor()
    tensor = to_tensor(img)
    young = netG_B2A(tensor)
    return (young.detach().permute(1, 2, 0).numpy()+1)/2


def extract_faces_opencv(image):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Load OpenCV's pre-trained Haar cascade
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")

    # Detect faces
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)

    face_crops = []
    for (x, y, w, h) in faces:
        # Apply padding safely
        y1, y2 = max(0, y - 50), min(image.shape[0], y + h)
        x1, x2 = max(0, x - 30), min(image.shape[1], x + w + 30)

        # Crop the face region
        face_crop = image[y1:y2, x1:x2]
        face_crops.append(face_crop)
    return face_crops