@LastEditTime: 2019-09-16 12:55:43
@Update: 
'''
import os
import cv2
import torch
import pickle
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from module import CRN, _pca_

torch.set_default_tensor_type('torch.cuda.FloatTensor')

## Load models
pca = _pca_.PCA()
model = pickle.load(open("models/final_model.pkl", 'rb'))
model = model.cuda()

datapath = "data/makeuporigin"
outputpath = "data/output"
makeuppath = "data/makeup"

with open("%s/multidetect.txt" % datapath, 'r') as f:
    detects = eval(f.read())

n = len(detects)
for i, (dirname, (_, [x1, y1, x2, y2], _)) in enumerate(detects.items()):

    print("[%d]/[%d]" % (i + 1, n))
Ejemplo n.º 2
0
def train(param):
    # train_dataset = registry.create('Dataset', param["train_dataset"]["name"])(**param["train_dataset"]["kwargs"])
    # valid_dataset = registry.create('Dataset', param["valid_dataset"]["name"])(**param["valid_dataset"]["kwargs"])
    # train_data_loader = data.DataLoader(train_dataset, **param["loader"])
    pca = _pca_.PCA(**param["PCA"]["kwargs"])

    # model_names = ["model_CRN_6_110.pkl",
    # "model_CRN_6_120.pkl",
    # "model_CRN_6_130.pkl",
    # "model_CRN_6_140.pkl",
    # "model_CRN_6_150.pkl"]

    model_names = ["final_model.pkl"]

    torch.set_default_tensor_type('torch.cuda.FloatTensor')

    selected_ids = [3, 903, 1015, 1012, 1008]

    for model_id, model_name in enumerate(model_names):

        model = pickle.load(open("./models/" + model_name, "rb"))
        model = model.cuda()

        print("starting validating...")

        parts = model_name.split(".")
        # model_path = "./output_beautified/" + parts[0]
        model_path = "./output"

        if not os.path.exists(model_path):
            os.mkdir(model_path)

        if param["valid"]:
            with torch.no_grad():
                print(param["valid"])
                for i in selected_ids:
                    # path = param["valid"] + "/beautified ({}).jpg".format(i)
                    # inputs = img_to_tensor(path)
                    # path = "./data/EECS442_Makeup_Go/result_original/original ({}).jpg".format(i)
                    # targets = img_to_tensor(path)
                    # inputs.unsqueeze_(0)
                    # targets.unsqueeze_(0)

                    # ground_truth = pca.get_components(inputs, targets, True)
                    # print("generating one image...")
                    # ground_truth = ground_truth[0]

                    # print("gound truth size ", ground_truth.shape)

                    # for j in range(8):
                    #     img = np.array(ground_truth[j]).transpose((1,2,0))
                    #     img[img>255] = 255
                    #     img[img<0] = 0
                    #     cm = None
                    #     fn = model_path + "/result_{}_{}.jpg".format(i, j)
                    #     save_image(img, cm, fn)

                    # path = param["valid"] + "/beautified({}).jpg".format(i)
                    # x = img_to_tensor(path).cuda()
                    # x.unsqueeze_(0)
                    # output = model(x)
                    # img = pca.generate_img(output, x)
                    # np_img = np.array(img[0])
                    # np_img[np_img>255] = 255
                    # np_img[np_img<0] = 0
                    # np_img = np_img.transpose((1, 2, 0)).astype(np.uint8)
                    # plt.close()
                    # plt.imshow(np_img)
                    # plt.savefig(model_path + "/result_{}.jpg".format(i))

                    path = param["valid"] + "/beautified ({}).jpg".format(i)
                    x = img_to_tensor(path).cuda()
                    sz = x.shape
                    x.unsqueeze_(0)
                    output = torch.squeeze(model(x), 0)
                    print("generating one image...")
                    img = pca.generate_img(output, x)
                    np_img = np.array(img[0])
                    np_img[np_img > 255] = 255
                    np_img[np_img < 0] = 0
                    np_img = np_img.transpose((1, 2, 0)).astype(np.uint8)
                    save_image(np_img, None,
                               model_path + "/result_{}.jpg".format(i))
Ejemplo n.º 3
0
def train(param):
    # TODO
    # Use PCA part to get a test dataset:
    # DATASET:                  1. train data: the beautified images
    # (Based on PCA Results ->) 2. The components of the difference between the beautified images and the ground truth
    # Check the size of dataset images, decide whether to use batch operation
    # return a tensor of the eigenvalues
    # DRAFT:

    train_dataset = registry.create(
        'Dataset',
        param["train_dataset"]["name"])(**param["train_dataset"]["kwargs"])
    # valid_dataset = registry.create('Dataset', param["valid_dataset"]["name"])(**param["valid_dataset"]["kwargs"])
    train_data_loader = data.DataLoader(train_dataset, **param["loader"])
    # valid_data_loader = data.DataLoader(valid_dataset, **param["loader"])
    # pin_memory=param.use_gpu,

    # eigenvalues = train_dataset.get_()
    # model = torch.nn.DataParallel(registry.create('Network', param.network.name)(**param.network.kwargs))
    model = CRN.CRN(**param["network"]["kwargs"])

    criterion = registry.create(
        'Loss', param["loss"]["name"])(**param["loss"]["kwargs"])
    optimizer = registry.create('Optimizer', param["optimizer"]["name"])(
        model.parameters(), **param["optimizer"]["kwargs"])
    lr_scheduler = registry.create('LRScheduler',
                                   param["lr_scheduler"]["name"])(
                                       optimizer,
                                       **param["lr_scheduler"]["kwargs"])

    pca = _pca_.PCA(**param["PCA"]["kwargs"])

    # TODO: checkpoint

    # if param.use_gpu:
    #     model = model.cuda()
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
    model = model.cuda()

    load_previous = False
    if load_previous:
        pca = pickle.load(open("trained_models/model_PCA_9_0.pkl", "rb"))
        model = pickle.load(open("trained_models/model_CRN_9_0.pkl", "rb"))

    is_train = True

    g_counter = 0

    if is_train:
        save_dir = "./models"
        # idx = int(os.listdir(save_dir)[-1][5:]) + 1
        save_dir = "./models/model_PCA_modified"
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        save_dir += "/"

        if not os.path.exists("./output_PCA_modified"):
            os.mkdir("./output_PCA_modified")

        for epoch in range(int(param["epoch"])):
            print("epoch:", epoch)
            counter = 0
            for inputs, targets in train_data_loader:
                lr_scheduler.step()
                optimizer.zero_grad()

                ground_truth = pca.get_components(inputs, targets, True)
                # print('########')
                # print(ground_truth.shape)
                # for current_data in ground_truth:
                #     for img in current_data:
                #         save_image(img/torch.max(img)*255, 'decomposed_ground_truth/' + str(g_counter) + '.jpg')
                #         g_counter += 1
                # print('########')
                # print('=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=')
                # print(ground_truth)
                # print('-0-0-0-0-0-0-0-0-0-0-0-0-0-0-0-')

                # print("get")

                output = model(inputs)
                # print(output)
                # print("model")
                # with torch.no_grad():
                #     img = pca.generate_img(output[0], inputs[0])
                #     save_image(img, "output_PCA_modified/train/result_{}_{}_before.jpg".format(epoch, counter))

                loss = criterion(output, ground_truth)

                if counter % 10 == 0:
                    # if not torch.isnan(loss):
                    #     # Fixme change the epoch +10 back
                    # with open(save_dir + 'model_PCA_{}_{}.pkl'.format(epoch+10, counter), 'wb') as out:
                    #     pickle.dump(pca, out, pickle.HIGHEST_PROTOCOL)

                    with open(
                            save_dir +
                            'model_CRN_{}_{}.pkl'.format(epoch, counter),
                            'wb') as out:
                        pickle.dump(model, out, pickle.HIGHEST_PROTOCOL)

                    print("loss {}:".format(counter), loss.item(), "\tlr:",
                          lr_scheduler.get_lr()[0])

                    print("starting validating...")
                    if param["valid"]:
                        with torch.no_grad():
                            # image = Image.open(param["valid"])
                            # x = TF.to_tensor(image).cuda()
                            x = img_to_tensor(param["valid"])
                            sz = x.shape
                            x.unsqueeze_(0)
                            output = torch.squeeze(model(x), 0)
                            print("generating one image...")
                            img = pca.generate_img(output, x)
                            np_img = np.array(img[0])
                            # print('before conversion')
                            # print('@@@@@@@@')
                            # print(np_img)
                            np_img[np_img > 255] = 255
                            np_img[np_img < 0] = 0
                            np_img = np_img.transpose(
                                (1, 2, 0)).astype(np.uint8)
                            plt.close()
                            plt.imshow(np_img)
                            plt.savefig(
                                "output_PCA_modified/result_{}_{}.jpg".format(
                                    epoch, counter))
                            # save_image(img, "output_PCA_modified/result_{}_{}.jpg".format(epoch, counter))

                loss.backward()

                # perform gradient clipping
                clip_grad_norm_(model.parameters(), 2)

                optimizer.step()
                # with torch.no_grad():
                #     output = torch.squeeze(model(inputs[0].unsqueeze_(0)), 0)
                #     img = pca.generate_img(output, inputs[0]).view(sz)
                #     save_image(img, "output_PCA_modified/train/result_{}_{}_after.jpg".format(epoch, counter))
                counter += 1

            # model.eval()
            # @TODO: EVALUATE!!!
            # valid_data_loader...
            # how to get the truth
            # ground_truth = pca.get_components(inputs, targets, False)
            # how to get img:
            # output = model(inputs)
            # img = pca.generate_img(output, inputs)

            # model.train()

            # torch.save({
            #     'model': model.state_dict(),
            #     'optimizer': optimizer.state_dict(),
            #     'lr_scheduler': lr_scheduler.state_dict()
            # },
            #     os.path.join("TODO!!!")
            # )
            # logging.debug('saving model done')

            # print("starting validating...")
            # if param["valid"]:
            #     with torch.no_grad():
            #         image = Image.open(param["valid"])
            #         x = TF.to_tensor(image).cuda()
            #         sz = x.shape
            #         x.unsqueeze_(0)
            #         output = torch.squeeze(model(x), 0)
            #         print("generating one image...")
            #         img = pca.generate_img(output, x).view(sz)
            #         save_image(img, "output_PCA_modified/result_{}.jpg".format(epoch))

    else:
        print("starting validating...")

        if param["valid"]:
            with torch.no_grad():
                image = Image.open(param["valid"])
                parts = param["valid"].split("/")
                parts = parts[-1].split(".")
                parts = parts[0].split("_")
                name = parts[0]
                x = TF.to_tensor(image).cuda()
                sz = x.shape
                x.unsqueeze_(0)
                output = model(x)
                print(output.shape)
                output = torch.squeeze(output, 0)
                print(output.shape)
                print("generating one image...")
                img = pca.generate_img(output, x).view(sz)
                print(img)
                save_image(img, "output/{}_result.jpg".format(name))