Delete src/photosketch

src/photosketch/README.md

@@ -1,57 +0,0 @@
-# Photo-Sketching: Inferring Contour Drawings from Images
-
-<p align="center"><img alt="Teaser" src="doc/teaser.jpg"></p>
-
-This repo contains the training & testing code for our sketch generator. We also provide a [[pre-trained model]](https://drive.google.com/file/d/1TQf-LyS8rRDDapdcTnEgWzYJllPgiXdj/view).
-
-For technical details and the dataset, please refer to the [**[paper]**](https://arxiv.org/abs/1901.00542) and the [**[project page]**](http://www.cs.cmu.edu/~mengtial/proj/sketch/).
-
-# Setting up
-
-The code is now updated to use PyTorch 0.4 and runs on Windows, Mac and Linux. For the obsolete version with PyTorch 0.3 (Linux only), please check out the branch [pytorch-0.3-obsolete](../../tree/pytorch-0.3-obsolete).
-
-Windows users should find the corresponding `.cmd` files instead of `.sh` files mentioned below.
-
-## One-line installation (with Conda environments)
-`conda env create -f environment.yml`
-
-Then activate the environment (sketch) and you are ready to go!
-
-See [here](https://conda.io/docs/user-guide/tasks/manage-environments.html) for more information about conda environments.
-
-## Manual installation
-See `environment.yml` for a list of dependencies.
-
-# Using the pre-trained model
-
-- Download the [pre-trained model](https://drive.google.com/file/d/1TQf-LyS8rRDDapdcTnEgWzYJllPgiXdj/view)
-- Modify the path in `scripts/test_pretrained.sh`
-- From the repo's **root directory**, run `scripts/test_pretrained.sh`
-
-It supports a folder of images as input.
-
-# Train & test on our contour drawing dataset
-
-- Download the images and the rendered sketch from the [project page](http://www.cs.cmu.edu/~mengtial/proj/sketch/)
-- Unzip and organize them into the following structure:
-<p align="center"><img alt="File structure" src="doc/file_structure.png"></p>
-
-- Modify the path in `scripts/train.sh` and `scripts/test.sh`
-- From the repo's **root directory**, run `scripts/train.sh` to train the model
-- From the repo's **root directory**, run `scripts/test.sh` to test on the val set or the test set (specified by the phase flag)
-
-## Citation
-If you use the code or the data for your research, please cite the paper:
-
-```
-@article{LIPS2019,
-  title={Photo-Sketching: Inferring Contour Drawings from Images},
-  author={Li, Mengtian and Lin, Zhe and M\v ech, Radom\'ir and and Yumer, Ersin and Ramanan, Deva},
-  journal={WACV},
-  year={2019}
-}
-```
-
-## Acknowledgement
-This code is based on an old version of [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/).
-

src/photosketch/data/base_data_loader.py

@@ -1,14 +0,0 @@
-
-class BaseDataLoader():
-    def __init__(self):
-        pass
-    
-    def initialize(self, opt):
-        self.opt = opt
-        pass
-
-    def load_data():
-        return None
-
-        
-        

src/photosketch/data/base_dataset.py

@@ -1,45 +0,0 @@
-import torch.utils.data as data
-from PIL import Image
-import torchvision.transforms as transforms
-
-class BaseDataset(data.Dataset):
-    def __init__(self):
-        super(BaseDataset, self).__init__()
-
-    def name(self):
-        return 'BaseDataset'
-
-    def initialize(self, opt):
-        pass
-
-def get_transform(opt):
-    transform_list = []
-    if opt.resize_or_crop == 'resize_and_crop':
-        osize = [opt.loadSize, opt.loadSize]
-        transform_list.append(transforms.Scale(osize, Image.BICUBIC))
-        transform_list.append(transforms.RandomCrop(opt.fineSize))
-    elif opt.resize_or_crop == 'crop':
-        transform_list.append(transforms.RandomCrop(opt.fineSize))
-    elif opt.resize_or_crop == 'scale_width':
-        transform_list.append(transforms.Lambda(
-            lambda img: __scale_width(img, opt.fineSize)))
-    elif opt.resize_or_crop == 'scale_width_and_crop':
-        transform_list.append(transforms.Lambda(
-            lambda img: __scale_width(img, opt.loadSize)))
-        transform_list.append(transforms.RandomCrop(opt.fineSize))
-
-    if opt.isTrain and not opt.no_flip:
-        transform_list.append(transforms.RandomHorizontalFlip())
-
-    transform_list += [transforms.ToTensor(),
-                       transforms.Normalize((0.5, 0.5, 0.5),
-                                            (0.5, 0.5, 0.5))]
-    return transforms.Compose(transform_list)
-
-def __scale_width(img, target_width):
-    ow, oh = img.size
-    if (ow == target_width):
-        return img
-    w = target_width
-    h = int(target_width * oh / ow)
-    return img.resize((w, h), Image.BICUBIC)

src/photosketch/data/custom_dataset_data_loader.py

@@ -1,44 +0,0 @@
-import torch.utils.data
-from data.base_data_loader import BaseDataLoader
-
-
-def CreateDataset(opt):
-    dataset = None
-    if opt.dataset_mode == '1_to_n':
-        from data.paired_1_to_n_dataset import Paired1ToNDataset
-        dataset = Paired1ToNDataset()
-    elif opt.dataset_mode == 'test_dir':
-        from data.test_dir_dataset import TestDirDataset
-        dataset = TestDirDataset()
-    else:
-        raise ValueError("Dataset [%s] not recognized." % opt.dataset_mode)
-
-    print("dataset [%s] was created" % (dataset.name()))
-    dataset.initialize(opt)
-    return dataset
-
-
-class CustomDatasetDataLoader(BaseDataLoader):
-    def name(self):
-        return 'CustomDatasetDataLoader'
-
-    def initialize(self, opt):
-        BaseDataLoader.initialize(self, opt)
-        self.dataset = CreateDataset(opt)
-        self.dataloader = torch.utils.data.DataLoader(
-            self.dataset,
-            batch_size=opt.batchSize,
-            shuffle=not opt.serial_batches,
-            num_workers=int(opt.nThreads))
-
-    def load_data(self):
-        return self
-
-    def __len__(self):
-        return min(len(self.dataset), self.opt.max_dataset_size)
-
-    def __iter__(self):
-        for i, data in enumerate(self.dataloader):
-            if i >= self.opt.max_dataset_size:
-                break
-            yield data

src/photosketch/data/data_loader.py

@@ -1,7 +0,0 @@
-
-def CreateDataLoader(opt):
-    from data.custom_dataset_data_loader import CustomDatasetDataLoader
-    data_loader = CustomDatasetDataLoader()
-    print(data_loader.name())
-    data_loader.initialize(opt)
-    return data_loader

src/photosketch/data/image_folder.py

@@ -1,75 +0,0 @@
-###############################################################################
-# Code from
-# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
-# Modified the original code so that it also loads images from the current
-# directory as well as the subdirectories
-###############################################################################
-
-import torch.utils.data as data
-
-from PIL import Image
-import os
-import os.path
-
-import platform
-
-if platform.system() == 'Windows':
-    IMG_EXTENSIONS = [
-        '.jpg', '.jpeg', '.png', '.ppm', '.bmp',
-    ]
-else:
-    IMG_EXTENSIONS = [
-        '.jpg', '.JPG', '.jpeg', '.JPEG',
-        '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
-    ]
-
-
-def is_image_file(filename):
-    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
-
-
-def make_dataset(dir):
-    images = []
-    assert os.path.isdir(dir), '%s is not a valid directory' % dir
-
-    for root, _, fnames in sorted(os.walk(dir)):
-        for fname in fnames:
-            if is_image_file(fname):
-                path = os.path.join(root, fname)
-                images.append(path)
-
-    return images
-
-
-def default_loader(path):
-    return Image.open(path).convert('RGB')
-
-
-class ImageFolder(data.Dataset):
-
-    def __init__(self, root, transform=None, return_paths=False,
-                 loader=default_loader):
-        imgs = make_dataset(root)
-        if len(imgs) == 0:
-            raise(RuntimeError("Found 0 images in: " + root + "\n"
-                               "Supported image extensions are: " +
-                               ",".join(IMG_EXTENSIONS)))
-
-        self.root = root
-        self.imgs = imgs
-        self.transform = transform
-        self.return_paths = return_paths
-        self.loader = loader
-
-    def __getitem__(self, index):
-        path = self.imgs[index]
-        img = self.loader(path)
-        if self.transform is not None:
-            img = self.transform(img)
-        if self.return_paths:
-            return img, path
-        else:
-            return img
-
-    def __len__(self):
-        return len(self.imgs)

src/photosketch/data/paired_1_to_n_dataset.py

@@ -1,99 +0,0 @@
-import os.path
-import random
-import torchvision.transforms as transforms
-import torch
-from data.base_dataset import BaseDataset
-from data.image_folder import make_dataset
-from PIL import Image
-from data.rotate_and_crop import rotate_and_crop
-
-class Paired1ToNDataset(BaseDataset):
-    def initialize(self, opt):
-        self.opt = opt
-        self.root = opt.dataroot
-        self.if_aug = opt.phase == 'train'
-        self.img_dir = os.path.join(self.root, 'image')
-        self.skt_dir = os.path.join(self.root, opt.render_dir, opt.aug_folder)
-        list_file = os.path.join(self.root, 'list', opt.phase + '.txt')
-        with open(list_file) as f:
-            content = f.readlines()
-        self.list = sorted([x.strip() for x in content])
-
-    def __getitem__(self, index):
-        N = self.opt.nGT
-        filename = self.list[index]
-
-        if_crop_or_rotate = self.if_aug and ((self.opt.crop or self.opt.rotate) and random.random() < 0.8)
-        # even with augmentation on, 20% chance does nothing
-        if_flip = self.if_aug and ((not self.opt.no_flip) and random.random() < 0.5)
-
-        pathA = os.path.join(self.img_dir, filename + '.jpg')
-        img = Image.open(pathA)
-        fine_size = self.opt.fineSize
-
-        if if_flip:
-            img = img.transpose(Image.FLIP_LEFT_RIGHT)
-        if self.if_aug and self.opt.color_jitter:
-            jitter_amount = 0.02
-            img = transforms.ColorJitter(jitter_amount, jitter_amount, jitter_amount, jitter_amount)(img)
-        if if_crop_or_rotate:
-            load_size = self.opt.loadSize
-            img = img.resize((load_size, load_size), Image.BICUBIC)
-            if self.opt.rotate:
-                rot_deg = 5*random.randint(-3, 3)
-                img = rotate_and_crop(img, rot_deg, True)
-            if self.opt.crop:
-                img = transforms.ToTensor()(img)
-                w_offset = random.randint(0, max(0, load_size - fine_size - 1))
-                h_offset = random.randint(0, max(0, load_size - fine_size - 1))
-                img = img[:, h_offset:h_offset + fine_size,
-                          w_offset:w_offset + fine_size]
-            else:
-                img = img.resize((fine_size, fine_size), Image.BICUBIC)
-                img = transforms.ToTensor()(img)
-        else:
-            img = img.resize((fine_size, fine_size), Image.BICUBIC)
-            img = transforms.ToTensor()(img)
-
-        if self.opt.inverse_gamma:
-            linear_mask = (img <= 0.04045).float()
-            exponential_mask = (img > 0.04045).float()
-            img = (img / 12.92 * linear_mask) + (((img + 0.055) / 1.055) ** 2.4) * exponential_mask
-
-        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(img)
-
-        B = N*[None]
-        for i in range(N):
-            pathB = os.path.join(self.skt_dir, '%s_%02d.png' % (filename, i+1))
-            img = Image.open(pathB)
-            if if_flip:
-                img = img.transpose(Image.FLIP_LEFT_RIGHT)
-            if if_crop_or_rotate:
-                load_size = self.opt.loadSize
-                img = img.resize((load_size, load_size), Image.BICUBIC)
-                if self.opt.rotate:
-                    img = rotate_and_crop(img, rot_deg, True)
-                if self.opt.crop:
-                    img = transforms.ToTensor()(img)
-                    img = img[:, h_offset:h_offset + fine_size,
-                              w_offset:w_offset + fine_size]
-                else:
-                    img.resize((fine_size, fine_size), Image.BICUBIC)
-                    img = transforms.ToTensor()(img)
-            else:
-                img = img.resize((fine_size, fine_size), Image.BICUBIC)
-                img = transforms.ToTensor()(img)
-
-            B[i] = transforms.Normalize((0.5,), (0.5,))(img)
-
-        B = torch.cat(B, 0)
-
-        return {'A': A, 'B': B,
-                'A_paths': pathA, 'B_paths': pathB}
-
-    def __len__(self):
-        return len(self.list)
-
-    def name(self):
-        return 'Paired1ToNDataset'
-

src/photosketch/data/rotate_and_crop.py

@@ -1,63 +0,0 @@
-import numpy as np
-from PIL import Image
-
-def perp(a):
-    # https://stackoverflow.com/questions/3252194/numpy-and-line-intersections
-    b = np.empty_like(a)
-    b[0] = -a[1]
-    b[1] = a[0]
-    return b
-
-def seg_intersect(a1,a2, b1,b2):
-    # https://stackoverflow.com/questions/3252194/numpy-and-line-intersections
-    da = a2-a1
-    db = b2-b1
-    dp = a1-b1
-    dap = perp(da)
-    denom = np.dot(dap, db)
-    num = np.dot(dap, dp)
-    return (num / denom.astype(float))*db + b1
-
-def rotate_and_crop(img, deg, same_size=False, interp=Image.BICUBIC):
-    # let the four corners of a rectangle to be ABCD, clockwise
-    if deg == 0:
-        return img
-
-    w, h = img.size
-
-    A = np.array([-w/2, h/2])
-    B = np.array([w/2, h/2])
-    C = np.array([w/2, -h/2])
-    D = np.array([-w/2, -h/2])
-
-    rad = np.radians(deg)
-    c, s = np.cos(rad), np.sin(rad)
-    R = np.array([[c, -s], [s, c]]).T
-
-    Arot = np.dot(A, R)
-    Brot = np.dot(B, R)
-    if deg > 0:
-        X = seg_intersect(A, C, Arot, Brot)
-        offset = X - A
-        offset[1] = -offset[1]
-    else:
-        X = seg_intersect(B, D, Arot, Brot)
-        offset = B - X
-    
-    if same_size:
-        wh_org = np.array([w, h])
-        wh = np.ceil(np.divide(np.square(wh_org), wh_org - 2*offset)).astype(np.int32)
-        offset = (wh - wh_org)/2
-        img = img.resize(wh, interp)
-        w = wh[0]
-        h = wh[1]
-    else:
-        offset = np.ceil(offset)
-    img = img.rotate(deg, interp)
-    return img.crop(
-            (offset[0], 
-                offset[1],
-                w - offset[0],
-                h - offset[1])
-        )
-

src/photosketch/data/test_dir_dataset.py

@@ -1,38 +0,0 @@
-import os.path
-from glob import glob
-import torchvision.transforms as transforms
-import torch
-from data.base_dataset import BaseDataset
-from data.image_folder import make_dataset, IMG_EXTENSIONS
-from PIL import Image
-
-
-class TestDirDataset(BaseDataset):
-    def initialize(self, opt):
-        self.opt = opt
-        self.in_dir = opt.dataroot
-        if opt.file_name:
-            self.list = [os.path.join(self.in_dir, opt.file_name)]
-        else:
-            self.list = []
-            for ext in IMG_EXTENSIONS:
-                self.list.extend(glob(os.path.join(self.in_dir, '*' + ext)))
-
-    def __getitem__(self, index):
-        file_path = self.list[index]
-        img = Image.open(file_path)
-        w, h = img.size
-        fine_size = self.opt.fineSize
-        img = img.resize((fine_size, fine_size), Image.BICUBIC)
-        img = transforms.ToTensor()(img)
-        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(img)
-        B = A.clone()[0]
-
-        return {'A': A, 'B': B, 'A_paths': file_path, 'B_paths': file_path, 'w': w, 'h': h}
-
-    def __len__(self):
-        return len(self.list)
-
-    def name(self):
-        return 'TestDirDataset'
-

src/photosketch/inference_utils.py

@@ -1,109 +0,0 @@
-import os
-import json
-
-from PIL import Image
-from .models.models import create_model
-from .data.data_loader import CreateDataLoader
-from .options.inference_options import InferenceOptions
-from torchvision import transforms
-import numpy as np
-
-class SketchGenerator:
-
-    def __init__(self, config_path, photosketch_dir, results_dir, img_dir, name='pretrained', nThreads=1, batchSize=1, serial_batches=True, no_flip=True):
-        self.config_path = config_path
-        self.photosketch_dir = photosketch_dir
-        self.results_dir = results_dir
-        self.img_dir = img_dir
-
-        self.opt = self._create_options(name=name, nThreads=nThreads, batchSize=batchSize, serial_batches=serial_batches, no_flip=no_flip)
-        self.model = create_model(self.opt)
-
-    def img_to_sketch(self, img_path):
-        input_data = self.preprocess_input(img_path, self.opt)
-        self.model.set_input(input_data)
-        img_path = self.model.get_image_paths()
-        self.model.test()
-        image_numpy = self.model.fake_B.detach()[0][0].cpu().float().numpy()
-        image_numpy = (image_numpy + 1) / 2.0 * 255.0
-        image_pil = Image.fromarray(image_numpy.astype(np.uint8))
-        image_pil = image_pil.resize((self.model.input_w[0], self.model.input_h[0]), Image.BICUBIC)
-        return image_pil 
-
-    def _read_default_options(self):
-        with open(self.config_path) as f:
-            opt = json.load(f)
-        return opt
-
-    def _create_options(self, name='pretrained', nThreads=1, batchSize=1, serial_batches=True, no_flip=True):
-        loaded_opt = self._read_default_options()
-        loaded_opt['checkpoints_dir'] = self.photosketch_dir
-        loaded_opt['results_dir'] = self.results_dir
-        loaded_opt['name'] = name
-        loaded_opt['dataroot'] = self.img_dir
-        opt = InferenceOptions().parse(loaded_opt)
-        opt.nThreads = nThreads
-        opt.batchSize = batchSize
-        opt.serial_batches = serial_batches
-        opt.no_flip = no_flip
-        return opt
-
-    def preprocess_input(self, image_path):
-        """
-        Given an image path, preprocess the image to be used as input for the model
-        """
-        img = Image.open(image_path)
-        w, h = img.size
-        fine_size = self.opt.fineSize
-        img = img.resize((fine_size, fine_size), Image.BICUBIC)
-        img = transforms.ToTensor()(img)
-        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(img).unsqueeze(0)
-        B = A.clone()[0]
-        return {'A': A, 'B': B, 'A_paths': [image_path], 'B_paths': [image_path], 'w': [w], 'h': [h]}
-
-    def preprocess_PIL(self, img):
-        """
-        Given a PIL image, preprocess the image to be used as input for the model
-        """
-        w, h = img.size
-        fine_size = self.opt.fineSize
-        img = img.resize((fine_size, fine_size), Image.BICUBIC)
-        img = transforms.ToTensor()(img)
-        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(img).unsqueeze(0)
-        B = A.clone()[0]
-        return {'A': A, 'B': B, 'A_paths': [''], 'B_paths': [""], 'w': [w], 'h': [h]}
-
-    def get_dataloader(self):
-        data_loader = CreateDataLoader(self.opt)
-        dataset = data_loader.load_data()
-        return dataset
-
-
-    def path_to_sketch(self, img_path: str) -> Image:
-        """
-        Given a single path to an image, generate a sketch and return the PIL image
-        """
-        input_data = self.preprocess_input(img_path)
-        self.model.set_input(input_data)
-        img_path = self.model.get_image_paths()
-        self.model.test()
-        image_numpy = self.model.fake_B.detach()[0][0].cpu().float().numpy()
-        image_numpy = (image_numpy + 1) / 2.0 * 255.0
-        image_pil = Image.fromarray(image_numpy.astype(np.uint8))
-        image_pil = image_pil.resize((self.model.input_w[0], self.model.input_h[0]), Image.BICUBIC)
-        return image_pil
-
-    def PIL_to_sketch(self, img: Image) -> Image:
-        """
-        Given a single PIL image, generate a sketch and return the PIL image
-        """
-        input_data = self.preprocess_PIL(img)
-        self.model.set_input(input_data)
-        img_path = self.model.get_image_paths()
-        self.model.test()
-        image_numpy = self.model.fake_B.detach()[0][0].cpu().float().numpy()
-        image_numpy = (image_numpy + 1) / 2.0 * 255.0
-        image_pil = Image.fromarray(image_numpy.astype(np.uint8))
-        image_pil = image_pil.resize((self.model.input_w[0], self.model.input_h[0]), Image.BICUBIC)
-        return image_pil
-

src/photosketch/models/base_model.py

@@ -1,69 +0,0 @@
-import os
-import torch
-
-
-class BaseModel():
-    def name(self):
-        return 'BaseModel'
-
-    def initialize(self, opt):
-        self.opt = opt
-        self.isTrain = opt.isTrain
-        self.device = torch.device("cuda" if opt.use_cuda else "cpu")
-        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
-
-    def set_input(self, input):
-        self.input = input
-
-    def forward(self):
-        pass
-
-    # used in test time, no backprop
-    def test(self):
-        pass
-
-    def get_image_paths(self):
-        pass
-
-    def optimize_parameters(self):
-        pass
-
-    def get_current_visuals(self):
-        return self.input
-
-    def get_current_errors(self):
-        return {}
-
-    def save(self, label):
-        pass
-
-    # helper saving function that can be used by subclasses
-    def save_network(self, network, network_label, epoch_label):
-        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
-        save_path = os.path.join(self.save_dir, save_filename)
-        torch.save(network.cpu().state_dict(), save_path)
-        network = network.to(self.device)
-
-    # helper loading function that can be used by subclasses
-    def load_network(self, network, network_label, epoch_label):
-        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
-        if self.opt.pretrain_path:
-            save_path = os.path.join(self.opt.pretrain_path, save_filename)
-        else:
-            save_path = os.path.join(self.save_dir, save_filename)
-        network.load_state_dict(torch.load(save_path))
-
-    # update learning rate (called once every epoch)
-    def update_learning_rate(self):
-        for scheduler in self.schedulers:
-            scheduler.step()
-        lr = self.optimizers[0].param_groups[0]['lr']
-        print('learning rate = %.7f' % lr)
-
-    def set_requires_grad(self, nets, requires_grad=False):
-        if not isinstance(nets, list):
-            nets = [nets]
-        for net in nets:
-            if net is not None:
-                for param in net.parameters():
-                    param.requires_grad = requires_grad

src/photosketch/models/models.py

@@ -1,12 +0,0 @@
-
-def create_model(opt):
-    model = None
-    print(opt.model)
-    if opt.model == 'pix2pix':
-        from .pix2pix_model import Pix2PixModel
-        model = Pix2PixModel()
-    else:
-        raise ValueError("Model [%s] not recognized." % opt.model)
-    model.initialize(opt)
-    print("model [%s] was created" % (model.name()))
-    return model

src/photosketch/models/networks.py

@@ -1,557 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.nn import init
-import functools
-from torch.autograd import Variable
-from torch.optim import lr_scheduler
-import numpy as np
-###############################################################################
-# Functions
-###############################################################################
-
-
-def weights_init_normal(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if classname.find('Conv') != -1:
-        init.normal_(m.weight.data, 0.0, 0.02)
-    elif classname.find('Linear') != -1:
-        init.normal_(m.weight.data, 0.0, 0.02)
-    elif classname.find('BatchNorm2d') != -1:
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_xavier(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if classname.find('Conv') != -1:
-        init.xavier_normal_(m.weight.data, gain=0.02)
-    elif classname.find('Linear') != -1:
-        init.xavier_normal_(m.weight.data, gain=0.02)
-    elif classname.find('BatchNorm2d') != -1:
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_kaiming(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if classname.find('Conv') != -1:
-        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
-    elif classname.find('Linear') != -1:
-        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
-    elif classname.find('BatchNorm2d') != -1:
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_orthogonal(m):
-    classname = m.__class__.__name__
-    print(classname)
-    if classname.find('Conv') != -1:
-        init.orthogonal_(m.weight.data, gain=1)
-    elif classname.find('Linear') != -1:
-        init.orthogonal_(m.weight.data, gain=1)
-    elif classname.find('BatchNorm2d') != -1:
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def init_weights(net, init_type='normal'):
-    print('initialization method [%s]' % init_type)
-    if init_type == 'normal':
-        net.apply(weights_init_normal)
-    elif init_type == 'xavier':
-        net.apply(weights_init_xavier)
-    elif init_type == 'kaiming':
-        net.apply(weights_init_kaiming)
-    elif init_type == 'orthogonal':
-        net.apply(weights_init_orthogonal)
-    else:
-        raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
-
-
-def get_norm_layer(norm_type='instance'):
-    if norm_type == 'batch':
-        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
-    elif norm_type == 'instance':
-        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
-    elif norm_type == 'none':
-        norm_layer = None
-    else:
-        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
-    return norm_layer
-
-
-def get_scheduler(optimizer, opt):
-    if opt.lr_policy == 'lambda':
-        def lambda_rule(epoch):
-            lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
-            return lr_l
-        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
-    elif opt.lr_policy == 'step':
-        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
-    elif opt.lr_policy == 'plateau':
-        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
-    else:
-        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
-    return scheduler
-
-
-def define_G(input_nc, output_nc, ngf, which_model_netG, norm='batch', use_dropout=False, init_type='normal'):
-    netG = None
-    norm_layer = get_norm_layer(norm_type=norm)
-
-    if which_model_netG == 'resnet_9blocks':
-        netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
-    elif which_model_netG == 'resnet_6blocks':
-        netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
-    elif which_model_netG == 'unet_128':
-        netG = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    elif which_model_netG == 'unet_256':
-        netG = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    else:
-        raise NotImplementedError('Generator model name [%s] is not recognized' % which_model_netG)
-    
-    init_weights(netG, init_type=init_type)
-    return netG
-
-
-def define_D(input_nc, ndf, which_model_netD,
-             n_layers_D=3, norm='batch', use_sigmoid=False, init_type='normal'):
-    netD = None
-    norm_layer = get_norm_layer(norm_type=norm)
-
-    if which_model_netD == 'basic':
-        netD = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
-    elif which_model_netD == 'n_layers':
-        netD = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
-    elif which_model_netD == 'pixel':
-        netD = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
-    elif which_model_netD == 'global':
-        netD = GlobalDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
-    elif which_model_netD == 'global_np':
-        netD = GlobalNPDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
-    else:
-        raise NotImplementedError('Discriminator model name [%s] is not recognized' %
-                                  which_model_netD)
-    init_weights(netD, init_type=init_type)
-    return netD
-
-
-def print_network(net):
-    num_params = 0
-    for param in net.parameters():
-        num_params += param.numel()
-    print(net)
-    print('Total number of parameters: %d' % num_params)
-
-
-##############################################################################
-# Classes
-##############################################################################
-
-
-# Defines the GAN loss which uses either LSGAN or the regular GAN.
-# When LSGAN is used, it is basically same as MSELoss,
-# but it abstracts away the need to create the target label tensor
-# that has the same size as the input
-class GANLoss(nn.Module):
-    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
-                 device='cpu'):
-        super(GANLoss, self).__init__()
-        self.device = device
-        self.real_label = target_real_label
-        self.fake_label = target_fake_label
-        self.real_label_var = None
-        self.fake_label_var = None
-        if use_lsgan:
-            self.loss = nn.MSELoss().to(device)
-        else:
-            self.loss = nn.BCELoss().to(device)
-
-    def get_target_tensor(self, input, target_is_real):
-        target_tensor = None
-        if target_is_real:
-            create_label = ((self.real_label_var is None) or
-                            (self.real_label_var.numel() != input.numel()))
-            if create_label:
-                self.real_label_var = torch.full(input.size(), self.real_label, requires_grad=False, device=self.device)
-            target_tensor = self.real_label_var
-        else:
-            create_label = ((self.fake_label_var is None) or
-                            (self.fake_label_var.numel() != input.numel()))
-            if create_label:
-                self.fake_label_var = torch.full(input.size(), self.fake_label, requires_grad=False, device=self.device)
-            target_tensor = self.fake_label_var
-        return target_tensor
-
-    def __call__(self, input, target_is_real):
-        target_tensor = self.get_target_tensor(input, target_is_real)
-        return self.loss(input, target_tensor)
-
-
-# Defines the generator that consists of Resnet blocks between a few
-# downsampling/upsampling operations.
-# Code and idea originally from Justin Johnson's architecture.
-# https://github.com/jcjohnson/fast-neural-style/
-class ResnetGenerator(nn.Module):
-    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
-        assert(n_blocks >= 0)
-        super(ResnetGenerator, self).__init__()
-        self.input_nc = input_nc
-        self.output_nc = output_nc
-        self.ngf = ngf
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        model = [nn.ReflectionPad2d(3),
-                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
-                           bias=use_bias),
-                 norm_layer(ngf),
-                 nn.ReLU(True)]
-
-        n_downsampling = 2
-        for i in range(n_downsampling):
-            mult = 2**i
-            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
-                                stride=2, padding=1, bias=use_bias),
-                      norm_layer(ngf * mult * 2),
-                      nn.ReLU(True)]
-
-        mult = 2**n_downsampling
-        for i in range(n_blocks):
-            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
-
-        for i in range(n_downsampling):
-            mult = 2**(n_downsampling - i)
-            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
-                                         kernel_size=3, stride=2,
-                                         padding=1, output_padding=1,
-                                         bias=use_bias),
-                      norm_layer(int(ngf * mult / 2)),
-                      nn.ReLU(True)]
-        model += [nn.ReflectionPad2d(3)]
-        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
-        model += [nn.Tanh()]
-
-        self.model = nn.Sequential(*model)
-
-    def forward(self, input):
-        return self.model(input)
-
-
-# Define a resnet block
-class ResnetBlock(nn.Module):
-    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
-        super(ResnetBlock, self).__init__()
-        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
-
-    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
-        conv_block = []
-        p = 0
-        if padding_type == 'reflect':
-            conv_block += [nn.ReflectionPad2d(1)]
-        elif padding_type == 'replicate':
-            conv_block += [nn.ReplicationPad2d(1)]
-        elif padding_type == 'zero':
-            p = 1
-        else:
-            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
-                       norm_layer(dim),
-                       nn.ReLU(True)]
-        if use_dropout:
-            conv_block += [nn.Dropout(0.5)]
-        else:
-            conv_block += [nn.Dropout(0)]
-
-        p = 0
-        if padding_type == 'reflect':
-            conv_block += [nn.ReflectionPad2d(1)]
-        elif padding_type == 'replicate':
-            conv_block += [nn.ReplicationPad2d(1)]
-        elif padding_type == 'zero':
-            p = 1
-        else:
-            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
-                       norm_layer(dim)]
-
-        return nn.Sequential(*conv_block)
-
-    def forward(self, x):
-        out = x + self.conv_block(x)
-        return out
-
-# Defines the Unet generator.
-# |num_downs|: number of downsamplings in UNet. For example,
-# if |num_downs| == 7, image of size 128x128 will become of size 1x1
-# at the bottleneck
-class UnetGenerator(nn.Module):
-    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
-                 norm_layer=nn.BatchNorm2d, use_dropout=False):
-        super(UnetGenerator, self).__init__()
-        
-        # construct unet structure
-        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
-        for i in range(num_downs - 5):
-            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
-        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
-
-        self.model = unet_block
-
-    def forward(self, input):
-        return self.model(input)
-
-
-# Defines the submodule with skip connection.
-# X -------------------identity---------------------- X
-#   |-- downsampling -- |submodule| -- upsampling --|
-class UnetSkipConnectionBlock(nn.Module):
-    def __init__(self, outer_nc, inner_nc, input_nc=None,
-                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
-        super(UnetSkipConnectionBlock, self).__init__()
-        self.outermost = outermost
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-        if input_nc is None:
-            input_nc = outer_nc
-        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
-                             stride=2, padding=1, bias=use_bias)
-        downrelu = nn.LeakyReLU(0.2, True)
-        downnorm = norm_layer(inner_nc)
-        uprelu = nn.ReLU(True)
-        upnorm = norm_layer(outer_nc)
-
-        if outermost:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1)
-            down = [downconv]
-            up = [uprelu, upconv, nn.Tanh()]
-            model = down + [submodule] + up
-        elif innermost:
-            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv]
-            up = [uprelu, upconv, upnorm]
-            model = down + up
-        else:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv, downnorm]
-            up = [uprelu, upconv, upnorm]
-
-            if use_dropout:
-                model = down + [submodule] + up + [nn.Dropout(0.5)]
-            else:
-                model = down + [submodule] + up
-
-        self.model = nn.Sequential(*model)
-
-    def forward(self, x):
-        if self.outermost:
-            return self.model(x)
-        else:
-            return torch.cat([x, self.model(x)], 1)
-
-
-# Defines the PatchGAN discriminator with the specified arguments.
-class NLayerDiscriminator(nn.Module):
-    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
-        super(NLayerDiscriminator, self).__init__()
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        # 256
-        kw = 4
-        padw = 1
-        sequence = [
-            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
-            nn.LeakyReLU(0.2, True)
-        ]
-
-        # 128
-
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):
-            nf_mult_prev = nf_mult
-            nf_mult = min(2**n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-        # 64
-        # 32
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2**n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                      kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # 31
-
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
-        # 30
-
-        if use_sigmoid:
-            sequence += [nn.Sigmoid()]
-
-        self.model = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        return self.model(input)
-
-class GlobalDiscriminator(nn.Module):
-    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
-        super(GlobalDiscriminator, self).__init__()
-        
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        # 256
-        kw = 4
-        padw = 1
-        sequence = [
-            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
-            nn.LeakyReLU(0.2, True)
-        ]
-
-        # 128
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):
-            nf_mult_prev = nf_mult
-            nf_mult = min(2**n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-        # 64
-        # 32
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2**n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                      kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # 16
-
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=2, padding=0)]
-        sequence += [nn.Conv2d(1, 1, kernel_size=7, stride=1, padding=0)]
-
-        if use_sigmoid:
-            sequence += [nn.Sigmoid()]
-
-        self.model = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        return self.model(input)
-
-class GlobalNPDiscriminator(nn.Module):
-    # no padding
-    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
-        super(GlobalNPDiscriminator, self).__init__()
-        
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        # 256
-        kw = [8, 3, 4]
-        padw = 0
-        sequence = [
-            nn.Conv2d(input_nc, ndf, kernel_size=kw[0], stride=2, padding=padw),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # 125
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):
-            nf_mult_prev = nf_mult
-            nf_mult = min(2**n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                          kernel_size=kw[n], stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-        # 62
-        # 30
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2**n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
-                      kernel_size=4, stride=2, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # 14
-
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=4, stride=2, padding=0)]
-        # 6
-        sequence += [nn.Conv2d(1, 1, kernel_size=6, stride=1, padding=0, bias=use_bias)]
-        # 1
-
-        if use_sigmoid:
-            sequence += [nn.Sigmoid()]
-
-        self.model = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        return self.model(input)
-
-class PixelDiscriminator(nn.Module):
-    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
-        super(PixelDiscriminator, self).__init__()
-        
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-            
-        self.net = [
-            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
-            nn.LeakyReLU(0.2, True),
-            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
-            norm_layer(ndf * 2),
-            nn.LeakyReLU(0.2, True),
-            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
-
-        if use_sigmoid:
-            self.net.append(nn.Sigmoid())
-
-        self.net = nn.Sequential(*self.net)
-
-    def forward(self, input):
-        return self.net(input)
-

src/photosketch/models/pix2pix_model.py

@@ -1,174 +0,0 @@
-import numpy as np
-import torch
-import os
-from collections import OrderedDict
-from torch.autograd import Variable
-from ..util import util
-from ..util.image_pool import ImagePool
-from .base_model import BaseModel
-from . import networks
-from PIL import Image
-
-
-class Pix2PixModel(BaseModel):
-    def name(self):
-        return 'Pix2PixModel'
-
-    def initialize(self, opt):
-        BaseModel.initialize(self, opt)
-        self.isTrain = opt.isTrain
-
-        # load/define networks
-        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
-                                      opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type)
-
-        if self.isTrain:
-            use_sigmoid = opt.no_lsgan
-            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf,
-                                          opt.which_model_netD,
-                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type)
-
-        if not self.isTrain or opt.continue_train:
-            self.load_network(self.netG, 'G', opt.which_epoch)
-            if self.isTrain:
-                self.load_network(self.netD, 'D', opt.which_epoch)
-
-        self.netG = self.netG.to(self.device)
-        if self.isTrain:
-            self.netD = self.netD.to(self.device)
-
-        if self.isTrain:
-            self.fake_AB_pool = ImagePool(opt.pool_size)
-            self.old_lr = opt.lr
-            # define loss functions
-            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, device=self.device)
-            self.criterionL1 = torch.nn.L1Loss().to(self.device)
-
-            # initialize optimizers
-            self.schedulers = []
-            self.optimizers = []
-            self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
-                                                lr=opt.lr, betas=(opt.beta1, 0.999))
-            self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
-                                                lr=opt.lr, betas=(opt.beta1, 0.999))
-            self.optimizers.append(self.optimizer_G)
-            self.optimizers.append(self.optimizer_D)
-            for optimizer in self.optimizers:
-                self.schedulers.append(networks.get_scheduler(optimizer, opt))
-
-        print('---------- Networks initialized -------------')
-        networks.print_network(self.netG)
-        if self.isTrain:
-            networks.print_network(self.netD)
-        print('-----------------------------------------------')
-
-    def set_input(self, input):
-        AtoB = self.opt.which_direction == 'AtoB'
-        self.input_A = input['A' if AtoB else 'B'].to(self.device)
-        self.input_B = input['B' if AtoB else 'A'].to(self.device)
-        self.image_paths = input['A_paths' if AtoB else 'B_paths']
-        if 'w' in input:
-            self.input_w = input['w']
-        if 'h' in input:
-            self.input_h = input['h']
-
-    def forward(self):
-        self.real_A = self.input_A
-        self.fake_B = self.netG(self.real_A)
-        self.real_B = self.input_B
-
-    # no backprop gradients
-    def test(self):
-        with torch.no_grad():
-            self.forward()
-
-    # get image paths
-    def get_image_paths(self):
-        return self.image_paths
-
-    def backward_D(self):
-        # Fake
-        # stop backprop to the generator by detaching fake_B
-        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1).detach())
-        pred_fake = self.netD(fake_AB.detach())
-        self.loss_D_fake = self.criterionGAN(pred_fake, False)
-
-        # Real
-        n = self.real_B.shape[1]
-        loss_D_real_set = torch.empty(n, device=self.device)
-        for i in range(n):
-            sel_B = self.real_B[:, i, :, :].unsqueeze(1)
-            real_AB = torch.cat((self.real_A, sel_B), 1)
-            pred_real = self.netD(real_AB)
-            loss_D_real_set[i] = self.criterionGAN(pred_real, True)
-        self.loss_D_real = torch.mean(loss_D_real_set)
-
-        # Combined loss
-        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5 * self.opt.lambda_G
-
-        self.loss_D.backward()
-
-    def backward_G(self):
-        # First, G(A) should fake the discriminator
-        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
-        pred_fake = self.netD(fake_AB)
-        self.loss_G_GAN = self.criterionGAN(pred_fake, True) * self.opt.lambda_G
-
-        # Second, G(A) = B
-        n = self.real_B.shape[1]
-        fake_B_expand = self.fake_B.expand(-1, n, -1, -1)
-        L1 = torch.abs(fake_B_expand - self.real_B)
-        L1 = L1.view(-1, n, self.real_B.shape[2]*self.real_B.shape[3])
-        L1 = torch.mean(L1, 2)
-        min_L1, min_idx = torch.min(L1, 1)
-        self.loss_G_L1 = torch.mean(min_L1) * self.opt.lambda_A
-        self.min_idx = min_idx
-
-        self.loss_G = self.loss_G_GAN + self.loss_G_L1
-
-        self.loss_G.backward()
-
-    def optimize_parameters(self):
-        self.forward()
-        # update D
-        self.set_requires_grad(self.netD, True)
-        self.optimizer_D.zero_grad()
-        self.backward_D()
-        self.optimizer_D.step()
-
-        # update G
-        self.set_requires_grad(self.netD, False)
-        self.optimizer_G.zero_grad()
-        self.backward_G()
-        self.optimizer_G.step()
-
-    def get_current_errors(self):
-        return OrderedDict([('G_GAN', self.loss_G_GAN.item()),
-                            ('G_L1', self.loss_G_L1.item()),
-                            ('D_real', self.loss_D_real.item()),
-                            ('D_fake', self.loss_D_fake.item())
-                            ])
-
-    def get_current_visuals(self):
-        real_A = util.tensor2im(self.real_A.detach())
-        fake_B = util.tensor2im(self.fake_B.detach())
-        if self.isTrain:
-            sel_B = self.real_B[:, self.min_idx[0], :, :]
-        else:
-            sel_B = self.real_B[:, 0, :, :]
-        real_B = util.tensor2im(sel_B.unsqueeze(1).detach())
-        return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('real_B', real_B)])
-
-    def save(self, label):
-        self.save_network(self.netG, 'G', label)
-        self.save_network(self.netD, 'D', label)
-
-    def write_image(self, out_dir):
-        image_numpy = self.fake_B.detach()[0][0].cpu().float().numpy()
-        image_numpy = (image_numpy + 1) / 2.0 * 255.0
-        image_pil = Image.fromarray(image_numpy.astype(np.uint8))
-        image_pil = image_pil.resize((self.input_w[0], self.input_h[0]), Image.BICUBIC)
-        name, _ = os.path.splitext(os.path.basename(self.image_paths[0]))
-        out_path = os.path.join(out_dir, name + self.opt.suffix + '.png')
-        image_pil.save(out_path)
-       

