guochengqian / tenet

[ICCP'22] Rethinking Learning-based Demosaicing, Denoising, and Super-Resolution Pipeline

TENet [ArXiv] [pixelshift200]

Rethinking Learning-based Demosaicing, Denoising, and Super-Resolution Pipeline

By Guocheng Qian*, Yuanhao Wang*, Jinjin Gu, Chao Dong, Wolfgang Heidrich, Bernard Ghanem, Jimmy S. Ren

The original name of this project is: "Trinity of Pixel Enhancement: a Joint Solution for Demosaicking, Denoising and Super-Resolution"

TENet

We insert the proposed pipeline DN -> SR -> DM into an end-to-end network constructed by RRDB for the joint DN, DM and SR. We leverage the detachable branch to provide the middle stage supervision.

PixelShift200 dataset

We employ advanced pixel shift technology to perform a full color sampling of the image. Pixel shift technology takes four samples of the same image, and physically controls the camera sensor to move one pixel horizontally or vertically at each sampling to capture all color information at each pixel. The pixel shift technology ensures that the sampled images follow the distribution of natural images sampled by the camera, and the full information of the color is completely obtained. In this way, the collected images are artifacts-free, which leads to better training results for demosaicing related tasks.

Download PixelShift200 and the original RAW files from pxielshift200 website for using this Google Drive Link

Environment installation

Clone this github repo and install the environment by:

git clone https://github.com/guochengqian/TENet
cd TENet
source install.sh
conda activate tenet

Data preparation

1. synthetic data preparation [Optional] Processing synthetic dataset if you are interested in synthetic benchmark.

   1. Download (DIV2K) dataset

   2. mkdir data && cd data

   3. Link DIV2K data into ./data/DIV2K, e.g. ln -s /data/lowlevel/DIV2K ./

   4. Crop DIV2K

      cd ../datasets
      python crop_imgs.py # crop train images
      python crop_imgs.py --src_dir ../data/DIV2K/DIV2K_val5_HR --save_dir ../data/DIV2K/DIV2K_val5_HR_sub # crop val5 images

2. PixelShift200 data preparation Processing PixelShift200 dataset if you are interested in PixelShift200 benchmark.

   1. Download Pixelshift200. They are .mat format, having 4 channels (R, Gr, Gb, B). Unzip the .zip file and put all folders inside into one folder called pixelshift200. For example, put here /data/lowlevel/pixelshift200.

   2. cd TENet && mkdir data && cd data

   3. Link PixelShift200 data into ./data/pixelshift200, e.g. ln -s /data/lowlevel/pixelshift200 pixelshift200

   4. Crop images into 512*512, and generate the text file that contains the location of each image:

      cd ../datasets
      python crop_pixelshift200.py 
      python generate_datalist_pixelshift.py 

Training

Train joint models

• DN+SR->DM (our TENet)

  python train.py --in_type noisy_lr_raw --mid_type raw --out_type linrgb --model tenet --n_gpus 4 --block rrdb  --n_blocks 12

• DN+DM+SR [Optional] (baseline: end to end without pipeline)

  python train.py --in_type noisy_lr_raw --mid_type None --out_type linrgb --model tenet --n_gpus 4 --block rrdb --n_blocks 12

  Note:

  1. --mid_type raw is to activate the auxiliary mid stage supervision. Here we add raw as the supervision, therefore, the pipeline will be DN+SR->DM

  2. for training on a different dataset, like DIV2K, change --dataset div2k

  3. for training with Gaussian noise model, add --noise_model g

  4. for using a different building block, such as NLSA --block nlsa, or EAM --block eam , or RRG --block rrg, or DRLM --block drlm or RRDB --block rrdb

  5. wandb is used by default. Set --no_wandb to not using wandb

Train sequential models [Optional] (ablation study)

# RawDN
python train.py --in_type noisy_raw --out_type raw --model resnet --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block rrdb --n_blocks 12

# RawSR
python train.py --in_type lr_raw --out_type raw --model resnet --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block rrdb --n_blocks 12

# DM
python train.py --in_type raw --out_type linrgb --model resnet --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block rrdb --n_blocks 12

# RGBDN
python train.py --in_type noisy_linrgb --out_type linrgb --model resnet --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block rrdb --n_blocks 12

# RGBSR
python train.py --in_type lr_linrgb --out_type linrgb --model resnet --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block rrdb --n_blocks 12

Train SOTA models [Optional]

We also support training previous work.

• JDSR

  python train.py --in_type noisy_lr_raw --mid_type None --out_type linrgb --model jdsr --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4  --use_wandb --block res --n_blocks 12 --channels 256

• JDnDmSR

  python train.py --in_type noisy_lr_raw --mid_type lr_raw --out_type linrgb --model jdndmsr --scale 2 --train_list datasets/train_pixelshift.txt --val_list datasets/val_pixelshift.txt --n_gpus 4 --n_blocks 2 --block rcab

Testing

bash script_all_pipelines.sh 

This script supports evaluation on all benchmarking datasets as well as the real-shot images for all possible pipelines. Check this script for details

Note: for the real shot images testing, you have to:

1. save the real-shot image as a readable raw image (like in .RAW, .ARW, .DNG format). For example, we use Lightroom mobile version to shot images on iPhone and save the photo in .DNG format.
2. Read the general metadata using RawPy and read the noise profiling metadata using Jeffrey's Image Metadata Viewer or metapicz.

Result

Citation

Please cite the following paper if you feel TENet and PixelShift200 is useful to your research

@article{qian2019rethink,
  title={Rethink the Pipeline of Demosaicking, Denoising, and Super-resolution},
  author={Qian, Guocheng and Wang, Yuanhao and Gu, Jinjin and Dong, Chao and  Heidrich, Wolfgang and Ghanem, Bernard and Ren, Jimmy S},
  journal={arXiv preprint arXiv:1905.02538},
  year={2019}
}

tenet's Issues

test bug?

I use the sh code to test your model, but it has a bug
"RuntimeError: invalid argument 0: Sizes of tensors must match except in deimension 1. Got 3 and 1 in dimension 0 at ......"

I try to modify the code, but other bug again.

so will you provide a well-done test code?

.ARW images as input

How do we give .ARW images as input to the network?

I tried converting them to mono-tiff using dcraw -6 -w -T -d image.ARW, but the output image appears to have a strong red tint.

tenet2-dn-ps200.path

Hi,
May you release your pixelshift 200 pretrained model?

Thanks

some problem about pixelshift image

Thank you very much for providing the pixelshift data set, but during the use process, I found that the collected pictures have serious false color problems in the edge area, which leads to serious false color problems in the training results. Is this caused by lens chromatic aberration? Is there any good solution?

How to extract the image data bayer data in the ARW file instead of exif and meta information

How to extract the image data bayer data in the ARW file instead of exif and meta information

RAW , linRGB, SRGB

Excuse me, as I know, the images are processed between linRGB and RAW in the code. I can't find where the gamma correction and the white balance are. However, the images presented in the paper are SRGB. Did I miss anything?

requirement

where is requirements.txt?

please upload the pretrained model

Hi, I just read about your paper and and loop up your project. You offer the instruction to test the result with your pretrained model, but I didn't find it in the project. Can you please upload the pretrained model? Thanks a lot!

test using provided model

Hello,
I take a mosaic image by running rgb2raw.py as imput to test the algorithm,
however, I get a big mosaic image.
rgb:

bayer raw image(input):

big mosaik image(output):

It seems wrong, and I don't the reason, can you help me?

Question about PixelShift200 data preparation

The generated .txt file train_pixelshift.txt (9444 Lines) and val_pixelshift.txt (20 Lines) are used for training. check them.

I found that pixelshift training db has 200 files, so running python crop_pixelshift200.py gives 200*28 = 5600 files.
And couldn't understand how can you get 20 val files.

How can I get 9444 files / 20 files for training / val?

Tensorboard monitor for recent tensorflow and pytorch

I attached the code for the guys who want to use recent tensorflow >= 2.x

import tensorflow as tf
import numpy as np

try:
    from StringIO import StringIO  # Python 2.7
except ImportError:
    from io import BytesIO         # Python 3.x


class Tf_Logger(object):

    def __init__(self, log_dir):
        """Create a summary writer logging to log_dir."""
        self.writer = tf.summary.create_file_writer(log_dir)


    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        with self.writer.as_default():
            tf.summary.scalar(tag, value, step)
            self.writer.flush()

    def image_summary(self, tag, images, step):
        """Log a list of images."""
        with self.writer.as_default():
            for i, img in enumerate(images):
                if len(img.shape) == 2:
                    img = img[np.newaxis,:,:,np.newaxis]
                elif len(img.shape) == 3:
                    img = img[np.newaxis,:,:,:]
                tf.summary.image('%s/%d' % (tag, i), img, step)
            self.writer.flush()
    
    def histo_summary(self, tag, values, step, bins=1000):
        """Log a histogram of the tensor of values."""

        # Create a histogram using numpy
        counts, bin_edges = np.histogram(values, bins=bins)

        # Fill the fields of the histogram proto
        hist = tf.HistogramProto()
        hist.min = float(np.min(values))
        hist.max = float(np.max(values))
        hist.num = int(np.prod(values.shape))
        hist.sum = float(np.sum(values))
        hist.sum_squares = float(np.sum(values**2))

        # Drop the start of the first bin
        bin_edges = bin_edges[1:]

        # Add bin edges and counts
        for edge in bin_edges:
            hist.bucket_limit.append(edge)
        for c in counts:
            hist.bucket.append(c)

        # Create and write Summary
        with self.writer.as_default():
            tf.summary.histogram(tag, hist, step)
            self.writer.flush()

Open dataset

Hi, I read your paper and thought it was a good job. I want to know if your PixelShift dataset will be made public. Thanks a lot.

where is the metadata of Pixelshift dataset ?

Hello, I can't find the metadata of Pixelshift dataset.

Can you release the RGB data of Pixelshift200 datasets?

Hi, I use your code in Prepro.py to product the RGB gt data, but it look like not the full color image? may u provide the full color RGB picture?

It seems the dataset link sucked.

Could you please update the dataset source web page?

thank you to give me pixelshift200 link to download ,but i find 。。。。。

thank you to give me pixelshift200 link to download ,but i find The data set given to me is not complete, there are only about 6 mat files, can you provide a complete data set, thank you very much

PreTrained Weights

Hello,
Can you please provide the preTrained Weights ?
Thank you

TypeError: Cannot set tensorflow.Summary.Value.simple_value to tensor(0.0915, device='cuda:0', grad_fn=<ThAddBackward>): tensor(0.0915, device='cuda:0', gred_fn=<ThAddBackward>) has type <calss 'torch.Tensor'>, but expected one of: numbers.Real for field Value.simple_value

how to convert rgb image to Bayer raw images (bayer pattern is rggb)

I want to test my images on your model but I don't know how to convert normal image to rggb.
Have you any script for that?

is the raw data 14 bit ?

in the load_dataset.py :
def getitem(self, index):
# read images, crop size

    # Note: RGGB is not RAW, it is a Full color sampled image.
    # (compose: Red, Green_red, Green_blue and blue channel);
    if self.phase == 'train':
        # for training dataset, the mat version if 4 (reads with scipy)
        matfile = loadmat(self.data_lists[index])
        rggb = np.asarray(matfile['ps4k']).astype(np.float32) / (2 ** self.bit - 1)
    else:
        # for testing dataset, the mat version if 7.3 (reads with h5py)
        with h5py.File(self.data_lists[index], 'r') as matfile:
            rggb = np.asarray(matfile['raw']).astype(np.float32) / (2 ** self.bit - 1)
            rggb = np.transpose(rggb, (2, 1, 0))
            matainfo = matfile['metadata']
            matainfo = {'colormatrix': np.transpose(matainfo['colormatrix']),
                        'red_gain': matainfo['red_gain'],
                        'blue_gain': matainfo['blue_gain']
                        }
            ccm, red_g, blue_g = process.metadata2tensor(matainfo)
            metadata = {'ccm': ccm, 'red_gain': red_g, 'blue_gain': blue_g}

    rggb = crop_img_np(rggb, self.patch_size, center_crop=self.phase != 'train')  # in PIL
    linrgb = np.stack((rggb[:, :, 0], np.mean(rggb[:, :, 1:3], axis=-1), rggb[:, :, 3]), axis=2)

    if self.phase == 'train':
        linrgb = aug_img_np(linrgb, random.randint(0, 7))
    linrgb = TF.to_tensor(linrgb)
    linrgb = torch.clamp(linrgb, 0., 1.)

请问raw data是14bit吗？当self.bit=14时，为什么读取的原始数据rggb会超过1呢？而self.bit=16时，数据最大是1，您是把数据映射到16bit了吗？

how to download 200 pixel shift dataset ？

200 pixel shift dataset how download ？ --03/06/2023

how to set noise_level(or sigma) if I want use it to process other RAW image?

how to set noise_level if I want use it to process other RAW image?
Can I use the default parameter value?
please help me,thank you!