pavancm / CONTRIQUE

Official implementation for "Image Quality Assessment using Contrastive Learning"
125 stars 12 forks source link

How to generate feat.npy for the score regressor training? #4

Closed CaptainEven closed 2 years ago

CaptainEven commented 2 years ago

I have made custom data set and labels with Matlab code and the below python scripts:

def gen_csv(dir_list,
            csv_path,
            ext=".bmp",
            data_mode="syn",
            write_mode="w"):
    """
    Generate CSV file from dir list
    :param data_mode: syn | ref | ugc
    :param write_mode: w | a+
    """
    if len(dir_list) == 0:
        print("[Err]: empty dir list.")
        return

    print("Data mode: ", data_mode)
    print("Write mode: ", write_mode)

    N_DISTORTIONS = 25
    N_CLS = N_DISTORTIONS * 5 + 1
    print("N_CLS: {:d}".format(N_CLS))
    cnt = 0
    with open(csv_path, write_mode, encoding="utf-8") as f:
        if write_mode != "a+":
            f.write(",Unnamed: 0,File_names,labels\n")
        for dir_path in dir_list:
            print("Processing {:s}...".format(dir_path))
            img_names = [x for x in os.listdir(dir_path) if x.endswith(ext)]
            for img_name in img_names:
                img_path = dir_path + "/" + img_name
                if not os.path.isfile(img_path):
                    continue

                print("processing {:s}".format(img_path))

                items = img_name.split(".")[0].split("_")
                # print(items)

                if len(items) > 3 and data_mode == "syn":
                    distort_type = int(items[-2])
                    distort_level = int(items[-1])

                f.write('{:d},{:d},{:s},"['.format(cnt, cnt, img_path))

                for i in range(N_CLS):
                    if data_mode == "syn":
                        if i == (distort_type - 1) * 5 + (distort_level - 1) + 1:
                            if i == N_CLS - 1:
                                f.write("1.0")
                            else:
                                f.write("1.0, ")
                        elif 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    elif data_mode == "ref":
                        if i == 0:
                            f.write("1.0, ")
                        elif 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    elif data_mode == "ugc":
                        if 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    else:
                        print("[Err]: invalid data mode, data mode should be:"
                              "syn | ref | ugc")
                f.write(']"\n')

                cnt += 1

the loss reduced from 11.7 to 3.4 now after 29 epoches... image

How to train the seond stage regressor?

pavancm commented 2 years ago

For obtaining feat.npy, calculate features for all images on which you want to train the linear regressor using the frozen encoder module as demonstrated in demo_score.py. Once the features are calculated the regressor can be trained using train_regressor.py.

CaptainEven commented 2 years ago

@pavancm Thanks for replying! By the way, is there any code for reference?

pavancm commented 2 years ago

Please refer to demo_feat.py for calculating features. The script calculates features for a single image, for multiple images rerun the script by changing image paths

zqgCoder commented 2 years ago

Hello, I have learned a lot from your discussion, but I have some questions, that is, I repeatedly run demo_feat.py to get feat.npy, then what is the internal format of score, thanks.

pavancm commented 2 years ago

score will be a numpy array with dimensions equal to the number of images you are using to train the regressor.