src/photosketch/options/base_options.py

@@ -1,83 +0,0 @@
-import argparse
-import os
-from ..util import util
-import torch
-
-
-class BaseOptions():
-    def __init__(self):
-        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-        self.initialized = False
-
-    def initialize(self):
-        self.parser.add_argument('--dataroot', required=True, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
-        self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
-        self.parser.add_argument('--loadSize', type=int, default=286, help='scale images to this size')
-        self.parser.add_argument('--fineSize', type=int, default=256, help='then crop to this size')
-        self.parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels')
-        self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
-        self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
-        self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
-        self.parser.add_argument('--which_model_netD', type=str, default='basic', help='selects model to use for netD')
-        self.parser.add_argument('--which_model_netG', type=str, default='resnet_9blocks', help='selects model to use for netG')
-        self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers')
-        self.parser.add_argument('--no-cuda', action='store_true', default=False, help='disable CUDA training (please use CUDA_VISIBLE_DEVICES to select GPU)')
-        self.parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
-        self.parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single]')
-        self.parser.add_argument('--model', type=str, default='cycle_gan',
-                                 help='chooses which model to use. cycle_gan, pix2pix, test')
-        self.parser.add_argument('--which_direction', type=str, default='AtoB', help='AtoB or BtoA')
-        self.parser.add_argument('--nThreads', default=6, type=int, help='# threads for loading data')
-        self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
-        self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
-        self.parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
-        self.parser.add_argument('--display_winsize', type=int, default=256,  help='display window size')
-        self.parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')
-        self.parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display')
-        self.parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
-        self.parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')
-        self.parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
-        self.parser.add_argument('--resize_or_crop', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
-        self.parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
-        self.parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
-        self.parser.add_argument('--render_dir', type=str, default='sketch-rendered')
-        self.parser.add_argument('--aug_folder', type=str, default='width-5')
-        self.parser.add_argument('--stroke_dir', type=str, default='')
-        self.parser.add_argument('--crop', action='store_true')
-        self.parser.add_argument('--rotate', action='store_true')
-        self.parser.add_argument('--color_jitter', action='store_true')
-        self.parser.add_argument('--stroke_no_couple', action='store_true', help='')
-        self.parser.add_argument('--pretrain_path', type=str, default='')
-        self.parser.add_argument('--nGT', type=int, default=5)
-        self.parser.add_argument('--rot_int_max', type=int, default=3)
-        self.parser.add_argument('--jitter_amount', type=float, default=0.02)
-        self.parser.add_argument('--inverse_gamma', action='store_true')
-        self.parser.add_argument('--img_mean', type=float, nargs='+')
-        self.parser.add_argument('--img_std', type=float, nargs='+')
-        self.parser.add_argument('--lst_file', type=str)
-        self.initialized = True
-
-    def parse(self):
-        if not self.initialized:
-            self.initialize()
-        self.opt = self.parser.parse_args()
-        self.opt.isTrain = self.isTrain   # train or test
-
-        self.opt.use_cuda = not self.opt.no_cuda and torch.cuda.is_available()
-        args = vars(self.opt)
-
-        print('------------ Options -------------')
-        for k, v in sorted(args.items()):
-            print('%s: %s' % (str(k), str(v)))
-        print('-------------- End ----------------')
-
-        # save to the disk
-        expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
-        util.mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, 'opt.txt')
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write('------------ Options -------------\n')
-            for k, v in sorted(args.items()):
-                opt_file.write('%s: %s\n' % (str(k), str(v)))
-            opt_file.write('-------------- End ----------------\n')
-        return self.opt

