import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, random_split
from torch.optim.lr_scheduler import ReduceLROnPlateau
import cv2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class BeautyScore(nn.Module):
def __init__(self, first_neuron):
super(BeautyScore, self).__init__()
self.first_out_channels = first_neuron
self.features = nn.Sequential(
# First Convolutional Block
nn.Conv2d(in_channels=3, out_channels=self.first_out_channels, kernel_size=3, padding=1), # dimension [batch_size, out_channel, 128, `128`] -> padding = 1
nn.ReLU(),
nn.BatchNorm2d(self.first_out_channels),
nn.MaxPool2d(2), # dimension [batch_size, out_channel, 64, 64]
# Second Convolutional Block
nn.Conv2d(in_channels=self.first_out_channels, out_channels=self.first_out_channels*2, kernel_size=3, padding=1), # dimension [batch_size, out_channel, 32, 32]
nn.ReLU(),
nn.BatchNorm2d(self.first_out_channels*2),
nn.MaxPool2d(2), # dimension [batch_size, out_channel*2, 32, 32]
# Third Convolutional Block
nn.Conv2d(in_channels=self.first_out_channels*2, out_channels=self.first_out_channels*4, kernel_size=3, padding=1), # dimension [batch_size, out_channel, 16, 16]
nn.ReLU(),
nn.BatchNorm2d(self.first_out_channels*4),
nn.MaxPool2d(2), # dimension [batch_size, out_channel*4, 16, 16]
)
# Calculate size of flattened features after the convolutional layers
self.flatten_size = self.first_out_channels * 4 * (128 // (2**3)) * (128 // (2**3)) # out_channel * (128 // 2^amount_of_max_pool) * (128 // 2^amount_of_max_pool)
self.classifier = nn.Sequential(
nn.Dropout(0.3),
nn.Linear(self.flatten_size, 256), # dimension [batch_size, 256]
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(256, 128), # dimension [batch_size, 128]
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(128, 1), # dimension [batch_size, 1]
nn.Sigmoid() # To get value from 0 to 1
)
def forward(self, x):
x = self.features(x)
x = x.reshape(x.size(0), -1) # Flatten the tensor
x = self.classifier(x)
return x
class Trainer:
def __init__(self, train_loader = None, val_loader = None):
self.model = BeautyScore(first_neuron=256)
self.train_loader = train_loader
self.val_loader = val_loader
self.criterion = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001)
self.scheduler = ReduceLROnPlateau(self.optimizer, mode='min', factor=0.1, patience=5, verbose=True)
self.num_epochs = 20
def load_data(self):
data_path = '/home/reynaldy/.cache/kagglehub/datasets/pranavchandane/scut-fbp5500-v2-facial-beauty-scores/versions/2/scut_fbp5500-cmprsd.npz'
data = np.load(data_path)
data['X'].shape, data['y'].shape
features_numpy = data['X'].astype(np.float32)
features_numpy = np.array([cv2.resize(img, (128, 128)) for img in features_numpy]) # Resize the images to 256x256
features = torch.tensor(features_numpy, dtype=torch.float32).to(device)
features = features.permute(0, 3, 1, 2).to(device)
label_numpy = data['y'].astype(np.float32)
labels = torch.tensor(label_numpy, dtype=torch.float32).to(device)
tensor_min = labels.min()
tensor_max = labels.max()
labels = (labels - tensor_min) / (tensor_max - tensor_min)
print("Finish loading data")
train_size = int(0.8 * len(features))
test_size = len(features) - train_size
train_dataset, test_dataset = random_split(TensorDataset(features, labels), [train_size, test_size])
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
return train_loader, val_loader
def train(self):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(device)
self.model.train()
running_loss = 0.0
train_loader, _= self.load_data()
for batch_idx, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device).float()
self.optimizer.zero_grad()
outputs = self.model(inputs)
loss = self.criterion(outputs.squeeze(), labels)
loss.backward()
self.optimizer.step()
running_loss += loss.item()
if (batch_idx + 1) % 20 == 0:
print(f"Batch {batch_idx + 1}/{len(train_loader)} Loss: {loss.item()}")
epoch_loss = running_loss / len(train_loader)
if self.scheduler:
self.scheduler.step(epoch_loss)
print(f"Training Loss: {epoch_loss:.4f}")
return epoch_loss
def validate(self):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(device)
self.model.eval()
running_loss = 0.0
_, val_loader = self.load_data()
with torch.no_grad():
for batch_idx, (inputs, labels) in enumerate(val_loader):
inputs, labels = inputs.to(device), labels.to(device).float()
outputs = self.model(inputs)
loss = self.criterion(outputs.squeeze(), labels)
running_loss += loss.item()
epoch_loss = running_loss / len(val_loader)
print(f"Validation Loss: {epoch_loss:.4f}")
return epoch_loss
def image_to_tensor(self, image_path):
image = cv2.imread(image_path)
image = cv2.resize(image, (128, 128))
image = torch.tensor(image, dtype=torch.float32).to(device)
image = image.permute(2, 0, 1).unsqueeze(0)
return image
def predict(self, inputs):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.load_state_dict(torch.load('best_model.pth', weights_only=True, map_location=torch.device('cpu')))
self.model.to(device)
self.model.eval()
inputs = inputs.to(device)
with torch.no_grad():
outputs = self.model(inputs)
return outputs
# if __name__ == "__main__":
# trainer = Trainer()
# # Test the model
# image_path = '6082308423334085331.jpg'
# image_tensor = trainer.image_to_tensor(image_path)
# prediction = trainer.predict(image_tensor)
# print(f"Predicted Beauty Score: {prediction.item() * 100}")