vaishnavk/Design2Style

1import os2import numpy as np3 4from PIL import Image5from config import *6from collections import OrderedDict7from models import VGG, U2NET8 9 10def open_image(path):11    return Image.open(path)12 13 14def load_image(img):15    img = open_image(img).convert()16    img = basic_transform(image=np.array(img))['image']17    img = img.unsqueeze(0)18 19    return img.to(DEVICE)20 21 22def load_essential_images(original_img, style_img):23    original_img = load_image(original_img)24    style_img = load_image(style_img)25 26    generated = original_img.clone().requires_grad_(True)27 28    return original_img, style_img, generated29 30 31def gram_matrix(feature_map):32    return feature_map.mm(feature_map.t())33 34 35def min_max_normalization(d):36    return (d - torch.min(d)) / (torch.max(d) - torch.min(d))37 38 39def checkpoint_exists(filename):40    return filename in os.listdir(CHECKPOINT_DIR)41 42 43def load_vgg_model():44    return VGG().to(DEVICE).eval()45 46 47def load_u2net_model(in_ch=3, out_ch=4, checkpoint=U2NET_CLOTHES_CHECKPOINT_FILE, ordered_dict=True):48    model = U2NET(in_ch, out_ch).to(DEVICE)49 50    if checkpoint_exists(checkpoint):51        checkpoint = torch.load(52            os.path.join(CHECKPOINT_DIR, checkpoint),53            map_location=DEVICE54        )55 56        if ordered_dict:57            state_dict = OrderedDict()58 59            for k, v in checkpoint.items():60                name = k[7:]61                state_dict[name] = v62 63            model.load_state_dict(state_dict)64        else:65            model.load_state_dict(checkpoint)66 67    model.eval()68 69    return model70