src/photosketch/options/inference_options.py

@@ -1,43 +0,0 @@
-from .base_options import BaseOptions
-from argparse import Namespace
-import os
-from ..util import util
-import torch
-
-
-class InferenceOptions(BaseOptions):
-    def initialize(self):
-        BaseOptions.initialize(self)
-        self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
-        self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
-        self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
-        self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
-        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
-        self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
-        self.parser.add_argument('--file_name', type=str, default='')
-        self.parser.add_argument('--suffix', type=str, default='')
-        self.isTrain = False
-
-    def parse(self, args_dict):
-        self.opt = Namespace(**args_dict)
-        self.opt.max_dataset_size = float(self.opt.max_dataset_size)
-        self.isTrain = False
-        self.opt.isTrain = self.isTrain
-        self.opt.use_cuda = not self.opt.no_cuda and torch.cuda.is_available()
-        args = vars(self.opt)
-
-        print('------------ Options -------------')
-        for k, v in sorted(args.items()):
-            print('%s: %s' % (str(k), str(v)))
-        print('-------------- End ----------------')
-
-        # save to the disk
-        expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
-        util.mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, 'inference_config.txt')
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write('------------ Options -------------\n')
-            for k, v in sorted(args.items()):
-                opt_file.write('%s: %s\n' % (str(k), str(v)))
-            opt_file.write('-------------- End ----------------\n')
-        return self.opt

