Add files from Neural-IMage-Assessment repo

.gitattributes

@@ -32,3 +32,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+snapshots/badpred.png filter=lfs diff=lfs merge=lfs -text
+snapshots/contrast.png filter=lfs diff=lfs merge=lfs -text
+snapshots/goodpred.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+./data/

LICENSE

@@ -0,0 +1,8 @@
+The MIT License (MIT) Copyright (c) 2020 Yunxiao Shi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+© 2020 GitHub, Inc.

README.md

@@ -1,12 +1,81 @@
----
-title: NIMA
-emoji: 👁
-colorFrom: red
-colorTo: indigo
-sdk: gradio
-sdk_version: 3.19.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## NIMA: Neural IMage Assessment
+
+[![Python 3.6+](https://img.shields.io/badge/Python-3.6%2B-blue)](https://www.python.org/)
+[![MIT License](https://img.shields.io/badge/MIT-License-brightgreen)](./LICENSE)
+
+This is a PyTorch implementation of the paper [NIMA: Neural IMage Assessment](https://arxiv.org/abs/1709.05424) (accepted at [IEEE Transactions on Image Processing](https://ieeexplore.ieee.org/document/8352823)) by Hossein Talebi and Peyman Milanfar. You can learn more from [this post at Google Research Blog](https://research.googleblog.com/2017/12/introducing-nima-neural-image-assessment.html).
+
+## Implementation Details
+
++ The model was trained on the [AVA (Aesthetic Visual Analysis) dataset](http://refbase.cvc.uab.es/files/MMP2012a.pdf) containing 255,500+ images. You can get it from [here](https://github.com/mtobeiyf/ava_downloader). ~~**Note: there may be some corrupted images in the dataset, remove them first before you start training**.~~ Use provided CSVs which have already done this for you.
+
++ Dataset is split into 229,981 images for training, 12,691 images for validation and 12,818 images for testing. 
+
++ An ImageNet pretrained VGG-16 is used as the base network. Should be easy to plug in the other two options (MobileNet and Inception-v2).
+
++ The learning rate setting differs from the original paper. Can't seem to get the model to converge using the original params. Also didn't do much hyper-param tuning therefore you could probably get better results. Other settings are all directly mirrored from the paper.
+
+## Requirements
+
+Code is written using [PyTorch](https://pytorch.org/get-started/locally/) 1.8.1 with [CUDA](https://developer.nvidia.com/cuda-toolkit) 11.1. You can recreate the environment I used with [conda](https://docs.conda.io/en/latest/miniconda.html) by
+```
+conda env create -f env.yml
+```  
+to install the dependancies.
+
+## Usage
+
+To start training on the AVA dataset, first download the dataset from the link above and decompress which should create a directory named ```images/```. Then download the curated annotation CSVs below
+which already splits the dataset (You can create your own split of course). Then do
+
+```python
+python main.py --img_path /path/to/images/ --train --train_csv_file /path/to/train_labels.csv --val_csv_file /path/to/val_labels.csv --conv_base_lr 5e-4 --dense_lr 5e-3 --decay --ckpt_path /path/to/ckpts --epochs 100 --early_stoppping_patience 10
+```
+
+For inference, do
+
+```python
+python -W ignore test.py --model /path/to/your_model --test_csv /path/to/test_labels.csv --test_images /path/to/images --predictions /path/to/save/predictions
+```
+
+See ```predictions/``` for dumped predictions as an example.
+
+## Training Statistics
+
+Training is done with early stopping. Here I set ```early_stopping_patience=10```.
+<p align="center">
+<img src="./snapshots/[email protected]">
+</p>
+
+## Pretrained Model
+
+~0.069 EMD on validation. Not fully converged yet (constrained by resources). To continue training, download the pretrained weights and add ```--warm_start --warm_start_epoch 34``` to your args.
+
+[Google Drive](https://drive.google.com/file/d/1w9Ig_d6yZqUZSR63kPjZLrEjJ1n845B_/view?usp=sharing)
+
+## Annotation CSV Files
+[Train](https://drive.google.com/file/d/1IBXPXPkCiTz04wWcoReJv4Nk06VsjSkI/view?usp=sharing) [Validation](https://drive.google.com/file/d/1tJfO1zFBoQYzd8kUo5PKeHTcdzBL7115/view?usp=sharing) [Test](https://drive.google.com/file/d/105UGnkglpKuusPhJaPnFSa2JlQV3du9O/view?usp=sharing)
+
+## Example Results
+
++ Here first shows some good predictions from the test set. Each image title starts with ground-truth rating followed by the predicted mean and std in the parentheses. 
+
+<p align="center">
+<img src="./snapshots/goodpred.png">
+</p>
+
++ Also some failure cases, it would seem that the model usually fails at images with low/high aesthetic ratings.
+
+<p align="center">
+<img src="./snapshots/badpred.png">
+</p>
+
++ The predicted aesthetic ratings from training on the AVA dataset are sensitive to contrast adjustments, preferring images with higher contrast. Below top row is the reference image with contrast ```c=1.0```,  while bottom images are enhanced with contrast ```[0.25, 0.75, 1.25, 1.75]```. Contrast adjustment is done using ```ImageEnhance.Contrast``` from ```PIL``` (in this case [pillow-simd](https://github.com/uploadcare/pillow-simd)). 
+
+<p align="center">
+<img src="./snapshots/contrast.png">
+</p>
+
+## License
+
+MIT

dataset/dataset.py

@@ -0,0 +1,66 @@
+"""
+file - dataset.py
+Customized dataset class to loop through the AVA dataset and apply needed image augmentations for training.
+
+Copyright (C) Yunxiao Shi 2017 - 2021
+NIMA is released under the MIT license. See LICENSE for the fill license text.
+"""
+
+import os
+
+import pandas as pd
+from PIL import Image
+
+import torch
+from torch.utils import data
+import torchvision.transforms as transforms
+
+
+class AVADataset(data.Dataset):
+    """AVA dataset
+
+    Args:
+        csv_file: a 11-column csv_file, column one contains the names of image files, column 2-11 contains the empiricial distributions of ratings
+        root_dir: directory to the images
+        transform: preprocessing and augmentation of the training images
+    """
+
+    def __init__(self, csv_file, root_dir, transform=None):
+        self.annotations = pd.read_csv(csv_file)
+        self.root_dir = root_dir
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx):
+        img_name = os.path.join(self.root_dir, str(self.annotations.iloc[idx, 0]) + '.jpg')
+        image = Image.open(img_name).convert('RGB')
+        annotations = self.annotations.iloc[idx, 1:].to_numpy()
+        annotations = annotations.astype('float').reshape(-1, 1)
+        sample = {'img_id': img_name, 'image': image, 'annotations': annotations}
+
+        if self.transform:
+            sample['image'] = self.transform(sample['image'])
+
+        return sample
+
+
+if __name__ == '__main__':
+
+    # sanity check
+    root = './data/images'
+    csv_file = './data/train_labels.csv'
+    train_transform = transforms.Compose([
+        transforms.Scale(256), 
+        transforms.RandomCrop(224), 
+        transforms.RandomHorizontalFlip(), 
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+    dset = AVADataset(csv_file=csv_file, root_dir=root, transform=train_transform)
+    train_loader = data.DataLoader(dset, batch_size=4, shuffle=True, num_workers=4)
+    for i, data in enumerate(train_loader):
+        images = data['image']
+        print(images.size())
+        labels = data['annotations']
+        print(labels.size())

env.yml

@@ -0,0 +1,106 @@
+name: torch1.8.1
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - blas=1.0=mkl
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2021.5.25=h06a4308_1
+  - certifi=2021.5.30=py38h06a4308_0
+  - cudatoolkit=11.1.74=h6bb024c_0
+  - cycler=0.10.0=py38_0
+  - dbus=1.13.18=hb2f20db_0
+  - expat=2.4.1=h2531618_2
+  - ffmpeg=4.3=hf484d3e_0
+  - fontconfig=2.13.1=h6c09931_0
+  - freetype=2.10.4=h5ab3b9f_0
+  - glib=2.68.2=h36276a3_0
+  - gmp=6.2.1=h2531618_2
+  - gnutls=3.6.15=he1e5248_0
+  - gst-plugins-base=1.14.0=h8213a91_2
+  - gstreamer=1.14.0=h28cd5cc_2
+  - icu=58.2=he6710b0_3
+  - intel-openmp=2021.2.0=h06a4308_610
+  - jpeg=9b=h024ee3a_2
+  - kiwisolver=1.3.1=py38h2531618_0
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.33.1=h53a641e_7
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libiconv=1.15=h63c8f33_5
+  - libidn2=2.3.1=h27cfd23_0
+  - libpng=1.6.37=hbc83047_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.2.0=h85742a9_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.0.3=h1bed415_2
+  - libuv=1.40.0=h7b6447c_0
+  - libwebp-base=1.2.0=h27cfd23_0
+  - libxcb=1.14=h7b6447c_0
+  - libxml2=2.9.10=hb55368b_3
+  - lz4-c=1.9.3=h2531618_0
+  - matplotlib=3.3.4=py38h06a4308_0
+  - matplotlib-base=3.3.4=py38h62a2d02_0
+  - mkl=2021.2.0=h06a4308_296
+  - mkl-service=2.3.0=py38h27cfd23_1
+  - mkl_fft=1.3.0=py38h42c9631_2
+  - mkl_random=1.2.1=py38ha9443f7_2
+  - ncurses=6.2=he6710b0_1
+  - nettle=3.7.2=hbbd107a_1
+  - ninja=1.10.2=hff7bd54_1
+  - numpy=1.20.2=py38h2d18471_0
+  - numpy-base=1.20.2=py38hfae3a4d_0
+  - olefile=0.46=py_0
+  - openh264=2.1.0=hd408876_0
+  - openssl=1.1.1k=h27cfd23_0
+  - pcre=8.44=he6710b0_0
+  - pip=21.1.1=py38h06a4308_0
+  - pyparsing=2.4.7=pyhd3eb1b0_0
+  - pyqt=5.9.2=py38h05f1152_4
+  - python=3.8.10=hdb3f193_7
+  - python-dateutil=2.8.1=pyhd3eb1b0_0
+  - pytorch=1.8.1=py3.8_cuda11.1_cudnn8.0.5_0
+  - qt=5.9.7=h5867ecd_1
+  - readline=8.1=h27cfd23_0
+  - setuptools=52.0.0=py38h06a4308_0
+  - sip=4.19.13=py38he6710b0_0
+  - six=1.15.0=py38h06a4308_0
+  - sqlite=3.35.4=hdfb4753_0
+  - tk=8.6.10=hbc83047_0
+  - torchaudio=0.8.1=py38
+  - torchvision=0.9.1=py38_cu111
+  - tornado=6.1=py38h27cfd23_0
+  - typing_extensions=3.7.4.3=pyha847dfd_0
+  - wheel=0.36.2=pyhd3eb1b0_0
+  - xz=5.2.5=h7b6447c_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.4.9=haebb681_0
+  - pip:
+    - absl-py==0.12.0
+    - cachetools==4.2.2
+    - chardet==4.0.0
+    - google-auth==1.30.1
+    - google-auth-oauthlib==0.4.4
+    - grpcio==1.38.0
+    - idna==2.10
+    - markdown==3.3.4
+    - oauthlib==3.1.1
+    - pandas==1.2.4
+    - pillow-simd==7.0.0.post3
+    - protobuf==3.17.2
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - pytz==2021.1
+    - requests==2.25.1
+    - requests-oauthlib==1.3.0
+    - rsa==4.7.2
+    - tensorboard==2.5.0
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.0
+    - tqdm==4.61.0
+    - urllib3==1.26.5
+    - werkzeug==2.0.1

main.py

@@ -0,0 +1,233 @@
+"""
+file - main.py
+Main script to train the aesthetic model on the AVA dataset.
+
+Copyright (C) Yunxiao Shi 2017 - 2021
+NIMA is released under the MIT license. See LICENSE for the fill license text.
+"""
+
+import argparse
+import os
+
+import numpy as np
+import matplotlib
+# matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+import torch
+import torch.autograd as autograd
+import torch.optim as optim
+
+import torchvision.transforms as transforms
+import torchvision.datasets as dsets
+import torchvision.models as models
+
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset.dataset import AVADataset
+
+from model.model import *
+
+
+def main(config):
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    writer = SummaryWriter()
+
+    train_transform = transforms.Compose([
+        transforms.Scale(256),
+        transforms.RandomCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225])])
+
+    val_transform = transforms.Compose([
+        transforms.Scale(256),
+        transforms.RandomCrop(224),
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225])])
+
+    base_model = models.vgg16(pretrained=True)
+    model = NIMA(base_model)
+
+    if config.warm_start:
+        model.load_state_dict(torch.load(os.path.join(config.ckpt_path, 'epoch-%d.pth' % config.warm_start_epoch)))
+        print('Successfully loaded model epoch-%d.pth' % config.warm_start_epoch)
+
+    if config.multi_gpu:
+        model.features = torch.nn.DataParallel(model.features, device_ids=config.gpu_ids)
+        model = model.to(device)
+    else:
+        model = model.to(device)
+
+    conv_base_lr = config.conv_base_lr
+    dense_lr = config.dense_lr
+    optimizer = optim.SGD([
+        {'params': model.features.parameters(), 'lr': conv_base_lr},
+        {'params': model.classifier.parameters(), 'lr': dense_lr}],
+        momentum=0.9
+        )
+
+    param_num = 0
+    for param in model.parameters():
+        if param.requires_grad:
+            param_num += param.numel()
+    print('Trainable params: %.2f million' % (param_num / 1e6))
+
+    if config.train:
+        trainset = AVADataset(csv_file=config.train_csv_file, root_dir=config.img_path, transform=train_transform)
+        valset = AVADataset(csv_file=config.val_csv_file, root_dir=config.img_path, transform=val_transform)
+
+        train_loader = torch.utils.data.DataLoader(trainset, batch_size=config.train_batch_size,
+            shuffle=True, num_workers=config.num_workers)
+        val_loader = torch.utils.data.DataLoader(valset, batch_size=config.val_batch_size,
+            shuffle=False, num_workers=config.num_workers)
+        # for early stopping
+        count = 0
+        init_val_loss = float('inf')
+        train_losses = []
+        val_losses = []
+        for epoch in range(config.warm_start_epoch, config.epochs):
+            batch_losses = []
+            for i, data in enumerate(train_loader):
+                images = data['image'].to(device)
+                labels = data['annotations'].to(device).float()
+                outputs = model(images)
+                outputs = outputs.view(-1, 10, 1)
+
+                optimizer.zero_grad()
+
+                loss = emd_loss(labels, outputs)
+                batch_losses.append(loss.item())
+
+                loss.backward()
+
+                optimizer.step()
+
+                print('Epoch: %d/%d | Step: %d/%d | Training EMD loss: %.4f' % (epoch + 1, config.epochs, i + 1, len(trainset) // config.train_batch_size + 1, loss.data[0]))
+                writer.add_scalar('batch train loss', loss.data[0], i + epoch * (len(trainset) // config.train_batch_size + 1))
+
+            avg_loss = sum(batch_losses) / (len(trainset) // config.train_batch_size + 1)
+            train_losses.append(avg_loss)
+            print('Epoch %d mean training EMD loss: %.4f' % (epoch + 1, avg_loss))
+
+            # exponetial learning rate decay
+            if config.decay:
+                if (epoch + 1) % 10 == 0:
+                    conv_base_lr = conv_base_lr * config.lr_decay_rate ** ((epoch + 1) / config.lr_decay_freq)
+                    dense_lr = dense_lr * config.lr_decay_rate ** ((epoch + 1) / config.lr_decay_freq)
+                    optimizer = optim.SGD([
+                        {'params': model.features.parameters(), 'lr': conv_base_lr},
+                        {'params': model.classifier.parameters(), 'lr': dense_lr}],
+                        momentum=0.9
+                    )
+
+            # do validation after each epoch
+            batch_val_losses = []
+            for data in val_loader:
+                images = data['image'].to(device)
+                labels = data['annotations'].to(device).float()
+                with torch.no_grad():
+                    outputs = model(images)
+                outputs = outputs.view(-1, 10, 1)
+                val_loss = emd_loss(labels, outputs)
+                batch_val_losses.append(val_loss.item())
+            avg_val_loss = sum(batch_val_losses) / (len(valset) // config.val_batch_size + 1)
+            val_losses.append(avg_val_loss)
+            print('Epoch %d completed. Mean EMD loss on val set: %.4f.' % (epoch + 1, avg_val_loss))
+            writer.add_scalars('epoch losses', {'epoch train loss': avg_loss, 'epoch val loss': avg_val_loss}, epoch + 1)
+
+            # Use early stopping to monitor training
+            if avg_val_loss < init_val_loss:
+                init_val_loss = avg_val_loss
+                # save model weights if val loss decreases
+                print('Saving model...')
+                if not os.path.exists(config.ckpt_path):
+                    os.makedirs(config.ckpt_path)
+                torch.save(model.state_dict(), os.path.join(config.ckpt_path, 'epoch-%d.pth' % (epoch + 1)))
+                print('Done.\n')
+                # reset count
+                count = 0
+            elif avg_val_loss >= init_val_loss:
+                count += 1
+                if count == config.early_stopping_patience:
+                    print('Val EMD loss has not decreased in %d epochs. Training terminated.' % config.early_stopping_patience)
+                    break
+
+        print('Training completed.')
+
+        '''
+        # use tensorboard to log statistics instead
+        if config.save_fig:
+            # plot train and val loss
+            epochs = range(1, epoch + 2)
+            plt.plot(epochs, train_losses, 'b-', label='train loss')
+            plt.plot(epochs, val_losses, 'g-', label='val loss')
+            plt.title('EMD loss')
+            plt.legend()
+            plt.savefig('./loss.png')
+        '''
+
+    if config.test:
+        model.eval()
+        # compute mean score
+        test_transform = val_transform
+        testset = AVADataset(csv_file=config.test_csv_file, root_dir=config.img_path, transform=val_transform)
+        test_loader = torch.utils.data.DataLoader(testset, batch_size=config.test_batch_size, shuffle=False, num_workers=config.num_workers)
+
+        mean_preds = []
+        std_preds = []
+        for data in test_loader:
+            image = data['image'].to(device)
+            output = model(image)
+            output = output.view(10, 1)
+            predicted_mean, predicted_std = 0.0, 0.0
+            for i, elem in enumerate(output, 1):
+                predicted_mean += i * elem
+            for j, elem in enumerate(output, 1):
+                predicted_std += elem * (j - predicted_mean) ** 2
+            predicted_std = predicted_std ** 0.5
+            mean_preds.append(predicted_mean)
+            std_preds.append(predicted_std)
+        # Do what you want with predicted and std...
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+
+    # input parameters
+    parser.add_argument('--img_path', type=str, default='./data/images')
+    parser.add_argument('--train_csv_file', type=str, default='./data/train_labels.csv')
+    parser.add_argument('--val_csv_file', type=str, default='./data/val_labels.csv')
+    parser.add_argument('--test_csv_file', type=str, default='./data/test_labels.csv')
+
+    # training parameters
+    parser.add_argument('--train',action='store_true')
+    parser.add_argument('--test', action='store_true')
+    parser.add_argument('--decay', action='store_true')
+    parser.add_argument('--conv_base_lr', type=float, default=5e-3)
+    parser.add_argument('--dense_lr', type=float, default=5e-4)
+    parser.add_argument('--lr_decay_rate', type=float, default=0.95)
+    parser.add_argument('--lr_decay_freq', type=int, default=10)
+    parser.add_argument('--train_batch_size', type=int, default=128)
+    parser.add_argument('--val_batch_size', type=int, default=128)
+    parser.add_argument('--test_batch_size', type=int, default=1)
+    parser.add_argument('--num_workers', type=int, default=2)
+    parser.add_argument('--epochs', type=int, default=100)
+
+    # misc
+    parser.add_argument('--ckpt_path', type=str, default='./ckpts')
+    parser.add_argument('--multi_gpu', action='store_true')
+    parser.add_argument('--gpu_ids', type=list, default=None)
+    parser.add_argument('--warm_start', action='store_true')
+    parser.add_argument('--warm_start_epoch', type=int, default=0)
+    parser.add_argument('--early_stopping_patience', type=int, default=10)
+    parser.add_argument('--save_fig', action='store_true')
+
+    config = parser.parse_args()
+
+    main(config)
+

model/model.py

@@ -0,0 +1,61 @@
+"""
+file - model.py
+Implements the aesthemic model and emd loss used in paper.
+
+Copyright (C) Yunxiao Shi 2017 - 2021
+NIMA is released under the MIT license. See LICENSE for the fill license text.
+"""
+
+import torch
+import torch.nn as nn
+
+class NIMA(nn.Module):
+
+    """Neural IMage Assessment model by Google"""
+    def __init__(self, base_model, num_classes=10):
+        super(NIMA, self).__init__()
+        self.features = base_model.features
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=0.75),
+            nn.Linear(in_features=25088, out_features=num_classes),
+            nn.Softmax())
+
+    def forward(self, x):
+        out = self.features(x)
+        out = out.view(out.size(0), -1)
+        out = self.classifier(out)
+        return out
+
+
+def single_emd_loss(p, q, r=2):
+    """
+    Earth Mover's Distance of one sample
+
+    Args:
+        p: true distribution of shape num_classes × 1
+        q: estimated distribution of shape num_classes × 1
+        r: norm parameter
+    """
+    assert p.shape == q.shape, "Length of the two distribution must be the same"
+    length = p.shape[0]
+    emd_loss = 0.0
+    for i in range(1, length + 1):
+        emd_loss += torch.abs(sum(p[:i] - q[:i])) ** r
+    return (emd_loss / length) ** (1. / r)
+
+
+def emd_loss(p, q, r=2):
+    """
+    Earth Mover's Distance on a batch
+
+    Args:
+        p: true distribution of shape mini_batch_size × num_classes × 1
+        q: estimated distribution of shape mini_batch_size × num_classes × 1
+        r: norm parameters
+    """
+    assert p.shape == q.shape, "Shape of the two distribution batches must be the same."
+    mini_batch_size = p.shape[0]
+    loss_vector = []
+    for i in range(mini_batch_size):
+        loss_vector.append(single_emd_loss(p[i], q[i], r=r))
+    return sum(loss_vector) / mini_batch_size

test.py

@@ -0,0 +1,85 @@
+"""
+file - test.py
+Simple quick script to evaluate model on test images.
+
+Copyright (C) Yunxiao Shi 2017 - 2021
+NIMA is released under the MIT license. See LICENSE for the fill license text.
+"""
+
+import argparse
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+import pandas as pd
+from tqdm import tqdm
+import torch
+import torchvision.models as models
+import torchvision.transforms as transforms
+
+from model.model import *
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', type=str, help='path to pretrained model')
+parser.add_argument('--test_csv', type=str, help='test csv file')
+parser.add_argument('--test_images', type=str, help='path to folder containing images')
+parser.add_argument('--workers', type=int, default=4, help='number of workers')
+parser.add_argument('--predictions', type=str, help='output file to store predictions')
+args = parser.parse_args()
+
+base_model = models.vgg16(pretrained=True)
+model = NIMA(base_model)
+
+try:
+    model.load_state_dict(torch.load(args.model))
+    print('successfully loaded model')
+except:
+    raise
+
+seed = 42
+torch.manual_seed(seed)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = model.to(device)
+
+model.eval()
+
+test_transform = transforms.Compose([
+    transforms.Scale(256), 
+    transforms.RandomCrop(224), 
+    transforms.ToTensor(), 
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+                         std=[0.229, 0.224, 0.225])
+    ])
+
+test_df = pd.read_csv(args.test_csv, header=None)
+test_imgs = test_df[0]
+pbar = tqdm(total=len(test_imgs))
+
+mean, std = 0.0, 0.0
+for i, img in enumerate(test_imgs):
+    im = Image.open(os.path.join(args.test_images, str(img) + '.jpg'))
+    im = im.convert('RGB')
+    imt = test_transform(im)
+    imt = imt.unsqueeze(dim=0)
+    imt = imt.to(device)
+    with torch.no_grad():
+        out = model(imt)
+    out = out.view(10, 1)
+    for j, e in enumerate(out, 1):
+        mean += j * e
+    for k, e in enumerate(out, 1):
+        std += e * (k - mean) ** 2
+    std = std ** 0.5
+    gt = test_df[test_df[0] == img].to_numpy()[:, 1:].reshape(10, 1)
+    gt_mean = 0.0
+    for l, e in enumerate(gt, 1):
+        gt_mean += l * e
+    # print(str(img) + ' mean: %.3f | std: %.3f | GT: %.3f' % (mean, std, gt_mean))
+    if not os.path.exists(args.predictions):
+        os.makedirs(args.predictions)
+    with open(os.path.join(args.predictions, 'pred.txt'), 'a') as f:
+        f.write(str(img) + ' mean: %.3f | std: %.3f | GT: %.3f\n' % (mean, std, gt_mean))
+    mean, std = 0.0, 0.0
+    pbar.update()