models.py

import numpy as np
import torch
import torch.nn as nn

# =============================================================================
# APACHE 2.0 / https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/LICENSE
# - Credit: https://github.com/christophschuhmann/improved-aesthetic-predictor
# - - I reused the trained model
# =============================================================================
class MLP(nn.Module):
    def __init__(self, input_size):
        super().__init__()
        self.input_size = input_size
        self.layers = nn.Sequential(
            nn.Linear(self.input_size, 1024), nn.Dropout(0.2),
            nn.Linear(1024, 128), nn.Dropout(0.2),
            nn.Linear(128, 64), nn.Dropout(0.1),
            nn.Linear(64, 16),
            nn.Linear(16, 1))
    def forward(self, x):
        return self.layers(x)

# =============================================================================
# Misc
# =============================================================================
import torchvision.transforms.functional as FT

#normalize tensor
def normalized_pt(b, axis=-1):
    l2 = torch.norm(b)
    l2[l2 == 0] = 1
    return b / l2.reshape(1,1)

#get transforms to try
def get_transform_to_params(mode='soft'):
    if mode == 'hard':
        a = 0.5
        transform_to_params = {}
        transform_to_params[FT.adjust_brightness] = {'brightness_factor':((a, 1/a),)}
        transform_to_params[FT.adjust_contrast] = {'contrast_factor':((a, 1/a),)}
        transform_to_params[FT.adjust_gamma] = {'gamma':((a, 1/a),), 'gain':((0.9,1.1),)}
        transform_to_params[FT.adjust_hue] = {'hue_factor':((-0.1, 0.1),)}
        transform_to_params[FT.invert] = {}
        transform_to_params[FT.solarize] = {'threshold':((0, 1),)}
        transform_to_params[FT.posterize] = {'bits': (np.random.randint,(0,9))}
        transform_to_params[FT.adjust_saturation] = {'saturation_factor':((a,1/a),)}
        transform_to_params[FT.adjust_sharpness] = {'sharpness_factor':((a, 1/a),)}
        transform_to_params[FT.autocontrast] = {}
        # transform_to_params[FT.gaussian_blur] = {'kernel_size': (np.random.randint,(1, 5))}
        transform_to_params[FT.equalize] = {}
        print('- HARD MODE')
    elif mode == 'soft':
        a = 0.9
        transform_to_params = {}
        transform_to_params[FT.adjust_brightness] = {'brightness_factor':((a, 1/a),)}
        transform_to_params[FT.adjust_contrast] = {'contrast_factor':((a, 1/a),)}
        transform_to_params[FT.adjust_gamma] = {'gamma':((a, 1/a),), 'gain':((0.9,1.1),)}
        transform_to_params[FT.adjust_saturation] = {'saturation_factor':((a,1/a),)}
        transform_to_params[FT.adjust_sharpness] = {'sharpness_factor':((a, 1/a),)}
        transform_to_params[FT.autocontrast] = {}
        transform_to_params[FT.gaussian_blur] = {'kernel_size': (np.random.randint,(1, 3))}
        transform_to_params[FT.equalize] = {}
        print('- SOFT MODE')
    return transform_to_params