src/photosketch/options/opt.json

@@ -1,56 +0,0 @@
-{
-    "aspect_ratio": 1.0,
-    "aug_folder": "width-5",
-    "batchSize": 1,
-    "checkpoints_dir": "",
-    "color_jitter": false,
-    "crop": false,
-    "dataroot": "examples/",
-    "dataset_mode": "test_dir",
-    "display_id": 1,
-    "display_port": 8097,
-    "display_server": "http://localhost",
-    "display_winsize": 256,
-    "file_name": "",
-    "fineSize": 256,
-    "how_many": 50,
-    "img_mean": null,
-    "img_std": null,
-    "init_type": "normal",
-    "input_nc": 3,
-    "inverse_gamma": false,
-    "isTrain": false,
-    "jitter_amount": 0.02,
-    "loadSize": 286,
-    "lst_file": null,
-    "max_dataset_size": "inf",
-    "model": "pix2pix",
-    "nGT": 5,
-    "nThreads": 6,
-    "n_layers_D": 3,
-    "name": "",
-    "ndf": 64,
-    "ngf": 64,
-    "no_cuda": false,
-    "no_dropout": true,
-    "no_flip": false,
-    "norm": "batch",
-    "ntest": "inf",
-    "output_nc": 1,
-    "phase": "test",
-    "pretrain_path": "",
-    "render_dir": "sketch-rendered",
-    "resize_or_crop": "resize_and_crop",
-    "results_dir": "",
-    "rot_int_max": 3,
-    "rotate": false,
-    "serial_batches": false,
-    "stroke_dir": "",
-    "stroke_no_couple": false,
-    "suffix": "",
-    "use_cuda": true,
-    "which_direction": "AtoB",
-    "which_epoch": "latest",
-    "which_model_netD": "basic",
-    "which_model_netG": "resnet_9blocks"
-}

