colorization icon indicating copy to clipboard operation
colorization copied to clipboard

get colored images from my pretrained caffemodel

Open ZiqiYe77 opened this issue 3 years ago • 13 comments

hello,thanks for your work. I want to use the pre-trained caffemodel that training from my own dataset to output the colored images.First I find colorize.py( Test-Time Python Script) and replace default model(colorization_release_v2.caffemodel) with my caffemodel(e.g. colornet_iter_51000.caffemodel) , but I found my caffemodel's network structure is different from the default model's, so the output image seems to be wrong. what should I do? If you have free time,please answer my question,thanks!

ZiqiYe77 avatar May 21 '21 02:05 ZiqiYe77

Hi, did you use the correct "prototxt" file? This file contains the architecture details

amoazeni75 avatar May 24 '21 07:05 amoazeni75

yes, I also replaced the pre-trained "prototxt" file with the training "prototxt" file, but another problem appeared,training pipeline can't output the colored images

ZiqiYe77 avatar May 24 '21 08:05 ZiqiYe77

I know. I had the same problem with demo_release.py (Pytorch version), but when I used this command: python ./colorize.py -img_in ./demo/imgs/ILSVRC2012_val_00041580.JPEG -img_out ./out.png and replaced the .caffemodel and .prototxt file inside the colorize.py, my problem solved.

amoazeni75 avatar May 24 '21 08:05 amoazeni75

I replaced the corresponding part in the colorize.py with ./models/colorization_train_val_v2.prototxt and colornet_iter_135000.caffemodel(caffemodel saved in training process). but the output image is from net.blobs['class8_ab'].data in colorize.py. this 'class8_ab' blobs is not in the structure of training caffemodel. Will the output result be correct? Can you tell me which file you replaced?

ZiqiYe77 avatar May 24 '21 08:05 ZiqiYe77

I use "./models/colorization_deploy_v2.prototxt" for .prototxt file. I think this file contains the lastest architecture updates. colorization_train_val_v2.prototxt is for training and validating not for testing. For testing the model on a single image you should use ./models/colorization_deploy_v2.prototxt file

amoazeni75 avatar May 24 '21 11:05 amoazeni75

so you use "./models/colorization_deploy_v2.prototxt" and the model you train(not default='./models/colorization_release_v2.caffemodel') and you get the right result?i just hope that the colored images I get are trained on the model with fine-tuned parameters based on my own dataset.

ZiqiYe77 avatar May 24 '21 12:05 ZiqiYe77

Yes, exactly. I also trained the model on another dataset. No worry!

amoazeni75 avatar May 24 '21 12:05 amoazeni75

I'm sorry to bother you again,I don't know if there was a problem with the pre-processing data in training process(I processed my images dataset into LMDB files,BGR images),I ran the colorize.py as you did, the output image looks like a BGR picture, not a RGB image! Have you encountered such a situation? Can you tell me how you preprocessed your data? This is the result i ran colorize.py: 图片4

ZiqiYe77 avatar May 24 '21 13:05 ZiqiYe77

No worry. It is interesting. No, I did not face this problem. Here is the script I used to prepare the LMDB dataset.

img_list = sorted(glob.glob(read_path))
dataset = []
data_size = 0
for i, v in enumerate(img_list):
    img = cv2.imread(v, cv2.IMREAD_UNCHANGED)
    dataset.append(img)
    data_size += img.nbytes
    print("read {}".format(v))
print(len(img_list))
env = lmdb.open(save_path, map_size=data_size * 10)

with env.begin(write=True) as txn:
    # txn is a Transaction object
    for i, v in enumerate(img_list):
        base_name = os.path.splitext(os.path.basename(v))[0]
        key = base_name.encode('ascii')
        data = dataset[i].tobytes()
        if dataset[i].ndim == 2:
            H, W = dataset[i].shape
            C = 1
        else:
            H, W, C = dataset[i].shape
        datum = caffe.proto.caffe_pb2.Datum()
        datum.channels = C
        datum.height = H
        datum.width = W
        datum.data = data
        txn.put(key, datum.SerializeToString())

amoazeni75 avatar May 24 '21 13:05 amoazeni75

thanks,I read your script. it seems that you did not preprocess the training data into BGR images,so this can directly get the correct result.

ZiqiYe77 avatar May 24 '21 14:05 ZiqiYe77

Your welcome. I used the BGR format. When you read an image using OpenCV, its channel order is BGR, not RGB; that means my LMDB contains images in BGR format.

amoazeni75 avatar May 24 '21 14:05 amoazeni75

oh! yes, thanks a lot!

ZiqiYe77 avatar May 24 '21 14:05 ZiqiYe77

Good luck!

amoazeni75 avatar May 24 '21 14:05 amoazeni75