baidut / paq2piq

PaQ2PiQ in PyTorch

PyTorch PaQ-2-PiQ: Patch Quality 2 Picture Quality Prediction

PyTorch implementation of PaQ2PiQ

Demo

Get Started

Check out demo.ipynb

get predicts from a file:

model = InferenceModel(RoIPoolModel(), 'models/RoIPoolModel.pth')
output = model.predict_from_file("images/Picture1.jpg")

predict from a PIL image:

model = InferenceModel(RoIPoolModel(), 'models/RoIPoolModel.pth')
image = Image.open("images/Picture1.jpg")
output = model.predict_from_pil_image(image)

The output would be a dictionary:

output['global_score'] # a float scale number indicating the predicted global quality
output['local_scores']  # a 20x20 numpy array indicating the predicted  local quality scores
output['category']  # From low to high quality: 'Bad', 'Poor', 'Fair', 'Good', 'Excellent'

Installing

git clone https://github.com/baidut/paq2piq
cd paq2piq
virtualenv -p python3.6 env
source ./env/bin/activate
pip install -r requirements.txt

Dataset

The model was trained on FLIVE. You can get it from here. (Feel free to create an issue here if you encountered any problem) For each image, we cropped three different-sized patches. The image data and patch location is taken as input while their scores as output. Here is an example:

Model

Used ResNet18 pretrained on ImageNet as backbone

Pre-trained model

Download

Train it with Pytorch-lightning

from pytorch_lightning_module import *
module = RoIPoolLightningModule()
trainer = pl.Trainer(gpus=[0])    
trainer.fit(module)

Train it with Pure-Pytorch

Change the settings here:

export PYTHONPATH=.
export PATH_TO_MODEL=models/RoIPoolModel.pth
export PATH_TO_IMAGES=/storage/DATA/images/
export PATH_TO_CSV=/storage/DATA/FLIVE/
export BATCH_SIZE=16
export NUM_WORKERS=2
export NUM_EPOCH=50
export INIT_LR=0.0001
export EXPERIMENT_DIR_NAME=/storage/experiment_n0001

Train model

python cli.py train_model --path_to_save_csv $PATH_TO_CSV \
                                --path_to_images $PATH_TO_IMAGES \
                                --batch_size $BATCH_SIZE \
                                --num_workers $NUM_WORKERS \
                                --num_epoch $NUM_EPOCH \
                                --init_lr $INIT_LR \
                                --experiment_dir_name $EXPERIMENT_DIR_NAME

Use tensorboard to tracking training progress

tensorboard --logdir .

Validate model on val and test datasets

python cli.py validate_model --path_to_model_state $PATH_TO_MODEL \
                                    --path_to_save_csv $PATH_TO_CSV \
                                    --path_to_images $PATH_TO_IMAGES \
                                    --batch_size $BATCH_SIZE \
                                    --num_workers $NUM_EPOCH

Get scores for one image

python cli.py get-image-score --path_to_model_state $PATH_TO_MODEL \
--path_to_image test_image.jpg

Contributing

Contributing are welcome

Acknowledgments

• PyTorch NIMA: Neural IMage Assessment
• https://github.com/vdouet/Discriminative-learning-rates-PyTorch

paq2piq's Issues

About train.csv|val.csv|test.csv

Can you provide the train.csv|val.csv|test.csv ?

About error when train the model !!!

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 333 and 285 in dimension 2 at /pytorch/aten/src/TH/generic/THTensor.cpp:612

Model doesn't work on multiprocessing

model = InferenceModel(RoIPoolModel(), 'models/RoIPoolModel.pth')
output = model.predict_from_file("images/Picture1.jpg")

Model gives different scores on running the script for the same image two times in multithreading , that is a big problem.
Ideally it should give same scoring at every run (with multithreading/multiprocessing).

The performance difference between the validation set and the test validation set

In your released paper, I find that the performance difference between the validation set and the test validation set is quite large, which is not normal. Could you explain why?

cannot produce the results in table 4

Hello
Thank you for the great work.
I am trying to re-produce the results in table 4 of the paper however I cannot.

• should we test the model only on the validation set of CLIVE and KonIQ or on the whole dataset?

• can you share the trained models,
  P2P-BM (Sec. 4.1) , P2P-RM (Sec. 4.2), and P2P-FM (Sec. 4.4), that can produce those results?

The results that I get from running the pre-trained model are far less than what is reported on the paper.

Would you please release the baseline model？

Performance on KonIQ10K and CLIVE

Hi, can you describe in details the evaluation protocols of CLIVE and KonIQ?

As reported in the paper, your model achieves 0.784 and 0.788 on CLIVE and KonIQ, respectively.

I use your pre-trained model to test on these two databases. In my experiments, images are randomly divided into training and testing sets under a 80%/20% ratio. But I can only get median SRCC across ten sessions on the test sets of CLIVE and KonIQ of 0.719 and 0.722, which are far from the numbers reported in your paper.

Problems in training model with pytorch version code

Excuse me? train.csv How to generate, in the python version of the code, I use the Train it with Pure-Pytorch training model, your model.py Roipoolmodel module y prediction output out in file_ Features are 1-D, but dataset.py The value of Y in the file is label_ cols = 'mos_ image', 'mos_ patch_ 1', 'mos_ patch_ 2', 'mos_ patch_ 3 ', which means 4 dimensions.Can you give a detailed description of the training process and configuration parameters using pytorch?
Second, what is the input of X when training with the pytorch version code, the whole picture, or three patches, or input together

nice work, new dataset

license of this work

Congraduaiton on the nice work first!

I am wondering if you can specify the license of this code base, or the released dataset?