src/photosketch/options/test_options.py

@@ -1,15 +0,0 @@
-from .base_options import BaseOptions
-
-
-class TestOptions(BaseOptions):
-    def initialize(self):
-        BaseOptions.initialize(self)
-        self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
-        self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
-        self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
-        self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
-        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
-        self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
-        self.parser.add_argument('--file_name', type=str, default='')
-        self.parser.add_argument('--suffix', type=str, default='')
-        self.isTrain = False

src/photosketch/options/train_options.py

@@ -1,31 +0,0 @@
-from .base_options import BaseOptions
-
-
-class TrainOptions(BaseOptions):
-    def initialize(self):
-        BaseOptions.initialize(self)
-        self.parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen')
-        self.parser.add_argument('--display_single_pane_ncols', type=int, default=0, help='if positive, display all images in a single visdom web panel with certain number of images per row.')
-        self.parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')
-        self.parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
-        self.parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
-        self.parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
-        self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
-        self.parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
-        self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
-        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
-        self.parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')
-        self.parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
-        self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
-        self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
-        self.parser.add_argument('--no_lsgan', action='store_true', help='do *not* use least square GAN, if false, use vanilla GAN')
-        self.parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
-        self.parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
-        self.parser.add_argument('--lambda_stroke', type=float, default=100, help='weight for stroke generation')
-        self.parser.add_argument('--lambda_G', type=float, default=1, help='weight for GAN')
-        self.parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')
-        self.parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
-        self.parser.add_argument('--lr_policy', type=str, default='lambda', help='learning rate policy: lambda|step|plateau')
-        self.parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
-        self.parser.add_argument('--identity', type=float, default=0.5, help='use identity mapping. Setting identity other than 1 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set optidentity = 0.1')
-        self.isTrain = True

