FBCNN: Flexible Blind JPEG Artifact Removal

FBCNN (Flexible Blind Convolutional Neural Network) is a PyTorch implementation for blind JPEG artifact removal. Unlike traditional methods that require a specific model for each quality factor or require DCT coefficients, FBCNN predicts an adjustable quality factor to allow users to control the trade-off between artifact removal and detail preservation. It is designed to handle:

• Single JPEG images (color and grayscale).
• Double JPEG images (both aligned and non-aligned).
• Real-world JPEG images.

Requirements:

• Python 3.7+
• PyTorch 1.7

Non-aligned double JPEG compression occurs when the 8x8 blocks of two successive JPEG compressions do not align (e.g., due to a one-pixel shift caused by cropping, resizing, or rotation).

Existing blind methods often fail on these images when the first quality factor (QF1) is less than or equal to the second (QF2). FBCNN addresses this using two approaches:

1. FBCNN-D: Uses a single JPEG degradation model combined with automatic dominant QF correction.
2. FBCNN-A: Uses a double JPEG degradation model for data augmentation during training, defined as: y = JPEG(shift(JPEG(x, QF1)), QF2)

This augmentation helps the model handle the misalignment common in real-world internet images (e.g., social media re-compression).

FBCNN introduces flexibility into blind JPEG artifact removal by decoupling the quality factor from the image.

How it works:

1. Decoupler Module: Extracts the quality factor from the input JPEG image.
2. Reconstructor Module: Uses a quality factor attention block to embed the predicted quality factor into the reconstruction process.

This allows users to manually adjust the predicted quality factor to find the optimal balance between removing artifacts and preserving image details.

FBCNN provides different entry points for testing depending on the image characteristics and the training configuration used. Use the corresponding script for your use case:

• Grayscale JPEG images: Standard grayscale testing.
• Grayscale JPEG images (Double JPEG trained): Testing grayscale images using a model trained with the double JPEG degradation model.
• Color JPEG images: Standard color image testing.
• Real-World JPEG images: Testing on real-world JPEG datasets.

# Grayscale JPEG images
python main_test_fbcnn_gray.py

# Grayscale JPEG images, trained with double JPEG degradation model
python main_test_fbcnn_gray_doublejpeg.py

# Color JPEG images
python main_test_fbcnn_color.py

# Real-World JPEG images
python main_test_fbcnn_color_real.py

To start training the model, execute the main_train_fbcnn.py script. Configuration for training is managed via .json files located in the ./options/ directory.

python main_train_fbcnn.py

If you want to test the model without local installation, you can use the integrated Gradio Web Demo hosted on Huggingface Spaces.

https://huggingface.co/spaces/danielsapit/JPEG_Artifacts_Removal

To use the FBCNN models, download the pre-trained weights from the official GitHub Release v1.0. The available models are:

• fbcnn_color.pth: For color image JPEG artifact removal.
• fbcnn_gray.pth: For grayscale image JPEG artifact removal.
• fbcnn_gray_double.pth: For grayscale images that have undergone double JPEG compression.