Parcourir la source

Merge branch 'release/v0.0.2'

Jérôme BUISINE il y a 4 ans
Parent
commit
714bc0055b
6 fichiers modifiés avec 378 ajouts et 51 suppressions
  1. 2 0
      .gitignore
  2. 9 0
      LICENSE
  3. 27 0
      README.md
  4. 297 0
      ganSynthesisImage_100.py
  5. 40 49
      noise_gan.ipynb
  6. 3 2
      tensorboard.ipynb

+ 2 - 0
.gitignore

@@ -1,3 +1,5 @@
+.vscode
+
 runs
 .python-version
 .ipynb_checkpoints

+ 9 - 0
LICENSE

@@ -0,0 +1,9 @@
+MIT License
+Copyright (c) 2019 prise-3d
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+

+ 27 - 0
README.md

@@ -0,0 +1,27 @@
+# Noise Generation
+
+## Description
+
+Study of how to generate noise filter from monte carlo rendering in synthesis images using Generative Adversarial Network approach. The aim of this project is to reproduce monte carlo noise obtained during rendering.
+
+## How to use ?
+
+```bash
+pip install -r requirements.txt
+```
+
+```bash
+python prepare_data.py
+```
+
+```bash
+python ganSynthesisImage_200.py
+```
+
+## Contributors
+
+* [jbuisine](https://github.com/jbuisine)
+
+## License
+
+[MIT](https://github.com/prise-3d/IPFML/blob/master/LICENSE)

+ 297 - 0
ganSynthesisImage_100.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python
+import random
+import argparse
+import cv2
+import os, sys, getopt
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from tensorboardX import SummaryWriter
+from PIL import Image
+
+import torchvision.utils as vutils
+
+import gym
+import gym.spaces
+
+import numpy as np
+
+log = gym.logger
+log.set_level(gym.logger.INFO)
+
+LATENT_VECTOR_SIZE = 200
+DISCR_FILTERS = 100
+GENER_FILTERS = 100
+BATCH_SIZE = 16
+
+# dimension input image will be rescaled
+IMAGE_SIZE = 100
+input_shape = (3, IMAGE_SIZE, IMAGE_SIZE)
+
+BACKUP_MODEL_NAME = "synthesis_{}_model.pt"
+BACKUP_FOLDER = "saved_models"
+BACKUP_EVERY_ITER = 1
+
+LEARNING_RATE = 0.0001
+REPORT_EVERY_ITER = 10
+SAVE_IMAGE_EVERY_ITER = 20
+MAX_ITERATION = 100000
+
+data_folder = 'synthesis_images/generated_blocks'
+
+class Discriminator(nn.Module):
+    def __init__(self, input_shape):
+        super(Discriminator, self).__init__()
+        # this pipe converges image into the single number
+        self.conv_pipe = nn.Sequential(
+            nn.Conv2d(in_channels=input_shape[0], out_channels=DISCR_FILTERS,
+                      kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=DISCR_FILTERS, out_channels=DISCR_FILTERS*2,
+                      kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(DISCR_FILTERS*2),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=DISCR_FILTERS * 2, out_channels=DISCR_FILTERS * 4,
+                      kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(DISCR_FILTERS * 4),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=DISCR_FILTERS * 4, out_channels=DISCR_FILTERS * 8,
+                      kernel_size=5, stride=2, padding=1),
+            nn.BatchNorm2d(DISCR_FILTERS * 8),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=DISCR_FILTERS * 8, out_channels=1,
+                      kernel_size=5, stride=1, padding=0),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        print("Discriminator output")
+        print("Shape : ", x.shape)
+        conv_out = self.conv_pipe(x)
+        print(conv_out.shape)
+        print(conv_out.view(-1, 1).squeeze(dim=1).shape)
+        return conv_out.view(-1, 1).squeeze(dim=1)
+
+
+class Generator(nn.Module):
+    def __init__(self, output_shape):
+        super(Generator, self).__init__()
+        # pipe deconvolves input vector into (3, 100, 100) image
+        self.pipe = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=LATENT_VECTOR_SIZE, out_channels=GENER_FILTERS * 8,
+                               kernel_size=6, stride=1, padding=0),
+            nn.BatchNorm2d(GENER_FILTERS * 8),
+            nn.ReLU(),
+
+            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 8, out_channels=GENER_FILTERS * 4,
+                               kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(GENER_FILTERS * 4),
+            nn.ReLU(),
+
+            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 4, out_channels=GENER_FILTERS * 2,
+                               kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(GENER_FILTERS * 2),
+            nn.ReLU(),
+
+            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 2, out_channels=GENER_FILTERS,
+                               kernel_size=4, stride=2, padding=0),
+            nn.BatchNorm2d(GENER_FILTERS),
+            nn.ReLU(),
+
+            nn.ConvTranspose2d(in_channels=GENER_FILTERS, out_channels=output_shape[0],
+                               kernel_size=4, stride=2, padding=1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        print("Generator output")
+        print(self.pipe(x).shape)
+        return self.pipe(x)
+
+# here we have to generate our batches from final or noisy synthesis images
+def iterate_batches(batch_size=BATCH_SIZE):
+
+    batch = []
+    images = os.listdir(data_folder)
+    nb_images = len(images)
+
+    while True:
+        i = random.randint(0, nb_images - 1)
+
+        img = Image.open(os.path.join(data_folder, images[i]))
+        img_arr = np.asarray(img)
+
+        new_obs = cv2.resize(img_arr, (IMAGE_SIZE, IMAGE_SIZE))
+        # transform (210, 160, 3) -> (3, 210, 160)
+        new_obs = np.moveaxis(new_obs, 2, 0)
+
+        batch.append(new_obs)
+
+        if len(batch) == batch_size:
+            # Normalising input between -1 to 1
+            batch_np = np.array(batch, dtype=np.float32) * 2.0 / 255.0 - 1.0
+            yield torch.tensor(batch_np)
+            batch.clear()
+
+
+if __name__ == "__main__":
+
+    save_model = False
+    load_model = False
+    p_cuda = False
+
+    #parser = argparse.ArgumentParser()
+    #parser.add_argument("--cuda", default=False, action='store_true', help="Enable cuda computation")
+    #args = parser.parse_args()
+
+    try:
+        opts, args = getopt.getopt(sys.argv[1:], "hflc", ["help=", "folder=", "load=", "cuda="])
+    except getopt.GetoptError:
+        # print help information and exit:
+        print('python ganSynthesisImage_100.py --folder folder_name_to_save --cuda 1')
+        print('python ganSynthesisImage_100.py --load model_name_to_load ')
+        sys.exit(2)
+    for o, a in opts:
+        if o in ("-h", "--help"):
+            print('python ganSynthesisImage_100.py --folder folder_name_to_save --cuda 1')
+            print('python ganSynthesisImage_100.py --load folder_name_to_load ')
+            sys.exit()
+        elif o in ("-f", "--folder"):
+            p_model_folder = a
+            save_model = True
+        elif o in ("-l", "--load"):
+            p_load = a
+            load_model = True
+        elif o in ("-c", "--cuda"):
+            p_cuda = int(a)
+        else:
+            assert False, "unhandled option"
+
+    if save_model and load_model:
+        raise Exception("Cannot save and load model. One argurment in only required.")
+    if not save_model and not load_model:
+        print('python ganSynthesisImage_100.py --folder folder_name_to_save --cuda 1')
+        print('python ganSynthesisImage_100.py --load folder_name_to_load ')
+        print("Need at least one argurment.")
+        sys.exit(2)
+
+    device = torch.device("cuda" if p_cuda else "cpu")
+    #envs = [InputWrapper(gym.make(name)) for name in ('Breakout-v0', 'AirRaid-v0', 'Pong-v0')]
+
+
+    # prepare folder names to save models
+    if save_model:
+
+        models_folder_path = os.path.join(BACKUP_FOLDER, p_model_folder)
+        dis_model_path = os.path.join(models_folder_path, BACKUP_MODEL_NAME.format('disc'))
+        gen_model_path = os.path.join(models_folder_path, BACKUP_MODEL_NAME.format('gen'))
+
+    if load_model:
+
+        models_folder_path = os.path.join(BACKUP_FOLDER, p_load)
+        dis_model_path = os.path.join(models_folder_path, BACKUP_MODEL_NAME.format('disc'))
+        gen_model_path = os.path.join(models_folder_path, BACKUP_MODEL_NAME.format('gen'))
+
+    # Construct model
+    net_discr = Discriminator(input_shape=input_shape).to(device)
+    net_gener = Generator(output_shape=input_shape).to(device)
+    print(net_discr)
+    print(net_gener)
+
+    objective = nn.BCELoss()
+    gen_optimizer = optim.Adam(params=net_gener.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
+    dis_optimizer = optim.Adam(params=net_discr.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
+    writer = SummaryWriter()
+
+    gen_losses = []
+    dis_losses = []
+    iter_no = 0
+
+    true_labels_v = torch.ones(BATCH_SIZE, dtype=torch.float32, device=device)
+    fake_labels_v = torch.zeros(BATCH_SIZE, dtype=torch.float32, device=device)
+
+
+    # load models checkpoint if exists
+    if load_model:
+        gen_checkpoint = torch.load(gen_model_path)
+
+        net_gener.load_state_dict(gen_checkpoint['model_state_dict'])
+        gen_optimizer.load_state_dict(gen_checkpoint['optimizer_state_dict'])
+        gen_losses = gen_checkpoint['gen_losses']
+        iteration = gen_checkpoint['iteration'] # retrieve only from the gen net the iteration number
+
+        dis_checkpoint = torch.load(dis_model_path)
+
+        net_discr.load_state_dict(dis_checkpoint['model_state_dict'])
+        dis_optimizer.load_state_dict(dis_checkpoint['optimizer_state_dict'])
+        dis_losses = dis_checkpoint['dis_losses']
+
+        iter_no = iteration
+
+    for batch_v in iterate_batches():
+
+        # generate extra fake samples, input is 4D: batch, filters, x, y
+        gen_input_v = torch.FloatTensor(BATCH_SIZE, LATENT_VECTOR_SIZE, 1, 1).normal_(0, 1).to(device)
+
+        # There we get data
+        batch_v = batch_v.to(device)
+
+        gen_output_v = net_gener(gen_input_v)
+
+        # train discriminator
+        dis_optimizer.zero_grad()
+        print("True disc")
+        dis_output_true_v = net_discr(batch_v)
+        print("Fake disc")
+        dis_output_fake_v = net_discr(gen_output_v.detach())
+        dis_loss = objective(dis_output_true_v, true_labels_v) + objective(dis_output_fake_v, fake_labels_v)
+        dis_loss.backward()
+        dis_optimizer.step()
+        dis_losses.append(dis_loss.item())
+
+        # train generator
+        gen_optimizer.zero_grad()
+        dis_output_v = net_discr(gen_output_v)
+        gen_loss_v = objective(dis_output_v, true_labels_v)
+        gen_loss_v.backward()
+        gen_optimizer.step()
+        gen_losses.append(gen_loss_v.item())
+
+        iter_no += 1
+        print("Iteration : ", iter_no)
+
+        if iter_no % REPORT_EVERY_ITER == 0:
+            log.info("Iter %d: gen_loss=%.3e, dis_loss=%.3e", iter_no, np.mean(gen_losses), np.mean(dis_losses))
+            writer.add_scalar("gen_loss", np.mean(gen_losses), iter_no)
+            writer.add_scalar("dis_loss", np.mean(dis_losses), iter_no)
+            gen_losses = []
+            dis_losses = []
+
+        if iter_no % SAVE_IMAGE_EVERY_ITER == 0:
+            writer.add_image("fake", vutils.make_grid(gen_output_v.data[:IMAGE_SIZE], normalize=True), iter_no)
+            writer.add_image("real", vutils.make_grid(batch_v.data[:IMAGE_SIZE], normalize=True), iter_no)
+
+        if iter_no % BACKUP_EVERY_ITER == 0:
+            if not os.path.exists(models_folder_path):
+                os.makedirs(models_folder_path)
+
+            torch.save({
+                        'iteration': iter_no,
+                        'model_state_dict': net_gener.state_dict(),
+                        'optimizer_state_dict': gen_optimizer.state_dict(),
+                        'gen_losses': gen_losses
+                    }, gen_model_path)
+
+            torch.save({
+                        'iteration': iter_no,
+                        'model_state_dict': net_discr.state_dict(),
+                        'optimizer_state_dict': dis_optimizer.state_dict(),
+                        'dis_losses': dis_losses
+                    }, dis_model_path)
+
+
+
+
+

+ 40 - 49
noise_gan.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
@@ -11,74 +11,65 @@
      "text": [
       "Discriminator(\n",
       "  (conv_pipe): Sequential(\n",
-      "    (0): Conv2d(3, 200, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (0): Conv2d(3, 100, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "    (1): ReLU()\n",
-      "    (2): Conv2d(200, 400, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (3): BatchNorm2d(400, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (2): Conv2d(100, 200, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (3): BatchNorm2d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (4): ReLU()\n",
-      "    (5): Conv2d(400, 800, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (6): BatchNorm2d(800, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (5): Conv2d(200, 400, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (6): BatchNorm2d(400, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (7): ReLU()\n",
-      "    (8): Conv2d(800, 1600, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (9): BatchNorm2d(1600, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (8): Conv2d(400, 800, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (9): BatchNorm2d(800, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (10): ReLU()\n",
-      "    (11): Conv2d(1600, 3200, kernel_size=(8, 8), stride=(2, 2), padding=(1, 1))\n",
-      "    (12): BatchNorm2d(3200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (11): Conv2d(800, 1600, kernel_size=(8, 8), stride=(2, 2), padding=(1, 1))\n",
+      "    (12): BatchNorm2d(1600, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (13): ReLU()\n",
-      "    (14): Conv2d(3200, 1, kernel_size=(4, 4), stride=(1, 1))\n",
+      "    (14): Conv2d(1600, 1, kernel_size=(4, 4), stride=(1, 1))\n",
       "    (15): Sigmoid()\n",
       "  )\n",
       ")\n",
       "Generator(\n",
       "  (pipe): Sequential(\n",
-      "    (0): ConvTranspose2d(400, 1000, kernel_size=(6, 6), stride=(1, 1))\n",
-      "    (1): BatchNorm2d(1000, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (0): ConvTranspose2d(200, 800, kernel_size=(6, 6), stride=(1, 1))\n",
+      "    (1): BatchNorm2d(800, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
-      "    (3): ConvTranspose2d(1000, 800, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (4): BatchNorm2d(800, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (3): ConvTranspose2d(800, 400, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (4): BatchNorm2d(400, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (5): ReLU()\n",
-      "    (6): ConvTranspose2d(800, 600, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (7): BatchNorm2d(600, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (6): ConvTranspose2d(400, 200, kernel_size=(6, 6), stride=(1, 1), padding=(1, 1))\n",
+      "    (7): BatchNorm2d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (8): ReLU()\n",
-      "    (9): ConvTranspose2d(600, 400, kernel_size=(6, 6), stride=(2, 2), padding=(1, 1))\n",
-      "    (10): BatchNorm2d(400, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (9): ConvTranspose2d(200, 100, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (10): BatchNorm2d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (11): ReLU()\n",
-      "    (12): ConvTranspose2d(400, 200, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (13): BatchNorm2d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "    (14): ReLU()\n",
-      "    (15): ConvTranspose2d(200, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
-      "    (16): Tanh()\n",
+      "    (12): ConvTranspose2d(100, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "    (13): Tanh()\n",
       "  )\n",
       ")\n",
-      "Iteration :  20\n",
-      "INFO: Iter 20: gen_loss=9.649e+00, dis_loss=5.266e-01\n",
-      "Iteration :  21\n",
-      "Iteration :  22\n",
-      "Iteration :  23\n",
-      "Iteration :  24\n",
-      "Iteration :  25\n",
-      "Iteration :  26\n",
-      "Iteration :  27\n",
-      "Iteration :  28\n",
-      "Iteration :  29\n",
-      "Iteration :  30\n",
-      "INFO: Iter 30: gen_loss=6.851e+00, dis_loss=1.911e-01\n",
-      "Iteration :  31\n",
-      "Iteration :  32\n",
-      "Iteration :  33\n",
-      "Iteration :  34\n",
-      "Iteration :  35\n",
-      "Iteration :  36\n",
-      "Iteration :  37\n",
-      "Iteration :  38\n",
-      "Iteration :  39\n",
-      "Iteration :  40\n",
-      "INFO: Iter 40: gen_loss=6.467e+00, dis_loss=3.223e-01\n"
+      "Generator output\n",
+      "torch.Size([16, 3, 60, 60])\n",
+      "Traceback (most recent call last):\n",
+      "  File \"ganSynthesisImage_100.py\", line 249, in <module>\n",
+      "    dis_output_true_v = net_discr(batch_v)\n",
+      "  File \"/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py\", line 489, in __call__\n",
+      "    result = self.forward(*input, **kwargs)\n",
+      "  File \"ganSynthesisImage_100.py\", line 75, in forward\n",
+      "    conv_out = self.conv_pipe(x)\n",
+      "  File \"/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py\", line 489, in __call__\n",
+      "    result = self.forward(*input, **kwargs)\n",
+      "  File \"/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py\", line 92, in forward\n",
+      "    input = module(input)\n",
+      "  File \"/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py\", line 489, in __call__\n",
+      "    result = self.forward(*input, **kwargs)\n",
+      "  File \"/opt/conda/lib/python3.6/site-packages/torch/nn/modules/conv.py\", line 320, in forward\n",
+      "    self.padding, self.dilation, self.groups)\n",
+      "RuntimeError: std::exception\n"
      ]
     }
    ],
    "source": [
-    "!python ganSynthesisImage_200.py --load test_model"
+    "!python ganSynthesisImage_100.py --folder gan_synthesis"
    ]
   },
   {

+ 3 - 2
tensorboard.ipynb

@@ -11,14 +11,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "TensorBoard 1.12.2 at http://7447cb2679c8:6006 (Press CTRL+C to quit)\n"
+      "TensorBoard 1.12.2 at http://309553e0ec4f:6006 (Press CTRL+C to quit)\n",
+      "^C\n"
      ]
     }
    ],