src/photosketch/pretrained/inference_config.txt

@@ -1,56 +0,0 @@
------------- Options -------------
-aspect_ratio: 1.0
-aug_folder: width-5
-batchSize: 1
-checkpoints_dir: src/photosketch
-color_jitter: False
-crop: False
-dataroot: src/photosketch/inputs
-dataset_mode: test_dir
-display_id: 1
-display_port: 8097
-display_server: http://localhost
-display_winsize: 256
-file_name: 
-fineSize: 256
-how_many: 50
-img_mean: None
-img_std: None
-init_type: normal
-input_nc: 3
-inverse_gamma: False
-isTrain: False
-jitter_amount: 0.02
-loadSize: 286
-lst_file: None
-max_dataset_size: inf
-model: pix2pix
-nGT: 5
-nThreads: 6
-n_layers_D: 3
-name: pretrained
-ndf: 64
-ngf: 64
-no_cuda: False
-no_dropout: True
-no_flip: False
-norm: batch
-ntest: inf
-output_nc: 1
-phase: test
-pretrain_path: 
-render_dir: sketch-rendered
-resize_or_crop: resize_and_crop
-results_dir: src/photosketch/outputs
-rot_int_max: 3
-rotate: False
-serial_batches: False
-stroke_dir: 
-stroke_no_couple: False
-suffix: 
-use_cuda: True
-which_direction: AtoB
-which_epoch: latest
-which_model_netD: basic
-which_model_netG: resnet_9blocks
--------------- End ----------------

src/photosketch/pretrained/latest_net_D.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1b40d1a853a9b2a4b4f1b618f095728651dc8c119dfa093eb84ee7884dfc2d44
-size 10896061

src/photosketch/pretrained/latest_net_G.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:442a5cf113b4b544698c1ea71dff20890685b821d0ae853b4442f9ed2fb98f5b
-size 45566034

src/photosketch/scripts/test.cmd

@@ -1,19 +0,0 @@
-@set "dataDir="
-
-python test.py ^
-    --name default ^
-    --dataroot %dataDir%\ContourDrawing\ ^
-    --phase val ^
-    --how_many 100 ^
-    --checkpoints_dir %dataDir%\Exp\PhotoSketch\Checkpoints\ ^
-    --results_dir %dataDir%\Exp\PhotoSketch\Results\ ^
-    --model pix2pix ^
-    --which_direction AtoB ^
-    --dataset_mode 1_to_n ^
-    --norm batch ^
-    --input_nc 3 ^
-    --output_nc 1 ^
-    --which_model_netG resnet_9blocks ^
-    --which_model_netD global_np ^
-    --aug_folder width-5 ^
-    --no_dropout ^

src/photosketch/scripts/test.sh

@@ -1,19 +0,0 @@
-dataDir=
-
-python test.py \
-    --name default \
-    --dataroot ${dataDir}/ContourDrawing/ \
-    --phase val \
-    --how_many 100 \
-    --checkpoints_dir ${dataDir}/Exp/PhotoSketch/Checkpoints/ \
-    --results_dir ${dataDir}/Exp/PhotoSketch/Results/ \
-    --model pix2pix \
-    --which_direction AtoB \
-    --dataset_mode 1_to_n \
-    --norm batch \
-    --input_nc 3 \
-    --output_nc 1 \
-    --which_model_netG resnet_9blocks \
-    --which_model_netD global_np \
-    --aug_folder width-5 \
-    --no_dropout \

src/photosketch/scripts/test_pretrained.cmd

@@ -1,15 +0,0 @@
-@set "dataDir="
-
-python test_pretrained.py ^
-    --name pretrained ^
-    --dataset_mode test_dir ^
-    --dataroot examples\ ^
-    --results_dir %dataDir%\Exp\PhotoSketch\Results\ ^
-    --checkpoints_dir %dataDir%\Exp\PhotoSketch\Checkpoints\ ^
-    --model pix2pix ^
-    --which_direction AtoB ^
-    --norm batch ^
-    --input_nc 3 ^
-    --output_nc 1 ^
-    --which_model_netG resnet_9blocks ^
-    --no_dropout ^

src/photosketch/scripts/test_pretrained.sh

@@ -1,15 +0,0 @@
-dataDir=/home/fgirella/sketch2img/PhotoSketch
-
-python test_pretrained.py \
-    --name pretrained \
-    --dataset_mode test_dir \
-    --dataroot examples/ \
-    --results_dir ${dataDir}/Runs/Results/ \
-    --checkpoints_dir ${dataDir}/ \
-    --model pix2pix \
-    --which_direction AtoB \
-    --norm batch \
-    --input_nc 3 \
-    --output_nc 1 \
-    --which_model_netG resnet_9blocks \
-    --no_dropout \

src/photosketch/scripts/train.cmd

@@ -1,22 +0,0 @@
-@set "dataDir="
-
-python train.py ^
-    --name default ^
-    --dataroot %dataDir%\ContourDrawing\ ^
-    --checkpoints_dir %dataDir%\Exp\PhotoSketch\Checkpoints\ ^
-    --model pix2pix ^
-    --which_direction AtoB ^
-    --dataset_mode 1_to_n ^
-    --no_lsgan ^
-    --norm batch ^
-    --pool_size 0 ^
-    --output_nc 1 ^
-    --which_model_netG resnet_9blocks ^
-    --which_model_netD global_np ^
-    --batchSize 2 ^
-    --lambda_A 200 ^
-    --lr 0.0002 ^
-    --aug_folder width-5 ^
-    --crop --rotate --color_jitter ^
-    --niter 400 ^
-    --niter_decay 400 ^

src/photosketch/scripts/train.sh

@@ -1,22 +0,0 @@
-dataDir=
-
-python train.py \
-    --name default \
-    --dataroot ${dataDir}/ContourDrawing/ \
-    --checkpoints_dir ${dataDir}/Exp/PhotoSketch/Checkpoints/ \
-    --model pix2pix \
-    --which_direction AtoB \
-    --dataset_mode 1_to_n \
-    --no_lsgan \
-    --norm batch \
-    --pool_size 0 \
-    --output_nc 1 \
-    --which_model_netG resnet_9blocks \
-    --which_model_netD global_np \
-    --batchSize 2 \
-    --lambda_A 200 \
-    --lr 0.0002 \
-    --aug_folder width-5 \
-    --crop --rotate --color_jitter \
-    --niter 400 \
-    --niter_decay 400 \

src/photosketch/test.py

@@ -1,37 +0,0 @@
-import time
-import os
-from options.test_options import TestOptions
-from data.data_loader import CreateDataLoader
-from models.models import create_model
-from util.visualizer import Visualizer
-from util import html
-
-def main():
-    opt = TestOptions().parse()
-    opt.nThreads = 1   # test code only supports nThreads = 1
-    opt.batchSize = 1  # test code only supports batchSize = 1
-    opt.serial_batches = True  # no shuffle
-    opt.no_flip = True  # no flip
-
-    data_loader = CreateDataLoader(opt)
-    dataset = data_loader.load_data()
-    model = create_model(opt)
-    visualizer = Visualizer(opt)
-    # create website
-    web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
-    webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
-    # test
-    for i, data in enumerate(dataset):
-        if i >= opt.how_many:
-            break
-        model.set_input(data)
-        model.test()
-        visuals = model.get_current_visuals()
-        img_path = model.get_image_paths()
-        print('%04d: process image... %s' % (i, img_path))
-        visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio)
-
-    webpage.save()
-
-if __name__ == '__main__':
-    main()

src/photosketch/test_pretrained.py

@@ -1,29 +0,0 @@
-import os
-from options.test_options import TestOptions
-from data.data_loader import CreateDataLoader
-from models.models import create_model
-
-def main():
-    opt = TestOptions().parse()
-    opt.nThreads = 1   # test code only supports nThreads = 1
-    opt.batchSize = 1  # test code only supports batchSize = 1
-    opt.serial_batches = True  # no shuffle
-    opt.no_flip = True  # no flip
-
-    if not os.path.isdir(opt.results_dir):
-        os.makedirs(opt.results_dir)
-
-    data_loader = CreateDataLoader(opt)
-    dataset = data_loader.load_data()
-    model = create_model(opt)
-
-    # test
-    for i, data in enumerate(dataset):
-        model.set_input(data)
-        img_path = model.get_image_paths()
-        print('Processing %04d (%s)' % (i+1, img_path[0]))
-        model.test()
-        model.write_image(opt.results_dir)
-
-if __name__ == '__main__':
-    main()

src/photosketch/train.py

@@ -1,61 +0,0 @@
-import time
-from options.train_options import TrainOptions
-from data.data_loader import CreateDataLoader
-from models.models import create_model
-from util.visualizer import Visualizer
-
-import torch, random
-torch.manual_seed(0)
-random.seed(0)
-
-def main():
-    opt = TrainOptions().parse()
-    data_loader = CreateDataLoader(opt)
-    dataset = data_loader.load_data()
-    dataset_size = len(data_loader)
-    print('#training images = %d' % dataset_size)
-
-    model = create_model(opt)
-    visualizer = Visualizer(opt)
-    total_steps = 0
-
-    for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
-        epoch_start_time = time.time()
-        epoch_iter = 0
-
-        for i, data in enumerate(dataset):
-            iter_start_time = time.time()
-            visualizer.reset()
-            total_steps += opt.batchSize
-            epoch_iter += opt.batchSize
-            model.set_input(data)
-            model.optimize_parameters()
-
-            if total_steps % opt.display_freq == 0:
-                save_result = total_steps % opt.update_html_freq == 0
-                visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)
-
-            if total_steps % opt.print_freq == 0:
-                errors = model.get_current_errors()
-                t = (time.time() - iter_start_time) / opt.batchSize
-                visualizer.print_current_errors(epoch, epoch_iter, errors, t)
-                if opt.display_id > 0:
-                    visualizer.plot_current_errors(epoch, float(epoch_iter)/dataset_size, opt, errors)
-
-            if total_steps % opt.save_latest_freq == 0:
-                print('saving the latest model (epoch %d, total_steps %d)' %
-                    (epoch, total_steps))
-                model.save('latest')
-
-        if epoch % opt.save_epoch_freq == 0:
-            print('saving the model at the end of epoch %d, iters %d' %
-                (epoch, total_steps))
-            model.save('latest')
-            model.save(epoch)
-
-        print('End of epoch %d / %d \t Time Taken: %d sec' %
-            (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
-        model.update_learning_rate()
-
-if __name__ == '__main__':
-    main()