mirror of https://github.com/kritiksoman/GIMP-ML
GIMP 3 development progress
DESKTOP-F04AGRR\Kritik Soman
3 years ago
parent
7898514a76
commit
9766f999a3
@ -1,8 +0,0 @@
|
||||
.DS_Store
|
||||
*~
|
||||
*$py.class
|
||||
*.so
|
||||
*.pyc
|
||||
gimpenv/
|
||||
weights/
|
||||
output/
|
||||
@ -0,0 +1,65 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="CsvFileAttributes">
|
||||
<option name="attributeMap">
|
||||
<map>
|
||||
<entry key="C:\Users\Kritik Soman\AppData\Roaming\JetBrains\PyCharmCE2020.2\scratches\backtranslate.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="C:\Users\Kritik Soman\AppData\Roaming\JetBrains\PyCharmCE2020.2\scratches\scratch.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="D:\PycharmProjects\gimp\plug-ins\python\foggify.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="D:\PycharmProjects\gimp\plug-ins\python\histogram-export.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="\gimpml\tools\complete_install.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value=":" />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="\gimpml\tools\model_info.csv">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="\tmp.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value=":" />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
<entry key="\tmponedrive.py">
|
||||
<value>
|
||||
<Attribute>
|
||||
<option name="separator" value="," />
|
||||
</Attribute>
|
||||
</value>
|
||||
</entry>
|
||||
</map>
|
||||
</option>
|
||||
</component>
|
||||
</project>
|
||||
@ -1,4 +1,4 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="Python 2.7 (gimpenv)" project-jdk-type="Python SDK" />
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8 (gimpenv3)" project-jdk-type="Python SDK" />
|
||||
</project>
|
||||
@ -0,0 +1,10 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/../GIMP-ML/.idea/GIMP-ML.iml" filepath="$PROJECT_DIR$/../GIMP-ML/.idea/GIMP-ML.iml" />
|
||||
<module fileurl="file://$PROJECT_DIR$/.idea/GIMP3-ML-pip.iml" filepath="$PROJECT_DIR$/.idea/GIMP3-ML-pip.iml" />
|
||||
<module fileurl="file://$PROJECT_DIR$/../gimp/.idea/gimp.iml" filepath="$PROJECT_DIR$/../gimp/.idea/gimp.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
@ -0,0 +1,7 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="$PROJECT_DIR$/../GIMP-ML" vcs="Git" />
|
||||
<mapping directory="$PROJECT_DIR$/../gimp" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
@ -1,34 +0,0 @@
|
||||
## GIMP-ML directory structure
|
||||
|
||||
```plaintext
|
||||
|-- GIMP-ML
|
||||
| |-- gimp-plugins
|
||||
| |-- gimpenv
|
||||
| |-- weights
|
||||
| |-- .........
|
||||
| |-- <plugin files and folders>
|
||||
| |-- .........
|
||||
| |-- installGimpML.sh
|
||||
|-- README.md
|
||||
|-- INSTALLATION.md
|
||||
```
|
||||
## Installation Methods
|
||||
Clone this repository: ```git clone https://github.com/kritiksoman/GIMP-ML.git``` <br>
|
||||
### Manual Setup
|
||||
[1] Install python 2.7 using ```sudo apt install python2-minimal``` or ```sudo apt install python-minimal``` if not already present. <br>
|
||||
[2] Download and install pip ```wget https://bootstrap.pypa.io/get-pip.py ```, followed by ```python get-pip.py```, if not already present. <br>
|
||||
[3] Install virtualenv package with ```python -m pip install --user virtualenv```. <br>
|
||||
[4] Create a virtual environment ```gimpenv``` in the ```gimp-plugins``` folder using ```python -m virtualenv gimpenv```.<br>
|
||||
[5] Activate the environment and install ```torchvision, "opencv-python<=4.3", numpy, future, torch, scipy, typing, enum, pretrainedmodels, requests``` using pip. <br>
|
||||
[6] Open GIMP and go to Preferences -> Folders -> Plug-ins, add the folder gimp-plugins and close GIMP. <br>
|
||||
[7] Download the weights folder from [link](https://drive.google.com/drive/folders/10IiBO4fuMiGQ-spBStnObbk9R-pGp6u8?usp=sharing) and move it inside ```gimp-plugins``` folder. <br>
|
||||
[8] Allow the python scripts to be executable using ```chmod +x *.py``` in the ```gimp-plugins``` folder.<br>
|
||||
[9] Run GIMP. <br>
|
||||
Note: See [```installGimpML.sh```](https://github.com/kritiksoman/GIMP-ML/blob/master/gimp-plugins/installGimpML.sh) if getting stuck.
|
||||
|
||||
|
||||
### Using shell script and update tool
|
||||
[1] Open terminal, go to GIMP-ML/gimp-plugins and run : ```bash installGimpML.sh```<br>
|
||||
[2] Open GIMP and go to Preferences -> Folders -> Plug-ins, add the folder gimp-plugins and restart GIMP. <br>
|
||||
[3] Go to Layer->GIMP-ML->update, click on ok with "update weights" set to yes and restart GIMP. (Weights ~ 1.5GB will be downloaded)<br>
|
||||
|
||||
@ -0,0 +1,2 @@
|
||||
include gimpml/plugins/colorpalette/color_palette.png
|
||||
include gimpml/tools/model_info.csv
|
||||
@ -0,0 +1,17 @@
|
||||
# from .kmeans import get_kmeans as kmeans
|
||||
# from .deblur import getdeblur as deblur
|
||||
# from .deepcolor import get_deepcolor as deepcolor
|
||||
# from .deepdehaze import get_dehaze as dehaze
|
||||
# from .deepdenoise import get_denoise as denoise
|
||||
# from .deepmatting import get_newalpha as matting
|
||||
# from .enlighten import get_enlighten as enlighten
|
||||
# from .facegen import get_newface as newface
|
||||
# from .faceparse import get_face as parseface
|
||||
# from .interpolateframes import get_inter as interpolateframe
|
||||
from .tools.monodepth import get_mono_depth as depth
|
||||
from .tools.complete_install import setup_python_weights
|
||||
# from .semseg import get_sem_seg as semseg
|
||||
# from .super_resolution import get_super as super
|
||||
# from .inpainting import get_inpaint as inpaint
|
||||
# from .syncWeights import sync as sync
|
||||
|
||||
@ -0,0 +1,44 @@
|
||||
import pickle
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
|
||||
plugin_loc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
base_loc = os.path.expanduser("~") + '/GIMP-ML/'
|
||||
# base_loc = "D:/PycharmProjects/"
|
||||
sys.path.extend([plugin_loc + 'MiDaS'])
|
||||
# data_path = "D:/PycharmProjects/GIMP3-ML-pip/gimpml/"
|
||||
|
||||
from mono_run import run_depth
|
||||
from monodepth_net import MonoDepthNet
|
||||
import MiDaS_utils as MiDaS_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
|
||||
def get_mono_depth(input_image, cFlag = False):
|
||||
image = input_image / 255.0
|
||||
out = run_depth(image, base_loc + 'weights/MiDaS/model.pt', MonoDepthNet, MiDaS_utils, target_w=640, f=cFlag)
|
||||
out = np.repeat(out[:, :, np.newaxis], 3, axis=2)
|
||||
d1, d2 = input_image.shape[:2]
|
||||
out = cv2.resize(out, (d2, d1))
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# # This will run when script is run as sub-process
|
||||
# dbfile = open(data_path + "data_input", 'rb')
|
||||
# data_input = pickle.load(dbfile)
|
||||
# dbfile.close()
|
||||
# # print(data)
|
||||
# data_output = {'args_input': {'processed': 1}, 'image_output': get_mono_depth(data_input['image'])}
|
||||
#
|
||||
# dbfile = open(data_path + "data_output", 'ab')
|
||||
# pickle.dump(data_output, dbfile) # source, destination
|
||||
# dbfile.close()
|
||||
|
||||
image = cv2.imread(os.path.join(base_loc, "cache.png"))[:, :, ::-1]
|
||||
output = get_mono_depth(image)
|
||||
cv2.imwrite(os.path.join(base_loc, 'cache.png'), output[:, :, ::-1])
|
||||
@ -0,0 +1,44 @@
|
||||
import pickle
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
|
||||
plugin_loc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
base_loc = os.path.expanduser("~") + '/GIMP-ML/'
|
||||
# base_loc = "D:/PycharmProjects/"
|
||||
sys.path.extend([plugin_loc + 'MiDaS'])
|
||||
# data_path = "D:/PycharmProjects/GIMP3-ML-pip/gimpml/"
|
||||
|
||||
from mono_run import run_depth
|
||||
from monodepth_net import MonoDepthNet
|
||||
import MiDaS_utils as MiDaS_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
|
||||
def get_mono_depth(input_image, cFlag = False):
|
||||
image = input_image / 255.0
|
||||
out = run_depth(image, base_loc + 'weights/MiDaS/model.pt', MonoDepthNet, MiDaS_utils, target_w=640, f=cFlag)
|
||||
out = np.repeat(out[:, :, np.newaxis], 3, axis=2)
|
||||
d1, d2 = input_image.shape[:2]
|
||||
out = cv2.resize(out, (d2, d1))
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# # This will run when script is run as sub-process
|
||||
# dbfile = open(data_path + "data_input", 'rb')
|
||||
# data_input = pickle.load(dbfile)
|
||||
# dbfile.close()
|
||||
# # print(data)
|
||||
# data_output = {'args_input': {'processed': 1}, 'image_output': get_mono_depth(data_input['image'])}
|
||||
#
|
||||
# dbfile = open(data_path + "data_output", 'ab')
|
||||
# pickle.dump(data_output, dbfile) # source, destination
|
||||
# dbfile.close()
|
||||
|
||||
image = cv2.imread(os.path.join(base_loc, "cache.png"))[:, :, ::-1]
|
||||
output = get_mono_depth(image)
|
||||
cv2.imwrite(os.path.join(base_loc, 'cache.png'), output[:, :, ::-1])
|
||||
@ -0,0 +1,238 @@
|
||||
#!/usr/bin/env python3
|
||||
#coding: utf-8
|
||||
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU General Public License as published by
|
||||
# the Free Software Foundation; either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""
|
||||
Extracts the monocular depth of the current layer.
|
||||
"""
|
||||
import pickle
|
||||
import csv
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
import gi
|
||||
gi.require_version('Gimp', '3.0')
|
||||
from gi.repository import Gimp
|
||||
gi.require_version('GimpUi', '3.0')
|
||||
from gi.repository import GimpUi
|
||||
from gi.repository import GObject
|
||||
from gi.repository import GLib
|
||||
from gi.repository import Gio
|
||||
gi.require_version('Gtk', '3.0')
|
||||
from gi.repository import Gtk
|
||||
|
||||
import gettext
|
||||
_ = gettext.gettext
|
||||
def N_(message): return message
|
||||
|
||||
import subprocess
|
||||
import pickle
|
||||
import os
|
||||
|
||||
def monodepth(procedure, image, drawable, force_cpu, progress_bar):
|
||||
|
||||
#
|
||||
# install_location = os.path.join(os.path.expanduser("~"), "GIMP-ML")
|
||||
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools")
|
||||
with open(os.path.join(config_path, 'gimp_ml_config.pkl'), 'rb') as file:
|
||||
data_output = pickle.load(file)
|
||||
weight_path = data_output["weight_path"]
|
||||
python_path = data_output["python_path"]
|
||||
# python_path = r"D:/PycharmProjects/gimpenv3/Scripts/python.exe"
|
||||
plugin_path = os.path.join(config_path, 'monodepth.py')
|
||||
|
||||
# plugin_path = r"D:/PycharmProjects/GIMP3-ML-pip/gimpml/tools/monodepth.py"
|
||||
|
||||
Gimp.context_push()
|
||||
image.undo_group_start()
|
||||
|
||||
interlace, compression = 0, 2
|
||||
Gimp.get_pdb().run_procedure('file-png-save', [
|
||||
GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
|
||||
GObject.Value(Gimp.Image, image),
|
||||
GObject.Value(GObject.TYPE_INT, 1),
|
||||
GObject.Value(Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, [drawable], 1)),
|
||||
GObject.Value(Gio.File, Gio.File.new_for_path(os.path.join(weight_path, 'cache.png'))),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, interlace),
|
||||
GObject.Value(GObject.TYPE_INT, compression),
|
||||
# write all PNG chunks except oFFs(ets)
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, False),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
])
|
||||
|
||||
subprocess.call([python_path, plugin_path])
|
||||
|
||||
result = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(os.path.join(weight_path, 'cache.png')))
|
||||
result_layer = result.get_active_layer()
|
||||
copy = Gimp.Layer.new_from_drawable(result_layer, image)
|
||||
copy.set_name("Mono Depth")
|
||||
copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)#DIFFERENCE_LEGACY
|
||||
image.insert_layer(copy, None, -1)
|
||||
|
||||
image.undo_group_end()
|
||||
Gimp.context_pop()
|
||||
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
|
||||
|
||||
|
||||
|
||||
def run(procedure, run_mode, image, layer, args, data):
|
||||
# gio_file = args.index(0)
|
||||
# bucket_size = args.index(0)
|
||||
force_cpu = args.index(1)
|
||||
# output_format = args.index(2)
|
||||
|
||||
progress_bar = None
|
||||
config = None
|
||||
|
||||
if run_mode == Gimp.RunMode.INTERACTIVE:
|
||||
|
||||
config = procedure.create_config()
|
||||
|
||||
# Set properties from arguments. These properties will be changed by the UI.
|
||||
#config.set_property("file", gio_file)
|
||||
#config.set_property("bucket_size", bucket_size)
|
||||
config.set_property("force_cpu", force_cpu)
|
||||
#config.set_property("output_format", output_format)
|
||||
config.begin_run(image, run_mode, args)
|
||||
|
||||
GimpUi.init("monodepth.py")
|
||||
use_header_bar = Gtk.Settings.get_default().get_property("gtk-dialogs-use-header")
|
||||
dialog = GimpUi.Dialog(use_header_bar=use_header_bar,
|
||||
title=_("Mono Depth..."))
|
||||
dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
|
||||
dialog.add_button("_OK", Gtk.ResponseType.OK)
|
||||
|
||||
vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL,
|
||||
homogeneous=False, spacing=10)
|
||||
dialog.get_content_area().add(vbox)
|
||||
vbox.show()
|
||||
|
||||
# Create grid to set all the properties inside.
|
||||
grid = Gtk.Grid()
|
||||
grid.set_column_homogeneous(False)
|
||||
grid.set_border_width(10)
|
||||
grid.set_column_spacing(10)
|
||||
grid.set_row_spacing(10)
|
||||
vbox.add(grid)
|
||||
grid.show()
|
||||
|
||||
# # Bucket size parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Bucket Size"))
|
||||
# grid.attach(label, 0, 1, 1, 1)
|
||||
# label.show()
|
||||
# spin = GimpUi.prop_spin_button_new(config, "bucket_size", step_increment=0.001, page_increment=0.1, digits=3)
|
||||
# grid.attach(spin, 1, 1, 1, 1)
|
||||
# spin.show()
|
||||
|
||||
# Force CPU parameter
|
||||
spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
|
||||
spin.set_tooltip_text(_("If checked, CPU is used for model inference."
|
||||
" Otherwise, GPU will be used if available."))
|
||||
grid.attach(spin, 1, 2, 1, 1)
|
||||
spin.show()
|
||||
|
||||
# # Output format parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Output Format"))
|
||||
# grid.attach(label, 0, 3, 1, 1)
|
||||
# label.show()
|
||||
# combo = GimpUi.prop_string_combo_box_new(config, "output_format", output_format_enum.get_tree_model(), 0, 1)
|
||||
# grid.attach(combo, 1, 3, 1, 1)
|
||||
# combo.show()
|
||||
|
||||
progress_bar = Gtk.ProgressBar()
|
||||
vbox.add(progress_bar)
|
||||
progress_bar.show()
|
||||
|
||||
dialog.show()
|
||||
if dialog.run() != Gtk.ResponseType.OK:
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.CANCEL,
|
||||
GLib.Error())
|
||||
|
||||
result = monodepth(procedure, image, layer, force_cpu, progress_bar)
|
||||
|
||||
# If the execution was successful, save parameters so they will be restored next time we show dialog.
|
||||
if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
|
||||
config.end_run(Gimp.PDBStatusType.SUCCESS)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class MonoDepth(Gimp.PlugIn):
|
||||
|
||||
## Parameters ##
|
||||
__gproperties__ = {
|
||||
# "filename": (str,
|
||||
# # TODO: I wanted this property to be a path (and not just str) , so I could use
|
||||
# # prop_file_chooser_button_new to open a file dialog. However, it fails without an error message.
|
||||
# # Gimp.ConfigPath,
|
||||
# _("Histogram _File"),
|
||||
# _("Histogram _File"),
|
||||
# "monodepth.csv",
|
||||
# # Gimp.ConfigPathType.FILE,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "file": (Gio.File,
|
||||
# _("Histogram _File"),
|
||||
# "Histogram export file",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "bucket_size": (float,
|
||||
# _("_Bucket Size"),
|
||||
# "Bucket Size",
|
||||
# 0.001, 1.0, 0.01,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
"force_cpu": (bool,
|
||||
_("Force _CPU"),
|
||||
"Force CPU",
|
||||
False,
|
||||
GObject.ParamFlags.READWRITE),
|
||||
# "output_format": (str,
|
||||
# _("Output format"),
|
||||
# "Output format: 'pixel count', 'normalized', 'percent'",
|
||||
# "pixel count",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
}
|
||||
|
||||
## GimpPlugIn virtual methods ##
|
||||
def do_query_procedures(self):
|
||||
self.set_translation_domain("gimp30-python",
|
||||
Gio.file_new_for_path(Gimp.locale_directory()))
|
||||
return ['monodepth']
|
||||
|
||||
def do_create_procedure(self, name):
|
||||
procedure = None
|
||||
if name == 'monodepth':
|
||||
procedure = Gimp.ImageProcedure.new(self, name, Gimp.PDBProcType.PLUGIN, run, None)
|
||||
procedure.set_image_types("*")
|
||||
procedure.set_documentation (
|
||||
N_("Extracts the monocular depth of the current layer."),
|
||||
globals()["__doc__"], # This includes the docstring, on the top of the file
|
||||
name)
|
||||
procedure.set_menu_label(N_("_Mono Depth..."))
|
||||
procedure.set_attribution("Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2021")
|
||||
procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
|
||||
|
||||
# procedure.add_argument_from_property(self, "file")
|
||||
# procedure.add_argument_from_property(self, "bucket_size")
|
||||
procedure.add_argument_from_property(self, "force_cpu")
|
||||
# procedure.add_argument_from_property(self, "output_format")
|
||||
|
||||
return procedure
|
||||
|
||||
|
||||
Gimp.main(MonoDepth.__gtype__, sys.argv)
|
||||
@ -0,0 +1,29 @@
|
||||
import cv2
|
||||
import pickle
|
||||
import subprocess
|
||||
|
||||
|
||||
def run_plugin(data_input, plugin_path):
|
||||
dbfile = open(data_path + "data_input", 'ab')
|
||||
pickle.dump(data_input, dbfile) # source, destination
|
||||
dbfile.close()
|
||||
subprocess.call([python_path, plugin_path])
|
||||
dbfile = open(data_path + "data_output", 'rb')
|
||||
data_output = pickle.load(dbfile)
|
||||
dbfile.close()
|
||||
return data_output
|
||||
|
||||
|
||||
image = cv2.imread('/Users/kritiksoman/Documents/GitHub/GIMP-ML-pip/sampleinput/inpaint.png')[:, :, ::-1]
|
||||
|
||||
python_path = "/Users/kritiksoman/GIMP-ML/gimpenv3/bin/python"
|
||||
plugin_path = "/Users/kritiksoman/Documents/GitHub/GIMP3-ML-pip/gimpml/monodepth.py"
|
||||
data_path = "/Users/kritiksoman/GIMP-ML/"
|
||||
|
||||
# save data and call plugin
|
||||
data_input = {'image': image, 'args_input': {'force_cpu': 0}}
|
||||
|
||||
data_output = run_plugin(data_input, plugin_path)
|
||||
# print(data_output)
|
||||
cv2.imwrite("output.png", data_output['image_output'])
|
||||
|
||||
@ -0,0 +1,70 @@
|
||||
"""
|
||||
Script will download weights and create gimp_ml_config.pkl, and print path to be added to GIMP
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import pickle
|
||||
import csv
|
||||
import hashlib
|
||||
import gdown
|
||||
|
||||
|
||||
def setup_python_weights(install_location=None):
|
||||
if not install_location:
|
||||
install_location = os.path.join(os.path.expanduser("~"), "GIMP-ML")
|
||||
if not os.path.isdir(install_location):
|
||||
os.mkdir(install_location)
|
||||
python_string = "python"
|
||||
if os.name == 'nt': # windows
|
||||
python_string += ".exe"
|
||||
python_path = os.path.join(os.path.dirname(sys.executable), python_string)
|
||||
with open(os.path.join(install_location, 'gimp_ml_config.pkl'), 'ab') as file:
|
||||
pickle.dump(python_path, file)
|
||||
# print(r"\n\n*******************", python_path)
|
||||
|
||||
weight_path = os.path.join(install_location, "weights")
|
||||
if not os.path.isdir(weight_path):
|
||||
os.mkdir(weight_path)
|
||||
|
||||
if os.name == 'nt': # windows
|
||||
print("Automatic downloading of weights not supported on Windows.")
|
||||
print("Please downloads weights folder from: \n"
|
||||
"https://drive.google.com/drive/folders/10IiBO4fuMiGQ-spBStnObbk9R-pGp6u8?usp=sharing")
|
||||
print("and place in: " + weight_path)
|
||||
else: # linux
|
||||
with open('model_info.csv') as csv_file:
|
||||
csv_reader = csv.reader(csv_file, delimiter=',')
|
||||
headings = next(csv_reader)
|
||||
line_count = 0
|
||||
for row in csv_reader:
|
||||
model = row[0]
|
||||
file_id = row[1]
|
||||
fileSize = float(row[2]) # in MB
|
||||
mFName = row[3]
|
||||
md5sum = row[4]
|
||||
if not os.path.isdir(os.path.join(weight_path, model)):
|
||||
os.mkdir(os.path.join(weight_path, model))
|
||||
destination = os.path.join(os.path.join(weight_path, model), mFName)
|
||||
if os.path.isfile(destination):
|
||||
md5_hash = hashlib.md5()
|
||||
a_file = open(destination, "rb")
|
||||
content = a_file.read()
|
||||
md5_hash.update(content)
|
||||
digest = md5_hash.hexdigest()
|
||||
if not os.path.isfile(destination) or (digest and digest != md5sum):
|
||||
try:
|
||||
gimp.progress_init("Downloading " + model + "(~" + str(fileSize) + "MB)...")
|
||||
except:
|
||||
print("\nDownloading " + model + "(~" + str(fileSize) + "MB)...")
|
||||
url = 'https://drive.google.com/uc?id={0}'.format(file_id)
|
||||
gdown.cached_download(url, destination, md5=md5sum)
|
||||
plugin_loc = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(os.path.join(plugin_loc, 'gimp_ml_config.pkl'), 'wb') as file:
|
||||
pickle.dump({"python_path": python_path, "weight_path": weight_path}, file)
|
||||
|
||||
print("Please add this path to Preferences-->Plug-ins : ",
|
||||
os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "plugins"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup_python_weights()
|
||||
@ -0,0 +1,50 @@
|
||||
import pickle
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
|
||||
|
||||
plugin_loc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
# base_loc = "D:/PycharmProjects/"
|
||||
sys.path.extend([plugin_loc + 'MiDaS'])
|
||||
# data_path = "D:/PycharmProjects/GIMP3-ML-pip/gimpml/"
|
||||
|
||||
from mono_run import run_depth
|
||||
from monodepth_net import MonoDepthNet
|
||||
import MiDaS_utils as MiDaS_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
|
||||
def get_mono_depth(input_image, cFlag = False):
|
||||
image = input_image / 255.0
|
||||
out = run_depth(image, os.path.join(weight_path, 'MiDaS', 'model.pt'), MonoDepthNet, MiDaS_utils, target_w=640, f=cFlag)
|
||||
out = np.repeat(out[:, :, np.newaxis], 3, axis=2)
|
||||
d1, d2 = input_image.shape[:2]
|
||||
out = cv2.resize(out, (d2, d1))
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# # This will run when script is run as sub-process
|
||||
# dbfile = open(data_path + "data_input", 'rb')
|
||||
# data_input = pickle.load(dbfile)
|
||||
# dbfile.close()
|
||||
# # print(data)
|
||||
# data_output = {'args_input': {'processed': 1}, 'image_output': get_mono_depth(data_input['image'])}
|
||||
#
|
||||
# dbfile = open(data_path + "data_output", 'ab')
|
||||
# pickle.dump(data_output, dbfile) # source, destination
|
||||
# dbfile.close()
|
||||
config_path = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(os.path.join(config_path, 'gimp_ml_config.pkl'), 'rb') as file:
|
||||
# with open('gimp_ml_config.pkl', 'rb') as file:
|
||||
data_output = pickle.load(file)
|
||||
# base_loc = os.path.expanduser("~") + '/GIMP-ML/'
|
||||
weight_path = data_output["weight_path"]
|
||||
image = cv2.imread(os.path.join(weight_path, "cache.png"))[:, :, ::-1]
|
||||
output = get_mono_depth(image)
|
||||
cv2.imwrite(os.path.join(weight_path, 'cache.png'), output[:, :, ::-1])
|
||||
Binary file not shown.
Binary file not shown.
@ -1,100 +0,0 @@
|
||||
## Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import os.path
|
||||
from data.base_dataset import BaseDataset, get_params, get_transform, normalize
|
||||
from data.image_folder import make_dataset, make_dataset_test
|
||||
from PIL import Image
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class AlignedDataset(BaseDataset):
|
||||
def initialize(self, opt):
|
||||
self.opt = opt
|
||||
self.root = opt.dataroot
|
||||
|
||||
### input A (label maps)
|
||||
if opt.isTrain or opt.use_encoded_image:
|
||||
dir_A = '_A' if self.opt.label_nc == 0 else '_label'
|
||||
self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
|
||||
self.A_paths = sorted(make_dataset(self.dir_A))
|
||||
self.AR_paths = make_dataset(self.dir_A)
|
||||
|
||||
### input A inter 1 (label maps)
|
||||
if opt.isTrain or opt.use_encoded_image:
|
||||
dir_A_inter_1 = '_label_inter_1'
|
||||
self.dir_A_inter_1 = os.path.join(opt.dataroot, opt.phase + dir_A_inter_1)
|
||||
self.A_paths_inter_1 = sorted(make_dataset(self.dir_A_inter_1))
|
||||
|
||||
### input A inter 2 (label maps)
|
||||
if opt.isTrain or opt.use_encoded_image:
|
||||
dir_A_inter_2 = '_label_inter_2'
|
||||
self.dir_A_inter_2 = os.path.join(opt.dataroot, opt.phase + dir_A_inter_2)
|
||||
self.A_paths_inter_2 = sorted(make_dataset(self.dir_A_inter_2))
|
||||
|
||||
### input A test (label maps)
|
||||
if not (opt.isTrain or opt.use_encoded_image):
|
||||
dir_A = '_A' if self.opt.label_nc == 0 else '_label'
|
||||
self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
|
||||
self.A_paths = sorted(make_dataset_test(self.dir_A))
|
||||
dir_AR = '_AR' if self.opt.label_nc == 0 else '_labelref'
|
||||
self.dir_AR = os.path.join(opt.dataroot, opt.phase + dir_AR)
|
||||
self.AR_paths = sorted(make_dataset_test(self.dir_AR))
|
||||
|
||||
### input B (real images)
|
||||
dir_B = '_B' if self.opt.label_nc == 0 else '_img'
|
||||
self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B)
|
||||
self.B_paths = sorted(make_dataset(self.dir_B))
|
||||
self.BR_paths = sorted(make_dataset(self.dir_B))
|
||||
|
||||
self.dataset_size = len(self.A_paths)
|
||||
|
||||
def __getitem__(self, index):
|
||||
### input A (label maps)
|
||||
A_path = self.A_paths[index]
|
||||
AR_path = self.AR_paths[index]
|
||||
A = Image.open(A_path)
|
||||
AR = Image.open(AR_path)
|
||||
|
||||
if self.opt.isTrain:
|
||||
A_path_inter_1 = self.A_paths_inter_1[index]
|
||||
A_path_inter_2 = self.A_paths_inter_2[index]
|
||||
A_inter_1 = Image.open(A_path_inter_1)
|
||||
A_inter_2 = Image.open(A_path_inter_2)
|
||||
|
||||
params = get_params(self.opt, A.size)
|
||||
if self.opt.label_nc == 0:
|
||||
transform_A = get_transform(self.opt, params)
|
||||
A_tensor = transform_A(A.convert('RGB'))
|
||||
if self.opt.isTrain:
|
||||
A_inter_1_tensor = transform_A(A_inter_1.convert('RGB'))
|
||||
A_inter_2_tensor = transform_A(A_inter_2.convert('RGB'))
|
||||
AR_tensor = transform_A(AR.convert('RGB'))
|
||||
else:
|
||||
transform_A = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
|
||||
A_tensor = transform_A(A) * 255.0
|
||||
if self.opt.isTrain:
|
||||
A_inter_1_tensor = transform_A(A_inter_1) * 255.0
|
||||
A_inter_2_tensor = transform_A(A_inter_2) * 255.0
|
||||
AR_tensor = transform_A(AR) * 255.0
|
||||
B_tensor = inst_tensor = feat_tensor = 0
|
||||
### input B (real images)
|
||||
B_path = self.B_paths[index]
|
||||
BR_path = self.BR_paths[index]
|
||||
B = Image.open(B_path).convert('RGB')
|
||||
BR = Image.open(BR_path).convert('RGB')
|
||||
transform_B = get_transform(self.opt, params)
|
||||
B_tensor = transform_B(B)
|
||||
BR_tensor = transform_B(BR)
|
||||
|
||||
if self.opt.isTrain:
|
||||
input_dict = {'inter_label_1': A_inter_1_tensor, 'label': A_tensor, 'inter_label_2': A_inter_2_tensor, 'label_ref': AR_tensor, 'image': B_tensor, 'image_ref': BR_tensor, 'path': A_path, 'path_ref': AR_path}
|
||||
else:
|
||||
input_dict = {'label': A_tensor, 'label_ref': AR_tensor, 'image': B_tensor, 'image_ref': BR_tensor, 'path': A_path, 'path_ref': AR_path}
|
||||
|
||||
return input_dict
|
||||
|
||||
def __len__(self):
|
||||
return len(self.A_paths) // self.opt.batchSize * self.opt.batchSize
|
||||
|
||||
def name(self):
|
||||
return 'AlignedDataset'
|
||||
@ -1,14 +0,0 @@
|
||||
|
||||
class BaseDataLoader():
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def initialize(self, opt):
|
||||
self.opt = opt
|
||||
pass
|
||||
|
||||
def load_data():
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@ -1,97 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
import torchvision.transforms as transforms
|
||||
import numpy as np
|
||||
import random
|
||||
|
||||
class BaseDataset(data.Dataset):
|
||||
def __init__(self):
|
||||
super(BaseDataset, self).__init__()
|
||||
|
||||
def name(self):
|
||||
return 'BaseDataset'
|
||||
|
||||
def initialize(self, opt):
|
||||
pass
|
||||
|
||||
def get_params(opt, size):
|
||||
w, h = size
|
||||
new_h = h
|
||||
new_w = w
|
||||
if opt.resize_or_crop == 'resize_and_crop':
|
||||
new_h = new_w = opt.loadSize
|
||||
elif opt.resize_or_crop == 'scale_width_and_crop':
|
||||
new_w = opt.loadSize
|
||||
new_h = opt.loadSize * h // w
|
||||
|
||||
x = random.randint(0, np.maximum(0, new_w - opt.fineSize))
|
||||
y = random.randint(0, np.maximum(0, new_h - opt.fineSize))
|
||||
|
||||
#flip = random.random() > 0.5
|
||||
flip = 0
|
||||
return {'crop_pos': (x, y), 'flip': flip}
|
||||
|
||||
def get_transform(opt, params, method=Image.BICUBIC, normalize=True, normalize_mask=False):
|
||||
transform_list = []
|
||||
if 'resize' in opt.resize_or_crop:
|
||||
osize = [opt.loadSize, opt.loadSize]
|
||||
transform_list.append(transforms.Scale(osize, method))
|
||||
elif 'scale_width' in opt.resize_or_crop:
|
||||
transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.loadSize, method)))
|
||||
|
||||
if 'crop' in opt.resize_or_crop:
|
||||
transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.fineSize)))
|
||||
|
||||
if opt.resize_or_crop == 'none':
|
||||
base = float(2 ** opt.n_downsample_global)
|
||||
if opt.netG == 'local':
|
||||
base *= (2 ** opt.n_local_enhancers)
|
||||
transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
|
||||
|
||||
if opt.isTrain and not opt.no_flip:
|
||||
transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
|
||||
|
||||
transform_list += [transforms.ToTensor()]
|
||||
|
||||
if normalize:
|
||||
transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
|
||||
(0.5, 0.5, 0.5))]
|
||||
if normalize_mask:
|
||||
transform_list += [transforms.Normalize((0, 0, 0),
|
||||
(1 / 255., 1 / 255., 1 / 255.))]
|
||||
|
||||
return transforms.Compose(transform_list)
|
||||
|
||||
def normalize():
|
||||
return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
|
||||
def __make_power_2(img, base, method=Image.BICUBIC):
|
||||
ow, oh = img.size
|
||||
h = int(round(oh / base) * base)
|
||||
w = int(round(ow / base) * base)
|
||||
if (h == oh) and (w == ow):
|
||||
return img
|
||||
return img.resize((w, h), method)
|
||||
|
||||
def __scale_width(img, target_width, method=Image.BICUBIC):
|
||||
ow, oh = img.size
|
||||
if (ow == target_width):
|
||||
return img
|
||||
w = target_width
|
||||
h = int(target_width * oh / ow)
|
||||
return img.resize((w, h), method)
|
||||
|
||||
def __crop(img, pos, size):
|
||||
ow, oh = img.size
|
||||
x1, y1 = pos
|
||||
tw = th = size
|
||||
if (ow > tw or oh > th):
|
||||
return img.crop((x1, y1, x1 + tw, y1 + th))
|
||||
return img
|
||||
|
||||
def __flip(img, flip):
|
||||
if flip:
|
||||
return img.transpose(Image.FLIP_LEFT_RIGHT)
|
||||
return img
|
||||
@ -1,31 +0,0 @@
|
||||
import torch.utils.data
|
||||
from data.base_data_loader import BaseDataLoader
|
||||
|
||||
|
||||
def CreateDataset(opt):
|
||||
dataset = None
|
||||
from data.aligned_dataset import AlignedDataset
|
||||
dataset = AlignedDataset()
|
||||
|
||||
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.dataloader
|
||||
|
||||
def __len__(self):
|
||||
return min(len(self.dataset), self.opt.max_dataset_size)
|
||||
@ -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
|
||||
@ -1,82 +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
|
||||
|
||||
IMG_EXTENSIONS = [
|
||||
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
||||
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
|
||||
]
|
||||
|
||||
|
||||
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
|
||||
|
||||
f = dir.split('/')[-1].split('_')[-1]
|
||||
print (dir, f)
|
||||
for i in range(len([name for name in os.listdir(dir) if os.path.isfile(os.path.join(dir, name))])):
|
||||
if f == 'label' or f == '1' or f == '2':
|
||||
img = str(i) + '.png'
|
||||
else:
|
||||
img = str(i) + '.jpg'
|
||||
path = os.path.join(dir, img)
|
||||
#print(path)
|
||||
images.append(path)
|
||||
return images
|
||||
|
||||
def make_dataset_test(dir):
|
||||
images = []
|
||||
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
||||
|
||||
f = dir.split('/')[-1].split('_')[-1]
|
||||
for i in range(len([name for name in os.listdir(dir) if os.path.isfile(os.path.join(dir, name))])):
|
||||
if f == 'label' or f == 'labelref':
|
||||
img = str(i) + '.png'
|
||||
else:
|
||||
img = str(i) + '.jpg'
|
||||
path = os.path.join(dir, img)
|
||||
#print(path)
|
||||
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)
|
||||
@ -1,94 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import os
|
||||
import torch
|
||||
import sys
|
||||
|
||||
class BaseModel(torch.nn.Module):
|
||||
def name(self):
|
||||
return 'BaseModel'
|
||||
|
||||
def initialize(self, opt):
|
||||
self.opt = opt
|
||||
self.gpu_ids = opt.gpu_ids
|
||||
self.isTrain = opt.isTrain
|
||||
self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
|
||||
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, gpu_ids):
|
||||
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)
|
||||
if len(gpu_ids) and torch.cuda.is_available():
|
||||
network.cuda()
|
||||
|
||||
# helper loading function that can be used by subclasses
|
||||
def load_network(self, network, network_label, epoch_label, save_dir=''):
|
||||
save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
|
||||
print (save_filename)
|
||||
if not save_dir:
|
||||
save_dir = self.save_dir
|
||||
save_path = os.path.join(save_dir, save_filename)
|
||||
if not os.path.isfile(save_path):
|
||||
print('%s not exists yet!' % save_path)
|
||||
if network_label == 'G':
|
||||
raise('Generator must exist!')
|
||||
else:
|
||||
#network.load_state_dict(torch.load(save_path))
|
||||
try:
|
||||
network.load_state_dict(torch.load(save_path))
|
||||
except:
|
||||
pretrained_dict = torch.load(save_path)
|
||||
model_dict = network.state_dict()
|
||||
try:
|
||||
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
|
||||
network.load_state_dict(pretrained_dict)
|
||||
if self.opt.verbose:
|
||||
print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
|
||||
except:
|
||||
print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
|
||||
for k, v in pretrained_dict.items():
|
||||
if v.size() == model_dict[k].size():
|
||||
model_dict[k] = v
|
||||
|
||||
if sys.version_info >= (3,0):
|
||||
not_initialized = set()
|
||||
else:
|
||||
from sets import Set
|
||||
not_initialized = Set()
|
||||
|
||||
for k, v in model_dict.items():
|
||||
if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
|
||||
not_initialized.add(k.split('.')[0])
|
||||
|
||||
print(sorted(not_initialized))
|
||||
network.load_state_dict(model_dict)
|
||||
|
||||
def update_learning_rate():
|
||||
pass
|
||||
@ -1,20 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import torch
|
||||
|
||||
def create_model(opt):
|
||||
if opt.model == 'pix2pixHD':
|
||||
from .pix2pixHD_model import Pix2PixHDModel, InferenceModel
|
||||
if opt.isTrain:
|
||||
model = Pix2PixHDModel()
|
||||
else:
|
||||
model = InferenceModel()
|
||||
|
||||
model.initialize(opt)
|
||||
if opt.verbose:
|
||||
print("model [%s] was created" % (model.name()))
|
||||
|
||||
if opt.isTrain and len(opt.gpu_ids):
|
||||
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
|
||||
|
||||
return model
|
||||
@ -1,818 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import functools
|
||||
from torch.autograd import Variable
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
###############################################################################
|
||||
# Functions
|
||||
###############################################################################
|
||||
def weights_init(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('Conv2d') != -1:
|
||||
m.weight.data.normal_(0.0, 0.02)
|
||||
elif classname.find('BatchNorm2d') != -1:
|
||||
m.weight.data.normal_(1.0, 0.02)
|
||||
m.bias.data.fill_(0)
|
||||
|
||||
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)
|
||||
else:
|
||||
raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
|
||||
return norm_layer
|
||||
|
||||
def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1,
|
||||
n_blocks_local=3, norm='instance', gpu_ids=[]):
|
||||
norm_layer = get_norm_layer(norm_type=norm)
|
||||
if netG == 'global':
|
||||
netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
|
||||
elif netG == 'local':
|
||||
netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global,
|
||||
n_local_enhancers, n_blocks_local, norm_layer)
|
||||
else:
|
||||
raise('generator not implemented!')
|
||||
print(netG)
|
||||
if len(gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
netG.cuda(gpu_ids[0])
|
||||
netG.apply(weights_init)
|
||||
return netG
|
||||
|
||||
def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]):
|
||||
norm_layer = get_norm_layer(norm_type=norm)
|
||||
netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
|
||||
print(netD)
|
||||
if len(gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
netD.cuda(gpu_ids[0])
|
||||
netD.apply(weights_init)
|
||||
return netD
|
||||
|
||||
def define_VAE(input_nc, gpu_ids=[]):
|
||||
netVAE = VAE(19, 32, 32, 1024)
|
||||
print(netVAE)
|
||||
if len(gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
netVAE.cuda(gpu_ids[0])
|
||||
return netVAE
|
||||
|
||||
def define_B(input_nc, output_nc, ngf, n_downsample_global=3, n_blocks_global=3, norm='instance', gpu_ids=[]):
|
||||
norm_layer = get_norm_layer(norm_type=norm)
|
||||
netB = BlendGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
|
||||
print(netB)
|
||||
if len(gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
netB.cuda(gpu_ids[0])
|
||||
netB.apply(weights_init)
|
||||
return netB
|
||||
|
||||
def print_network(net):
|
||||
if isinstance(net, list):
|
||||
net = net[0]
|
||||
num_params = 0
|
||||
for param in net.parameters():
|
||||
num_params += param.numel()
|
||||
print(net)
|
||||
print('Total number of parameters: %d' % num_params)
|
||||
|
||||
##############################################################################
|
||||
# Losses
|
||||
##############################################################################
|
||||
class GANLoss(nn.Module):
|
||||
def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
|
||||
tensor=torch.FloatTensor):
|
||||
super(GANLoss, self).__init__()
|
||||
self.real_label = target_real_label
|
||||
self.fake_label = target_fake_label
|
||||
self.real_label_var = None
|
||||
self.fake_label_var = None
|
||||
self.Tensor = tensor
|
||||
if use_lsgan:
|
||||
self.loss = nn.MSELoss()
|
||||
else:
|
||||
self.loss = nn.BCELoss()
|
||||
|
||||
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:
|
||||
real_tensor = self.Tensor(input.size()).fill_(self.real_label)
|
||||
self.real_label_var = Variable(real_tensor, requires_grad=False)
|
||||
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:
|
||||
fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
|
||||
self.fake_label_var = Variable(fake_tensor, requires_grad=False)
|
||||
target_tensor = self.fake_label_var
|
||||
return target_tensor
|
||||
|
||||
def __call__(self, input, target_is_real):
|
||||
if isinstance(input[0], list):
|
||||
loss = 0
|
||||
for input_i in input:
|
||||
pred = input_i[-1]
|
||||
target_tensor = self.get_target_tensor(pred, target_is_real)
|
||||
loss += self.loss(pred, target_tensor)
|
||||
return loss
|
||||
else:
|
||||
target_tensor = self.get_target_tensor(input[-1], target_is_real)
|
||||
return self.loss(input[-1], target_tensor)
|
||||
|
||||
class VGGLoss(nn.Module):
|
||||
def __init__(self, gpu_ids):
|
||||
super(VGGLoss, self).__init__()
|
||||
self.vgg = Vgg19().cuda()
|
||||
self.criterion = nn.L1Loss()
|
||||
self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
|
||||
|
||||
def forward(self, x, y):
|
||||
x_vgg, y_vgg = self.vgg(x), self.vgg(y)
|
||||
loss = 0
|
||||
for i in range(len(x_vgg)):
|
||||
loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
|
||||
return loss
|
||||
|
||||
##############################################################################
|
||||
# Generator
|
||||
##############################################################################
|
||||
class GlobalGenerator(nn.Module):
|
||||
def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
|
||||
padding_type='reflect'):
|
||||
assert(n_blocks >= 0)
|
||||
super(GlobalGenerator, self).__init__()
|
||||
activation = nn.ReLU(True)
|
||||
|
||||
model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
|
||||
### downsample
|
||||
for i in range(n_downsampling):
|
||||
mult = 2**i
|
||||
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
|
||||
norm_layer(ngf * mult * 2), activation]
|
||||
|
||||
### resnet blocks
|
||||
mult = 2**n_downsampling
|
||||
for i in range(n_blocks):
|
||||
model += [ResnetBlock(ngf * mult, norm_type='adain', padding_type=padding_type)]
|
||||
### upsample
|
||||
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),
|
||||
norm_layer(int(ngf * mult / 2)), activation]
|
||||
model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
|
||||
self.model = nn.Sequential(*model)
|
||||
|
||||
# style encoder
|
||||
self.enc_style = StyleEncoder(5, 3, 16, self.get_num_adain_params(self.model), norm='none', activ='relu', pad_type='reflect')
|
||||
# label encoder
|
||||
self.enc_label = LabelEncoder(5, 19, 16, 64, norm='none', activ='relu', pad_type='reflect')
|
||||
|
||||
def assign_adain_params(self, adain_params, model):
|
||||
# assign the adain_params to the AdaIN layers in model
|
||||
for m in model.modules():
|
||||
if m.__class__.__name__ == "AdaptiveInstanceNorm2d":
|
||||
mean = adain_params[:, :m.num_features]
|
||||
std = adain_params[:, m.num_features:2*m.num_features]
|
||||
m.bias = mean.contiguous().view(-1)
|
||||
m.weight = std.contiguous().view(-1)
|
||||
if adain_params.size(1) > 2*m.num_features:
|
||||
adain_params = adain_params[:, 2*m.num_features:]
|
||||
|
||||
def get_num_adain_params(self, model):
|
||||
# return the number of AdaIN parameters needed by the model
|
||||
num_adain_params = 0
|
||||
for m in model.modules():
|
||||
if m.__class__.__name__ == "AdaptiveInstanceNorm2d":
|
||||
num_adain_params += 2*m.num_features
|
||||
return num_adain_params
|
||||
|
||||
def forward(self, input, input_ref, image_ref):
|
||||
fea1, fea2 = self.enc_label(input_ref)
|
||||
adain_params = self.enc_style((image_ref, fea1, fea2))
|
||||
self.assign_adain_params(adain_params, self.model)
|
||||
return self.model(input)
|
||||
|
||||
class BlendGenerator(nn.Module):
|
||||
def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=3, norm_layer=nn.BatchNorm2d,
|
||||
padding_type='reflect'):
|
||||
assert(n_blocks >= 0)
|
||||
super(BlendGenerator, self).__init__()
|
||||
activation = nn.ReLU(True)
|
||||
|
||||
model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
|
||||
### downsample
|
||||
for i in range(n_downsampling):
|
||||
mult = 2**i
|
||||
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
|
||||
norm_layer(ngf * mult * 2), activation]
|
||||
|
||||
### resnet blocks
|
||||
mult = 2**n_downsampling
|
||||
for i in range(n_blocks):
|
||||
model += [ResnetBlock(ngf * mult, norm_type='in', padding_type=padding_type)]
|
||||
|
||||
### upsample
|
||||
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),
|
||||
norm_layer(int(ngf * mult / 2)), activation]
|
||||
model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Sigmoid()]
|
||||
self.model = nn.Sequential(*model)
|
||||
|
||||
def forward(self, input1, input2):
|
||||
m = self.model(torch.cat([input1, input2], 1))
|
||||
return input1 * m + input2 * (1-m), m
|
||||
|
||||
# Define the Multiscale Discriminator.
|
||||
class MultiscaleDiscriminator(nn.Module):
|
||||
def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d,
|
||||
use_sigmoid=False, num_D=3, getIntermFeat=False):
|
||||
super(MultiscaleDiscriminator, self).__init__()
|
||||
self.num_D = num_D
|
||||
self.n_layers = n_layers
|
||||
self.getIntermFeat = getIntermFeat
|
||||
|
||||
for i in range(num_D):
|
||||
netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
|
||||
if getIntermFeat:
|
||||
for j in range(n_layers+2):
|
||||
setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))
|
||||
else:
|
||||
setattr(self, 'layer'+str(i), netD.model)
|
||||
|
||||
self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
|
||||
|
||||
def singleD_forward(self, model, input):
|
||||
if self.getIntermFeat:
|
||||
result = [input]
|
||||
for i in range(len(model)):
|
||||
result.append(model[i](result[-1]))
|
||||
return result[1:]
|
||||
else:
|
||||
return [model(input)]
|
||||
|
||||
def forward(self, input):
|
||||
num_D = self.num_D
|
||||
result = []
|
||||
input_downsampled = input
|
||||
for i in range(num_D):
|
||||
if self.getIntermFeat:
|
||||
model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
|
||||
else:
|
||||
model = getattr(self, 'layer'+str(num_D-1-i))
|
||||
result.append(self.singleD_forward(model, input_downsampled))
|
||||
if i != (num_D-1):
|
||||
input_downsampled = self.downsample(input_downsampled)
|
||||
return result
|
||||
|
||||
# Define 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, getIntermFeat=False):
|
||||
super(NLayerDiscriminator, self).__init__()
|
||||
self.getIntermFeat = getIntermFeat
|
||||
self.n_layers = n_layers
|
||||
|
||||
kw = 4
|
||||
padw = int(np.ceil((kw-1.0)/2))
|
||||
sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
|
||||
|
||||
nf = ndf
|
||||
for n in range(1, n_layers):
|
||||
nf_prev = nf
|
||||
nf = min(nf * 2, 512)
|
||||
sequence += [[
|
||||
nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
|
||||
norm_layer(nf), nn.LeakyReLU(0.2, True)
|
||||
]]
|
||||
|
||||
nf_prev = nf
|
||||
nf = min(nf * 2, 512)
|
||||
sequence += [[
|
||||
nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
|
||||
norm_layer(nf),
|
||||
nn.LeakyReLU(0.2, True)
|
||||
]]
|
||||
|
||||
sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
|
||||
|
||||
if use_sigmoid:
|
||||
sequence += [[nn.Sigmoid()]]
|
||||
|
||||
if getIntermFeat:
|
||||
for n in range(len(sequence)):
|
||||
setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
|
||||
else:
|
||||
sequence_stream = []
|
||||
for n in range(len(sequence)):
|
||||
sequence_stream += sequence[n]
|
||||
self.model = nn.Sequential(*sequence_stream)
|
||||
|
||||
def forward(self, input):
|
||||
if self.getIntermFeat:
|
||||
res = [input]
|
||||
for n in range(self.n_layers+2):
|
||||
model = getattr(self, 'model'+str(n))
|
||||
res.append(model(res[-1]))
|
||||
return res[1:]
|
||||
else:
|
||||
return self.model(input)
|
||||
|
||||
from torchvision import models
|
||||
class Vgg19(torch.nn.Module):
|
||||
def __init__(self, requires_grad=False):
|
||||
super(Vgg19, self).__init__()
|
||||
vgg_pretrained_features = models.vgg19(pretrained=True).features
|
||||
self.slice1 = torch.nn.Sequential()
|
||||
self.slice2 = torch.nn.Sequential()
|
||||
self.slice3 = torch.nn.Sequential()
|
||||
self.slice4 = torch.nn.Sequential()
|
||||
self.slice5 = torch.nn.Sequential()
|
||||
for x in range(2):
|
||||
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(2, 7):
|
||||
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(7, 12):
|
||||
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(12, 21):
|
||||
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(21, 30):
|
||||
self.slice5.add_module(str(x), vgg_pretrained_features[x])
|
||||
if not requires_grad:
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, X):
|
||||
h_relu1 = self.slice1(X)
|
||||
h_relu2 = self.slice2(h_relu1)
|
||||
h_relu3 = self.slice3(h_relu2)
|
||||
h_relu4 = self.slice4(h_relu3)
|
||||
h_relu5 = self.slice5(h_relu4)
|
||||
out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
|
||||
return out
|
||||
|
||||
# Define the MaskVAE
|
||||
class VAE(nn.Module):
|
||||
def __init__(self, nc, ngf, ndf, latent_variable_size):
|
||||
super(VAE, self).__init__()
|
||||
#self.cuda = True
|
||||
self.nc = nc
|
||||
self.ngf = ngf
|
||||
self.ndf = ndf
|
||||
self.latent_variable_size = latent_variable_size
|
||||
|
||||
# encoder
|
||||
self.e1 = nn.Conv2d(nc, ndf, 4, 2, 1)
|
||||
self.bn1 = nn.BatchNorm2d(ndf)
|
||||
|
||||
self.e2 = nn.Conv2d(ndf, ndf*2, 4, 2, 1)
|
||||
self.bn2 = nn.BatchNorm2d(ndf*2)
|
||||
|
||||
self.e3 = nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)
|
||||
self.bn3 = nn.BatchNorm2d(ndf*4)
|
||||
|
||||
self.e4 = nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)
|
||||
self.bn4 = nn.BatchNorm2d(ndf*8)
|
||||
|
||||
self.e5 = nn.Conv2d(ndf*8, ndf*16, 4, 2, 1)
|
||||
self.bn5 = nn.BatchNorm2d(ndf*16)
|
||||
|
||||
self.e6 = nn.Conv2d(ndf*16, ndf*32, 4, 2, 1)
|
||||
self.bn6 = nn.BatchNorm2d(ndf*32)
|
||||
|
||||
self.e7 = nn.Conv2d(ndf*32, ndf*64, 4, 2, 1)
|
||||
self.bn7 = nn.BatchNorm2d(ndf*64)
|
||||
|
||||
self.fc1 = nn.Linear(ndf*64*4*4, latent_variable_size)
|
||||
self.fc2 = nn.Linear(ndf*64*4*4, latent_variable_size)
|
||||
|
||||
# decoder
|
||||
self.d1 = nn.Linear(latent_variable_size, ngf*64*4*4)
|
||||
|
||||
self.up1 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd1 = nn.ReplicationPad2d(1)
|
||||
self.d2 = nn.Conv2d(ngf*64, ngf*32, 3, 1)
|
||||
self.bn8 = nn.BatchNorm2d(ngf*32, 1.e-3)
|
||||
|
||||
self.up2 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd2 = nn.ReplicationPad2d(1)
|
||||
self.d3 = nn.Conv2d(ngf*32, ngf*16, 3, 1)
|
||||
self.bn9 = nn.BatchNorm2d(ngf*16, 1.e-3)
|
||||
|
||||
self.up3 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd3 = nn.ReplicationPad2d(1)
|
||||
self.d4 = nn.Conv2d(ngf*16, ngf*8, 3, 1)
|
||||
self.bn10 = nn.BatchNorm2d(ngf*8, 1.e-3)
|
||||
|
||||
self.up4 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd4 = nn.ReplicationPad2d(1)
|
||||
self.d5 = nn.Conv2d(ngf*8, ngf*4, 3, 1)
|
||||
self.bn11 = nn.BatchNorm2d(ngf*4, 1.e-3)
|
||||
|
||||
self.up5 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd5 = nn.ReplicationPad2d(1)
|
||||
self.d6 = nn.Conv2d(ngf*4, ngf*2, 3, 1)
|
||||
self.bn12 = nn.BatchNorm2d(ngf*2, 1.e-3)
|
||||
|
||||
self.up6 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd6 = nn.ReplicationPad2d(1)
|
||||
self.d7 = nn.Conv2d(ngf*2, ngf, 3, 1)
|
||||
self.bn13 = nn.BatchNorm2d(ngf, 1.e-3)
|
||||
|
||||
self.up7 = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.pd7 = nn.ReplicationPad2d(1)
|
||||
self.d8 = nn.Conv2d(ngf, nc, 3, 1)
|
||||
|
||||
self.leakyrelu = nn.LeakyReLU(0.2)
|
||||
self.relu = nn.ReLU()
|
||||
#self.sigmoid = nn.Sigmoid()
|
||||
self.maxpool = nn.MaxPool2d((2, 2), (2, 2))
|
||||
|
||||
def encode(self, x):
|
||||
h1 = self.leakyrelu(self.bn1(self.e1(x)))
|
||||
h2 = self.leakyrelu(self.bn2(self.e2(h1)))
|
||||
h3 = self.leakyrelu(self.bn3(self.e3(h2)))
|
||||
h4 = self.leakyrelu(self.bn4(self.e4(h3)))
|
||||
h5 = self.leakyrelu(self.bn5(self.e5(h4)))
|
||||
h6 = self.leakyrelu(self.bn6(self.e6(h5)))
|
||||
h7 = self.leakyrelu(self.bn7(self.e7(h6)))
|
||||
h7 = h7.view(-1, self.ndf*64*4*4)
|
||||
return self.fc1(h7), self.fc2(h7)
|
||||
|
||||
def reparametrize(self, mu, logvar):
|
||||
std = logvar.mul(0.5).exp_()
|
||||
#if self.cuda:
|
||||
eps = torch.cuda.FloatTensor(std.size()).normal_()
|
||||
#else:
|
||||
# eps = torch.FloatTensor(std.size()).normal_()
|
||||
eps = Variable(eps)
|
||||
return eps.mul(std).add_(mu)
|
||||
|
||||
def decode(self, z):
|
||||
h1 = self.relu(self.d1(z))
|
||||
h1 = h1.view(-1, self.ngf*64, 4, 4)
|
||||
h2 = self.leakyrelu(self.bn8(self.d2(self.pd1(self.up1(h1)))))
|
||||
h3 = self.leakyrelu(self.bn9(self.d3(self.pd2(self.up2(h2)))))
|
||||
h4 = self.leakyrelu(self.bn10(self.d4(self.pd3(self.up3(h3)))))
|
||||
h5 = self.leakyrelu(self.bn11(self.d5(self.pd4(self.up4(h4)))))
|
||||
h6 = self.leakyrelu(self.bn12(self.d6(self.pd5(self.up5(h5)))))
|
||||
h7 = self.leakyrelu(self.bn13(self.d7(self.pd6(self.up6(h6)))))
|
||||
return self.d8(self.pd7(self.up7(h7)))
|
||||
|
||||
def get_latent_var(self, x):
|
||||
mu, logvar = self.encode(x)
|
||||
z = self.reparametrize(mu, logvar)
|
||||
return z, mu, logvar.mul(0.5).exp_()
|
||||
|
||||
def forward(self, x):
|
||||
mu, logvar = self.encode(x)
|
||||
z = self.reparametrize(mu, logvar)
|
||||
res = self.decode(z)
|
||||
|
||||
return res, x, mu, logvar
|
||||
|
||||
# style encode part
|
||||
class StyleEncoder(nn.Module):
|
||||
def __init__(self, n_downsample, input_dim, dim, style_dim, norm, activ, pad_type):
|
||||
super(StyleEncoder, self).__init__()
|
||||
self.model = []
|
||||
self.model_middle = []
|
||||
self.model_last = []
|
||||
self.model += [ConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
for i in range(2):
|
||||
self.model += [ConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
dim *= 2
|
||||
for i in range(n_downsample - 2):
|
||||
self.model_middle += [ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
self.model_last += [nn.AdaptiveAvgPool2d(1)] # global average pooling
|
||||
self.model_last += [nn.Conv2d(dim, style_dim, 1, 1, 0)]
|
||||
|
||||
self.model = nn.Sequential(*self.model)
|
||||
self.model_middle = nn.Sequential(*self.model_middle)
|
||||
self.model_last = nn.Sequential(*self.model_last)
|
||||
|
||||
self.output_dim = dim
|
||||
|
||||
self.sft1 = SFTLayer()
|
||||
self.sft2 = SFTLayer()
|
||||
|
||||
def forward(self, x):
|
||||
fea = self.model(x[0])
|
||||
fea = self.sft1((fea, x[1]))
|
||||
fea = self.model_middle(fea)
|
||||
fea = self.sft2((fea, x[2]))
|
||||
return self.model_last(fea)
|
||||
|
||||
# label encode part
|
||||
class LabelEncoder(nn.Module):
|
||||
def __init__(self, n_downsample, input_dim, dim, style_dim, norm, activ, pad_type):
|
||||
super(LabelEncoder, self).__init__()
|
||||
self.model = []
|
||||
self.model_last = [nn.ReLU()]
|
||||
self.model += [ConvBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
self.model += [ConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
dim *= 2
|
||||
self.model += [ConvBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation='none', pad_type=pad_type)]
|
||||
dim *= 2
|
||||
for i in range(n_downsample - 3):
|
||||
self.model_last += [ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type=pad_type)]
|
||||
self.model_last += [ConvBlock(dim, dim, 4, 2, 1, norm=norm, activation='none', pad_type=pad_type)]
|
||||
self.model = nn.Sequential(*self.model)
|
||||
self.model_last = nn.Sequential(*self.model_last)
|
||||
self.output_dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
fea = self.model(x)
|
||||
return fea, self.model_last(fea)
|
||||
|
||||
# Define the basic block
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, input_dim ,output_dim, kernel_size, stride,
|
||||
padding=0, norm='none', activation='relu', pad_type='zero'):
|
||||
super(ConvBlock, self).__init__()
|
||||
self.use_bias = True
|
||||
# initialize padding
|
||||
if pad_type == 'reflect':
|
||||
self.pad = nn.ReflectionPad2d(padding)
|
||||
elif pad_type == 'replicate':
|
||||
self.pad = nn.ReplicationPad2d(padding)
|
||||
elif pad_type == 'zero':
|
||||
self.pad = nn.ZeroPad2d(padding)
|
||||
else:
|
||||
assert 0, "Unsupported padding type: {}".format(pad_type)
|
||||
|
||||
# initialize normalization
|
||||
norm_dim = output_dim
|
||||
if norm == 'bn':
|
||||
self.norm = nn.BatchNorm2d(norm_dim)
|
||||
elif norm == 'in':
|
||||
#self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
|
||||
self.norm = nn.InstanceNorm2d(norm_dim)
|
||||
elif norm == 'ln':
|
||||
self.norm = LayerNorm(norm_dim)
|
||||
elif norm == 'adain':
|
||||
self.norm = AdaptiveInstanceNorm2d(norm_dim)
|
||||
elif norm == 'none' or norm == 'sn':
|
||||
self.norm = None
|
||||
else:
|
||||
assert 0, "Unsupported normalization: {}".format(norm)
|
||||
|
||||
# initialize activation
|
||||
if activation == 'relu':
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
elif activation == 'lrelu':
|
||||
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||
elif activation == 'prelu':
|
||||
self.activation = nn.PReLU()
|
||||
elif activation == 'selu':
|
||||
self.activation = nn.SELU(inplace=True)
|
||||
elif activation == 'tanh':
|
||||
self.activation = nn.Tanh()
|
||||
elif activation == 'none':
|
||||
self.activation = None
|
||||
else:
|
||||
assert 0, "Unsupported activation: {}".format(activation)
|
||||
|
||||
# initialize convolution
|
||||
if norm == 'sn':
|
||||
self.conv = SpectralNorm(nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias))
|
||||
else:
|
||||
self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(self.pad(x))
|
||||
if self.norm:
|
||||
x = self.norm(x)
|
||||
if self.activation:
|
||||
x = self.activation(x)
|
||||
return x
|
||||
|
||||
class LinearBlock(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, norm='none', activation='relu'):
|
||||
super(LinearBlock, self).__init__()
|
||||
use_bias = True
|
||||
# initialize fully connected layer
|
||||
if norm == 'sn':
|
||||
self.fc = SpectralNorm(nn.Linear(input_dim, output_dim, bias=use_bias))
|
||||
else:
|
||||
self.fc = nn.Linear(input_dim, output_dim, bias=use_bias)
|
||||
|
||||
# initialize normalization
|
||||
norm_dim = output_dim
|
||||
if norm == 'bn':
|
||||
self.norm = nn.BatchNorm1d(norm_dim)
|
||||
elif norm == 'in':
|
||||
self.norm = nn.InstanceNorm1d(norm_dim)
|
||||
elif norm == 'ln':
|
||||
self.norm = LayerNorm(norm_dim)
|
||||
elif norm == 'none' or norm == 'sn':
|
||||
self.norm = None
|
||||
else:
|
||||
assert 0, "Unsupported normalization: {}".format(norm)
|
||||
|
||||
# initialize activation
|
||||
if activation == 'relu':
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
elif activation == 'lrelu':
|
||||
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||
elif activation == 'prelu':
|
||||
self.activation = nn.PReLU()
|
||||
elif activation == 'selu':
|
||||
self.activation = nn.SELU(inplace=True)
|
||||
elif activation == 'tanh':
|
||||
self.activation = nn.Tanh()
|
||||
elif activation == 'none':
|
||||
self.activation = None
|
||||
else:
|
||||
assert 0, "Unsupported activation: {}".format(activation)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.fc(x)
|
||||
if self.norm:
|
||||
out = self.norm(out)
|
||||
if self.activation:
|
||||
out = self.activation(out)
|
||||
return out
|
||||
|
||||
# Define a resnet block
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, dim, norm_type, padding_type, use_dropout=False):
|
||||
super(ResnetBlock, self).__init__()
|
||||
self.conv_block = self.build_conv_block(dim, norm_type, padding_type, use_dropout)
|
||||
|
||||
def build_conv_block(self, dim, norm_type, padding_type, use_dropout):
|
||||
conv_block = []
|
||||
conv_block += [ConvBlock(dim ,dim, 3, 1, 1, norm=norm_type, activation='relu', pad_type=padding_type)]
|
||||
conv_block += [ConvBlock(dim ,dim, 3, 1, 1, norm=norm_type, activation='none', pad_type=padding_type)]
|
||||
|
||||
return nn.Sequential(*conv_block)
|
||||
|
||||
def forward(self, x):
|
||||
out = x + self.conv_block(x)
|
||||
return out
|
||||
|
||||
class SFTLayer(nn.Module):
|
||||
def __init__(self):
|
||||
super(SFTLayer, self).__init__()
|
||||
self.SFT_scale_conv1 = nn.Conv2d(64, 64, 1)
|
||||
self.SFT_scale_conv2 = nn.Conv2d(64, 64, 1)
|
||||
self.SFT_shift_conv1 = nn.Conv2d(64, 64, 1)
|
||||
self.SFT_shift_conv2 = nn.Conv2d(64, 64, 1)
|
||||
|
||||
def forward(self, x):
|
||||
scale = self.SFT_scale_conv2(F.leaky_relu(self.SFT_scale_conv1(x[1]), 0.1, inplace=True))
|
||||
shift = self.SFT_shift_conv2(F.leaky_relu(self.SFT_shift_conv1(x[1]), 0.1, inplace=True))
|
||||
return x[0] * scale + shift
|
||||
|
||||
class ConvBlock_SFT(nn.Module):
|
||||
def __init__(self, dim, norm_type, padding_type, use_dropout=False):
|
||||
super(ResnetBlock_SFT, self).__init__()
|
||||
self.sft1 = SFTLayer()
|
||||
self.conv1 = ConvBlock(dim ,dim, 4, 2, 1, norm=norm_type, activation='none', pad_type=padding_type)
|
||||
|
||||
def forward(self, x):
|
||||
fea = self.sft1((x[0], x[1]))
|
||||
fea = F.relu(self.conv1(fea), inplace=True)
|
||||
return (x[0] + fea, x[1])
|
||||
|
||||
class ConvBlock_SFT_last(nn.Module):
|
||||
def __init__(self, dim, norm_type, padding_type, use_dropout=False):
|
||||
super(ResnetBlock_SFT_last, self).__init__()
|
||||
self.sft1 = SFTLayer()
|
||||
self.conv1 = ConvBlock(dim ,dim, 4, 2, 1, norm=norm_type, activation='none', pad_type=padding_type)
|
||||
|
||||
def forward(self, x):
|
||||
fea = self.sft1((x[0], x[1]))
|
||||
fea = F.relu(self.conv1(fea), inplace=True)
|
||||
return x[0] + fea
|
||||
|
||||
# Definition of normalization layer
|
||||
class AdaptiveInstanceNorm2d(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5, momentum=0.1):
|
||||
super(AdaptiveInstanceNorm2d, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
# weight and bias are dynamically assigned
|
||||
self.weight = None
|
||||
self.bias = None
|
||||
# just dummy buffers, not used
|
||||
self.register_buffer('running_mean', torch.zeros(num_features))
|
||||
self.register_buffer('running_var', torch.ones(num_features))
|
||||
|
||||
def forward(self, x):
|
||||
assert self.weight is not None and self.bias is not None, "Please assign weight and bias before calling AdaIN!"
|
||||
b, c = x.size(0), x.size(1)
|
||||
running_mean = self.running_mean.repeat(b)
|
||||
running_var = self.running_var.repeat(b)
|
||||
|
||||
# Apply instance norm
|
||||
x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])
|
||||
|
||||
out = F.batch_norm(
|
||||
x_reshaped, running_mean, running_var, self.weight, self.bias,
|
||||
True, self.momentum, self.eps)
|
||||
|
||||
return out.view(b, c, *x.size()[2:])
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + '(' + str(self.num_features) + ')'
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5, affine=True):
|
||||
super(LayerNorm, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.affine = affine
|
||||
self.eps = eps
|
||||
|
||||
if self.affine:
|
||||
self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_())
|
||||
self.beta = nn.Parameter(torch.zeros(num_features))
|
||||
|
||||
def forward(self, x):
|
||||
shape = [-1] + [1] * (x.dim() - 1)
|
||||
# print(x.size())
|
||||
if x.size(0) == 1:
|
||||
# These two lines run much faster in pytorch 0.4 than the two lines listed below.
|
||||
mean = x.view(-1).mean().view(*shape)
|
||||
std = x.view(-1).std().view(*shape)
|
||||
else:
|
||||
mean = x.view(x.size(0), -1).mean(1).view(*shape)
|
||||
std = x.view(x.size(0), -1).std(1).view(*shape)
|
||||
|
||||
x = (x - mean) / (std + self.eps)
|
||||
|
||||
if self.affine:
|
||||
shape = [1, -1] + [1] * (x.dim() - 2)
|
||||
x = x * self.gamma.view(*shape) + self.beta.view(*shape)
|
||||
return x
|
||||
|
||||
def l2normalize(v, eps=1e-12):
|
||||
return v / (v.norm() + eps)
|
||||
|
||||
class SpectralNorm(nn.Module):
|
||||
"""
|
||||
Based on the paper "Spectral Normalization for Generative Adversarial Networks" by Takeru Miyato, Toshiki Kataoka, Masanori Koyama, Yuichi Yoshida
|
||||
and the Pytorch implementation https://github.com/christiancosgrove/pytorch-spectral-normalization-gan
|
||||
"""
|
||||
def __init__(self, module, name='weight', power_iterations=1):
|
||||
super(SpectralNorm, self).__init__()
|
||||
self.module = module
|
||||
self.name = name
|
||||
self.power_iterations = power_iterations
|
||||
if not self._made_params():
|
||||
self._make_params()
|
||||
|
||||
def _update_u_v(self):
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
for _ in range(self.power_iterations):
|
||||
v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
|
||||
u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
|
||||
|
||||
# sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||
|
||||
def _made_params(self):
|
||||
try:
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
return True
|
||||
except AttributeError:
|
||||
return False
|
||||
|
||||
|
||||
def _make_params(self):
|
||||
w = getattr(self.module, self.name)
|
||||
|
||||
height = w.data.shape[0]
|
||||
width = w.view(height, -1).data.shape[1]
|
||||
|
||||
u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
|
||||
v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
|
||||
u.data = l2normalize(u.data)
|
||||
v.data = l2normalize(v.data)
|
||||
w_bar = nn.Parameter(w.data)
|
||||
|
||||
del self.module._parameters[self.name]
|
||||
|
||||
self.module.register_parameter(self.name + "_u", u)
|
||||
self.module.register_parameter(self.name + "_v", v)
|
||||
self.module.register_parameter(self.name + "_bar", w_bar)
|
||||
|
||||
def forward(self, *args):
|
||||
self._update_u_v()
|
||||
return self.module.forward(*args)
|
||||
@ -1,326 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import numpy as np
|
||||
import torch
|
||||
import os
|
||||
from torch.autograd import Variable
|
||||
from util.image_pool import ImagePool
|
||||
from .base_model import BaseModel
|
||||
from . import networks
|
||||
|
||||
def generate_discrete_label(inputs, label_nc):
|
||||
pred_batch = []
|
||||
size = inputs.size()
|
||||
for input in inputs:
|
||||
input = input.view(1, label_nc, size[2], size[3])
|
||||
pred = np.squeeze(input.data.max(1)[1].cpu().numpy(), axis=0)
|
||||
pred_batch.append(pred)
|
||||
|
||||
pred_batch = np.array(pred_batch)
|
||||
pred_batch = torch.from_numpy(pred_batch)
|
||||
label_map = []
|
||||
for p in pred_batch:
|
||||
p = p.view(1, 512, 512)
|
||||
label_map.append(p)
|
||||
label_map = torch.stack(label_map, 0)
|
||||
size = label_map.size()
|
||||
oneHot_size = (size[0], label_nc, size[2], size[3])
|
||||
if torch.cuda.is_available():
|
||||
input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
|
||||
else:
|
||||
input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long(), 1.0)
|
||||
|
||||
return input_label
|
||||
|
||||
class Pix2PixHDModel(BaseModel):
|
||||
def name(self):
|
||||
return 'Pix2PixHDModel'
|
||||
|
||||
def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
|
||||
flags = (True, use_gan_feat_loss, use_vgg_loss, True, use_gan_feat_loss, use_vgg_loss, True, True, True, True)
|
||||
def loss_filter(g_gan, g_gan_feat, g_vgg, gb_gan, gb_gan_feat, gb_vgg, d_real, d_fake, d_blend):
|
||||
return [l for (l,f) in zip((g_gan,g_gan_feat,g_vgg,gb_gan,gb_gan_feat,gb_vgg,d_real,d_fake,d_blend),flags) if f]
|
||||
return loss_filter
|
||||
|
||||
def initialize(self, opt):
|
||||
BaseModel.initialize(self, opt)
|
||||
if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
|
||||
torch.backends.cudnn.benchmark = True
|
||||
self.isTrain = opt.isTrain
|
||||
input_nc = opt.label_nc if opt.label_nc != 0 else opt.input_nc
|
||||
|
||||
##### define networks
|
||||
# Generator network
|
||||
netG_input_nc = input_nc
|
||||
# Main Generator
|
||||
self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG,
|
||||
opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers,
|
||||
opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)
|
||||
|
||||
# Discriminator network
|
||||
if self.isTrain:
|
||||
use_sigmoid = opt.no_lsgan
|
||||
netD_input_nc = input_nc + opt.output_nc
|
||||
netB_input_nc = opt.output_nc * 2
|
||||
self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid,
|
||||
opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
|
||||
self.netB = networks.define_B(netB_input_nc, opt.output_nc, 32, 3, 3, opt.norm, gpu_ids=self.gpu_ids)
|
||||
|
||||
if self.opt.verbose:
|
||||
print('---------- Networks initialized -------------')
|
||||
|
||||
# load networks
|
||||
if not self.isTrain or opt.continue_train or opt.load_pretrain:
|
||||
pretrained_path = '' if not self.isTrain else opt.load_pretrain
|
||||
print (pretrained_path)
|
||||
self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
|
||||
if self.isTrain:
|
||||
self.load_network(self.netB, 'B', opt.which_epoch, pretrained_path)
|
||||
self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)
|
||||
|
||||
# set loss functions and optimizers
|
||||
if self.isTrain:
|
||||
if opt.pool_size > 0 and (len(self.gpu_ids)) > 1:
|
||||
raise NotImplementedError("Fake Pool Not Implemented for MultiGPU")
|
||||
self.fake_pool = ImagePool(opt.pool_size)
|
||||
self.old_lr = opt.lr
|
||||
|
||||
# define loss functions
|
||||
self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
|
||||
|
||||
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
|
||||
self.criterionFeat = torch.nn.L1Loss()
|
||||
if not opt.no_vgg_loss:
|
||||
self.criterionVGG = networks.VGGLoss(self.gpu_ids)
|
||||
|
||||
# Names so we can breakout loss
|
||||
self.loss_names = self.loss_filter('G_GAN','G_GAN_Feat','G_VGG','GB_GAN','GB_GAN_Feat','GB_VGG','D_real','D_fake','D_blend')
|
||||
# initialize optimizers
|
||||
# optimizer G
|
||||
if opt.niter_fix_global > 0:
|
||||
import sys
|
||||
if sys.version_info >= (3,0):
|
||||
finetune_list = set()
|
||||
else:
|
||||
from sets import Set
|
||||
finetune_list = Set()
|
||||
|
||||
params_dict = dict(self.netG.named_parameters())
|
||||
params = []
|
||||
for key, value in params_dict.items():
|
||||
if key.startswith('model' + str(opt.n_local_enhancers)):
|
||||
params += [value]
|
||||
finetune_list.add(key.split('.')[0])
|
||||
print('------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global)
|
||||
print('The layers that are finetuned are ', sorted(finetune_list))
|
||||
else:
|
||||
params = list(self.netG.parameters())
|
||||
|
||||
self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
|
||||
|
||||
# optimizer D
|
||||
params = list(self.netD.parameters())
|
||||
self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
|
||||
|
||||
# optimizer G + B
|
||||
params = list(self.netG.parameters()) + list(self.netB.parameters())
|
||||
self.optimizer_GB = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
|
||||
|
||||
def encode_input(self, inter_label_map_1, label_map, inter_label_map_2, real_image, label_map_ref, real_image_ref, infer=False):
|
||||
|
||||
if self.opt.label_nc == 0:
|
||||
if torch.cuda.is_available():
|
||||
input_label = label_map.data.cuda()
|
||||
inter_label_1 = inter_label_map_1.data.cuda()
|
||||
inter_label_2 = inter_label_map_2.data.cuda()
|
||||
input_label_ref = label_map_ref.data.cuda()
|
||||
else:
|
||||
input_label = label_map.data
|
||||
inter_label_1 = inter_label_map_1.data
|
||||
inter_label_2 = inter_label_map_2.data
|
||||
input_label_ref = label_map_ref.data
|
||||
|
||||
else:
|
||||
# create one-hot vector for label map
|
||||
size = label_map.size()
|
||||
oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
|
||||
if torch.cuda.is_available():
|
||||
input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
|
||||
inter_label_1 = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
inter_label_1 = inter_label_1.scatter_(1, inter_label_map_1.data.long().cuda(), 1.0)
|
||||
inter_label_2 = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
inter_label_2 = inter_label_2.scatter_(1, inter_label_map_2.data.long().cuda(), 1.0)
|
||||
input_label_ref = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label_ref = input_label_ref.scatter_(1, label_map_ref.data.long().cuda(), 1.0)
|
||||
else:
|
||||
input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long(), 1.0)
|
||||
inter_label_1 = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
inter_label_1 = inter_label_1.scatter_(1, inter_label_map_1.data.long(), 1.0)
|
||||
inter_label_2 = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
inter_label_2 = inter_label_2.scatter_(1, inter_label_map_2.data.long(), 1.0)
|
||||
input_label_ref = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label_ref = input_label_ref.scatter_(1, label_map_ref.data.long(), 1.0)
|
||||
|
||||
if self.opt.data_type == 16:
|
||||
input_label = input_label.half()
|
||||
inter_label_1 = inter_label_1.half()
|
||||
inter_label_2 = inter_label_2.half()
|
||||
input_label_ref = input_label_ref.half()
|
||||
|
||||
input_label = Variable(input_label, volatile=infer)
|
||||
inter_label_1 = Variable(inter_label_1, volatile=infer)
|
||||
inter_label_2 = Variable(inter_label_2, volatile=infer)
|
||||
input_label_ref = Variable(input_label_ref, volatile=infer)
|
||||
if torch.cuda.is_available():
|
||||
real_image = Variable(real_image.data.cuda())
|
||||
real_image_ref = Variable(real_image_ref.data.cuda())
|
||||
else:
|
||||
real_image = Variable(real_image.data)
|
||||
real_image_ref = Variable(real_image_ref.data)
|
||||
|
||||
return inter_label_1, input_label, inter_label_2, real_image, input_label_ref, real_image_ref
|
||||
|
||||
def encode_input_test(self, label_map, label_map_ref, real_image_ref, infer=False,f=False):
|
||||
|
||||
if self.opt.label_nc == 0:
|
||||
if torch.cuda.is_available():
|
||||
input_label = label_map.data.cuda()
|
||||
input_label_ref = label_map_ref.data.cuda()
|
||||
else:
|
||||
input_label = label_map.data
|
||||
input_label_ref = label_map_ref.data
|
||||
|
||||
else:
|
||||
# create one-hot vector for label map
|
||||
size = label_map.size()
|
||||
oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
|
||||
if torch.cuda.is_available() and not f:
|
||||
input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
|
||||
input_label_ref = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label_ref = input_label_ref.scatter_(1, label_map_ref.data.long().cuda(), 1.0)
|
||||
real_image_ref = Variable(real_image_ref.data.cuda())
|
||||
|
||||
else:
|
||||
input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label = input_label.scatter_(1, label_map.data.long(), 1.0)
|
||||
input_label_ref = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
|
||||
input_label_ref = input_label_ref.scatter_(1, label_map_ref.data.long(), 1.0)
|
||||
real_image_ref = Variable(real_image_ref.data)
|
||||
|
||||
|
||||
if self.opt.data_type == 16:
|
||||
input_label = input_label.half()
|
||||
input_label_ref = input_label_ref.half()
|
||||
|
||||
input_label = Variable(input_label, volatile=infer)
|
||||
input_label_ref = Variable(input_label_ref, volatile=infer)
|
||||
|
||||
return input_label, input_label_ref, real_image_ref
|
||||
|
||||
def discriminate(self, input_label, test_image, use_pool=False):
|
||||
input_concat = torch.cat((input_label, test_image.detach()), dim=1)
|
||||
if use_pool:
|
||||
fake_query = self.fake_pool.query(input_concat)
|
||||
return self.netD.forward(fake_query)
|
||||
else:
|
||||
return self.netD.forward(input_concat)
|
||||
|
||||
def forward(self, inter_label_1, label, inter_label_2, image, label_ref, image_ref, infer=False):
|
||||
|
||||
# Encode Inputs
|
||||
inter_label_1, input_label, inter_label_2, real_image, input_label_ref, real_image_ref = self.encode_input(inter_label_1, label, inter_label_2, image, label_ref, image_ref)
|
||||
|
||||
fake_inter_1 = self.netG.forward(inter_label_1, input_label, real_image)
|
||||
fake_image = self.netG.forward(input_label, input_label, real_image)
|
||||
fake_inter_2 = self.netG.forward(inter_label_2, input_label, real_image)
|
||||
|
||||
blend_image, alpha = self.netB.forward(fake_inter_1, fake_inter_2)
|
||||
|
||||
# Fake Detection and Loss
|
||||
pred_fake_pool = self.discriminate(input_label, fake_image, use_pool=True)
|
||||
loss_D_fake = self.criterionGAN(pred_fake_pool, False)
|
||||
pred_blend_pool = self.discriminate(input_label, blend_image, use_pool=True)
|
||||
loss_D_blend = self.criterionGAN(pred_blend_pool, False)
|
||||
|
||||
# Real Detection and Loss
|
||||
pred_real = self.discriminate(input_label, real_image)
|
||||
loss_D_real = self.criterionGAN(pred_real, True)
|
||||
|
||||
# GAN loss (Fake Passability Loss)
|
||||
pred_fake = self.netD.forward(torch.cat((input_label, fake_image), dim=1))
|
||||
loss_G_GAN = self.criterionGAN(pred_fake, True)
|
||||
pred_blend = self.netD.forward(torch.cat((input_label, blend_image), dim=1))
|
||||
loss_GB_GAN = self.criterionGAN(pred_blend, True)
|
||||
|
||||
# GAN feature matching loss
|
||||
loss_G_GAN_Feat = 0
|
||||
loss_GB_GAN_Feat = 0
|
||||
if not self.opt.no_ganFeat_loss:
|
||||
feat_weights = 4.0 / (self.opt.n_layers_D + 1)
|
||||
D_weights = 1.0 / self.opt.num_D
|
||||
for i in range(self.opt.num_D):
|
||||
for j in range(len(pred_fake[i])-1):
|
||||
loss_G_GAN_Feat += D_weights * feat_weights * \
|
||||
self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
|
||||
loss_GB_GAN_Feat += D_weights * feat_weights * \
|
||||
self.criterionFeat(pred_blend[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
|
||||
|
||||
# VGG feature matching loss
|
||||
loss_G_VGG = 0
|
||||
loss_GB_VGG = 0
|
||||
if not self.opt.no_vgg_loss:
|
||||
loss_G_VGG += self.criterionVGG(fake_image, real_image) * self.opt.lambda_feat
|
||||
loss_GB_VGG += self.criterionVGG(blend_image, real_image) * self.opt.lambda_feat
|
||||
|
||||
# Only return the fake_B image if necessary to save BW
|
||||
return [ self.loss_filter( loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_GB_GAN, loss_GB_GAN_Feat, loss_GB_VGG, loss_D_real, loss_D_fake, loss_D_blend ), None if not infer else fake_inter_1, fake_image, fake_inter_2, blend_image, alpha, real_image, inter_label_1, input_label, inter_label_2 ]
|
||||
|
||||
def inference(self, label, label_ref, image_ref,cFlag):
|
||||
|
||||
# Encode Inputs
|
||||
image_ref = Variable(image_ref)
|
||||
input_label, input_label_ref, real_image_ref = self.encode_input_test(Variable(label), Variable(label_ref), image_ref, infer=True,f=cFlag)
|
||||
|
||||
if torch.__version__.startswith('0.4'):
|
||||
with torch.no_grad():
|
||||
fake_image = self.netG.forward(input_label, input_label_ref, real_image_ref)
|
||||
else:
|
||||
fake_image = self.netG.forward(input_label, input_label_ref, real_image_ref)
|
||||
return fake_image
|
||||
|
||||
def save(self, which_epoch):
|
||||
self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
|
||||
self.save_network(self.netD, 'D', which_epoch, self.gpu_ids)
|
||||
self.save_network(self.netB, 'B', which_epoch, self.gpu_ids)
|
||||
|
||||
def update_fixed_params(self):
|
||||
# after fixing the global generator for a number of iterations, also start finetuning it
|
||||
params = list(self.netG.parameters())
|
||||
if self.gen_features:
|
||||
params += list(self.netE.parameters())
|
||||
self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
|
||||
if self.opt.verbose:
|
||||
print('------------ Now also finetuning global generator -----------')
|
||||
|
||||
def update_learning_rate(self):
|
||||
lrd = self.opt.lr / self.opt.niter_decay
|
||||
lr = self.old_lr - lrd
|
||||
for param_group in self.optimizer_D.param_groups:
|
||||
param_group['lr'] = lr
|
||||
for param_group in self.optimizer_G.param_groups:
|
||||
param_group['lr'] = lr
|
||||
if self.opt.verbose:
|
||||
print('update learning rate: %f -> %f' % (self.old_lr, lr))
|
||||
self.old_lr = lr
|
||||
|
||||
class InferenceModel(Pix2PixHDModel):
|
||||
def forward(self, inp):
|
||||
label = inp
|
||||
return self.inference(label)
|
||||
|
||||
|
||||
@ -1,89 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
import argparse
|
||||
import os
|
||||
from util import util
|
||||
import torch
|
||||
|
||||
class BaseOptions():
|
||||
def __init__(self):
|
||||
self.parser = argparse.ArgumentParser()
|
||||
self.initialized = False
|
||||
|
||||
def initialize(self):
|
||||
# experiment specifics
|
||||
self.parser.add_argument('--name', type=str, default='label2face_512p', help='name of the experiment. It decides where to store samples and models')
|
||||
self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
|
||||
self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
|
||||
self.parser.add_argument('--model', type=str, default='pix2pixHD', help='which model to use')
|
||||
self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
|
||||
self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator')
|
||||
self.parser.add_argument('--data_type', default=32, type=int, choices=[8, 16, 32], help="Supported data type i.e. 8, 16, 32 bit")
|
||||
self.parser.add_argument('--verbose', action='store_true', default=False, help='toggles verbose')
|
||||
|
||||
# input/output sizes
|
||||
self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
|
||||
self.parser.add_argument('--loadSize', type=int, default=512, help='scale images to this size')
|
||||
self.parser.add_argument('--fineSize', type=int, default=512, help='then crop to this size')
|
||||
self.parser.add_argument('--label_nc', type=int, default=19, help='# of input label channels')
|
||||
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')
|
||||
|
||||
# for setting inputs
|
||||
self.parser.add_argument('--dataroot', type=str, default='../Data_preprocessing/')
|
||||
self.parser.add_argument('--resize_or_crop', type=str, default='scale_width', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
|
||||
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('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')
|
||||
self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data')
|
||||
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.')
|
||||
|
||||
# for displays
|
||||
self.parser.add_argument('--display_winsize', type=int, default=512, help='display window size')
|
||||
self.parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed')
|
||||
|
||||
# for generator
|
||||
self.parser.add_argument('--netG', type=str, default='global', help='selects model to use for netG')
|
||||
self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
|
||||
self.parser.add_argument('--n_downsample_global', type=int, default=4, help='number of downsampling layers in netG')
|
||||
self.parser.add_argument('--n_blocks_global', type=int, default=4, help='number of residual blocks in the global generator network')
|
||||
self.parser.add_argument('--n_blocks_local', type=int, default=3, help='number of residual blocks in the local enhancer network')
|
||||
self.parser.add_argument('--n_local_enhancers', type=int, default=1, help='number of local enhancers to use')
|
||||
self.parser.add_argument('--niter_fix_global', type=int, default=0, help='number of epochs that we only train the outmost local enhancer')
|
||||
|
||||
self.initialized = True
|
||||
|
||||
def parse(self, save=True):
|
||||
if not self.initialized:
|
||||
self.initialize()
|
||||
self.opt = self.parser.parse_args()
|
||||
self.opt.isTrain = self.isTrain # train or test
|
||||
|
||||
str_ids = self.opt.gpu_ids.split(',')
|
||||
self.opt.gpu_ids = []
|
||||
for str_id in str_ids:
|
||||
id = int(str_id)
|
||||
if id >= 0:
|
||||
self.opt.gpu_ids.append(id)
|
||||
|
||||
# set gpu ids
|
||||
# if len(self.opt.gpu_ids) > 0:
|
||||
# torch.cuda.set_device(self.opt.gpu_ids[0])
|
||||
|
||||
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)
|
||||
if save and not self.opt.continue_train:
|
||||
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
|
||||
@ -1,19 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
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=1000, help='how many test images to run')
|
||||
self.parser.add_argument('--cluster_path', type=str, default='features_clustered_010.npy', help='the path for clustered results of encoded features')
|
||||
self.parser.add_argument('--use_encoded_image', action='store_true', help='if specified, encode the real image to get the feature map')
|
||||
self.parser.add_argument("--export_onnx", type=str, help="export ONNX model to a given file")
|
||||
self.parser.add_argument("--engine", type=str, help="run serialized TRT engine")
|
||||
self.parser.add_argument("--onnx", type=str, help="run ONNX model via TRT")
|
||||
self.isTrain = False
|
||||
@ -1,36 +0,0 @@
|
||||
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
|
||||
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
||||
from .base_options import BaseOptions
|
||||
|
||||
class TrainOptions(BaseOptions):
|
||||
def initialize(self):
|
||||
BaseOptions.initialize(self)
|
||||
# for displays
|
||||
self.parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen')
|
||||
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=1000, help='frequency of saving the latest results')
|
||||
self.parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving checkpoints at the end of epochs')
|
||||
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('--debug', action='store_true', help='only do one epoch and displays at each iteration')
|
||||
|
||||
# for training
|
||||
self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
|
||||
self.parser.add_argument('--load_pretrain', type=str, default='./checkpoints/label2face_512p', help='load the pretrained model from the specified location')
|
||||
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('--phase', type=str, default='train', help='train, val, test, etc')
|
||||
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.00005, help='initial learning rate for adam')
|
||||
|
||||
# for discriminators
|
||||
self.parser.add_argument('--num_D', type=int, default=2, help='number of discriminators to use')
|
||||
self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers')
|
||||
self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
|
||||
self.parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')
|
||||
self.parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')
|
||||
self.parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss')
|
||||
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('--pool_size', type=int, default=0, help='the size of image buffer that stores previously generated images')
|
||||
|
||||
self.isTrain = True
|
||||
@ -1,31 +0,0 @@
|
||||
import random
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
class ImagePool():
|
||||
def __init__(self, pool_size):
|
||||
self.pool_size = pool_size
|
||||
if self.pool_size > 0:
|
||||
self.num_imgs = 0
|
||||
self.images = []
|
||||
|
||||
def query(self, images):
|
||||
if self.pool_size == 0:
|
||||
return images
|
||||
return_images = []
|
||||
for image in images.data:
|
||||
image = torch.unsqueeze(image, 0)
|
||||
if self.num_imgs < self.pool_size:
|
||||
self.num_imgs = self.num_imgs + 1
|
||||
self.images.append(image)
|
||||
return_images.append(image)
|
||||
else:
|
||||
p = random.uniform(0, 1)
|
||||
if p > 0.5:
|
||||
random_id = random.randint(0, self.pool_size-1)
|
||||
tmp = self.images[random_id].clone()
|
||||
self.images[random_id] = image
|
||||
return_images.append(tmp)
|
||||
else:
|
||||
return_images.append(image)
|
||||
return_images = Variable(torch.cat(return_images, 0))
|
||||
return return_images
|
||||
@ -1,107 +0,0 @@
|
||||
from __future__ import print_function
|
||||
|
||||
print ('?')
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
# import numpy as np
|
||||
import os
|
||||
|
||||
# Converts a Tensor into a Numpy array
|
||||
# |imtype|: the desired type of the converted numpy array
|
||||
def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
|
||||
if isinstance(image_tensor, list):
|
||||
image_numpy = []
|
||||
for i in range(len(image_tensor)):
|
||||
image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
|
||||
return image_numpy
|
||||
image_numpy = image_tensor.cpu().float().numpy()
|
||||
#if normalize:
|
||||
# image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
|
||||
#else:
|
||||
# image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
|
||||
image_numpy = (image_numpy + 1) / 2.0
|
||||
image_numpy = np.clip(image_numpy, 0, 1)
|
||||
if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:
|
||||
image_numpy = image_numpy[:,:,0]
|
||||
|
||||
return image_numpy
|
||||
|
||||
# Converts a one-hot tensor into a colorful label map
|
||||
def tensor2label(label_tensor, n_label, imtype=np.uint8):
|
||||
if n_label == 0:
|
||||
return tensor2im(label_tensor, imtype)
|
||||
label_tensor = label_tensor.cpu().float()
|
||||
if label_tensor.size()[0] > 1:
|
||||
label_tensor = label_tensor.max(0, keepdim=True)[1]
|
||||
label_tensor = Colorize(n_label)(label_tensor)
|
||||
#label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
|
||||
label_numpy = label_tensor.numpy()
|
||||
label_numpy = label_numpy / 255.0
|
||||
|
||||
return label_numpy
|
||||
|
||||
def save_image(image_numpy, image_path):
|
||||
image_pil = Image.fromarray(image_numpy)
|
||||
image_pil.save(image_path)
|
||||
|
||||
def mkdirs(paths):
|
||||
if isinstance(paths, list) and not isinstance(paths, str):
|
||||
for path in paths:
|
||||
mkdir(path)
|
||||
else:
|
||||
mkdir(paths)
|
||||
|
||||
def mkdir(path):
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
###############################################################################
|
||||
# Code from
|
||||
# https://github.com/ycszen/pytorch-seg/blob/master/transform.py
|
||||
# Modified so it complies with the Citscape label map colors
|
||||
###############################################################################
|
||||
def uint82bin(n, count=8):
|
||||
"""returns the binary of integer n, count refers to amount of bits"""
|
||||
return ''.join([str((n >> y) & 1) for y in range(count-1, -1, -1)])
|
||||
|
||||
def labelcolormap(N):
|
||||
if N == 35: # cityscape
|
||||
cmap = np.array([( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), (111, 74, 0), ( 81, 0, 81),
|
||||
(128, 64,128), (244, 35,232), (250,170,160), (230,150,140), ( 70, 70, 70), (102,102,156), (190,153,153),
|
||||
(180,165,180), (150,100,100), (150,120, 90), (153,153,153), (153,153,153), (250,170, 30), (220,220, 0),
|
||||
(107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255, 0, 0), ( 0, 0,142), ( 0, 0, 70),
|
||||
( 0, 60,100), ( 0, 0, 90), ( 0, 0,110), ( 0, 80,100), ( 0, 0,230), (119, 11, 32), ( 0, 0,142)],
|
||||
dtype=np.uint8)
|
||||
else:
|
||||
cmap = np.zeros((N, 3), dtype=np.uint8)
|
||||
for i in range(N):
|
||||
r, g, b = 0, 0, 0
|
||||
id = i
|
||||
for j in range(7):
|
||||
str_id = uint82bin(id)
|
||||
r = r ^ (np.uint8(str_id[-1]) << (7-j))
|
||||
g = g ^ (np.uint8(str_id[-2]) << (7-j))
|
||||
b = b ^ (np.uint8(str_id[-3]) << (7-j))
|
||||
id = id >> 3
|
||||
cmap[i, 0] = r
|
||||
cmap[i, 1] = g
|
||||
cmap[i, 2] = b
|
||||
return cmap
|
||||
|
||||
class Colorize(object):
|
||||
def __init__(self, n=35):
|
||||
self.cmap = labelcolormap(n)
|
||||
self.cmap = torch.from_numpy(self.cmap[:n])
|
||||
|
||||
def __call__(self, gray_image):
|
||||
size = gray_image.size()
|
||||
color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
|
||||
|
||||
for label in range(0, len(self.cmap)):
|
||||
mask = (label == gray_image[0]).cpu()
|
||||
color_image[0][mask] = self.cmap[label][0]
|
||||
color_image[1][mask] = self.cmap[label][1]
|
||||
color_image[2][mask] = self.cmap[label][2]
|
||||
|
||||
return color_image
|
||||
@ -1,8 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/.idea/EnlightenGAN.iml" filepath="$PROJECT_DIR$/.idea/EnlightenGAN.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
@ -1,6 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="$PROJECT_DIR$/../.." vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2020 hzwer
|
||||
|
||||
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.
|
||||
@ -1,117 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from model.warplayer import warp
|
||||
|
||||
|
||||
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(out_planes),
|
||||
)
|
||||
|
||||
|
||||
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(out_planes),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes, stride=1):
|
||||
super(ResBlock, self).__init__()
|
||||
if in_planes == out_planes and stride == 1:
|
||||
self.conv0 = nn.Identity()
|
||||
else:
|
||||
self.conv0 = nn.Conv2d(in_planes, out_planes,
|
||||
3, stride, 1, bias=False)
|
||||
self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
|
||||
self.conv2 = conv_wo_act(out_planes, out_planes, 3, 1, 1)
|
||||
self.relu1 = nn.PReLU(1)
|
||||
self.relu2 = nn.PReLU(out_planes)
|
||||
self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
|
||||
self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.conv0(x)
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
w = x.mean(3, True).mean(2, True)
|
||||
w = self.relu1(self.fc1(w))
|
||||
w = torch.sigmoid(self.fc2(w))
|
||||
x = self.relu2(x * w + y)
|
||||
return x
|
||||
|
||||
|
||||
class IFBlock(nn.Module):
|
||||
def __init__(self, in_planes, scale=1, c=64):
|
||||
super(IFBlock, self).__init__()
|
||||
self.scale = scale
|
||||
self.conv0 = conv(in_planes, c, 3, 2, 1)
|
||||
self.res0 = ResBlock(c, c)
|
||||
self.res1 = ResBlock(c, c)
|
||||
self.res2 = ResBlock(c, c)
|
||||
self.res3 = ResBlock(c, c)
|
||||
self.res4 = ResBlock(c, c)
|
||||
self.res5 = ResBlock(c, c)
|
||||
self.conv1 = nn.Conv2d(c, 8, 3, 1, 1)
|
||||
self.up = nn.PixelShuffle(2)
|
||||
|
||||
def forward(self, x):
|
||||
if self.scale != 1:
|
||||
x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear",
|
||||
align_corners=False)
|
||||
x = self.conv0(x)
|
||||
x = self.res0(x)
|
||||
x = self.res1(x)
|
||||
x = self.res2(x)
|
||||
x = self.res3(x)
|
||||
x = self.res4(x)
|
||||
x = self.res5(x)
|
||||
x = self.conv1(x)
|
||||
flow = self.up(x)
|
||||
if self.scale != 1:
|
||||
flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear",
|
||||
align_corners=False)
|
||||
return flow
|
||||
|
||||
|
||||
class IFNet(nn.Module):
|
||||
def __init__(self, cFlag):
|
||||
super(IFNet, self).__init__()
|
||||
self.block0 = IFBlock(6, scale=4, c=192)
|
||||
self.block1 = IFBlock(8, scale=2, c=128)
|
||||
self.block2 = IFBlock(8, scale=1, c=64)
|
||||
self.cFlag = cFlag
|
||||
|
||||
def forward(self, x):
|
||||
x = F.interpolate(x, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False)
|
||||
flow0 = self.block0(x)
|
||||
F1 = flow0
|
||||
warped_img0 = warp(x[:, :3], F1, self.cFlag)
|
||||
warped_img1 = warp(x[:, 3:], -F1, self.cFlag)
|
||||
flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1), 1))
|
||||
F2 = (flow0 + flow1)
|
||||
warped_img0 = warp(x[:, :3], F2, self.cFlag)
|
||||
warped_img1 = warp(x[:, 3:], -F2, self.cFlag)
|
||||
flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2), 1))
|
||||
F3 = (flow0 + flow1 + flow2)
|
||||
return F3, [F1, F2, F3]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
||||
img1 = torch.tensor(np.random.normal(
|
||||
0, 1, (3, 3, 256, 256))).float().to(device)
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
flownet = IFNet()
|
||||
flow, _ = flownet(imgs)
|
||||
print(flow.shape)
|
||||
@ -1,115 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from model.warplayer import warp
|
||||
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(out_planes),
|
||||
)
|
||||
|
||||
|
||||
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(out_planes),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes, stride=1):
|
||||
super(ResBlock, self).__init__()
|
||||
if in_planes == out_planes and stride == 1:
|
||||
self.conv0 = nn.Identity()
|
||||
else:
|
||||
self.conv0 = nn.Conv2d(in_planes, out_planes,
|
||||
3, stride, 1, bias=False)
|
||||
self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
|
||||
self.conv2 = conv_wo_act(out_planes, out_planes, 3, 1, 1)
|
||||
self.relu1 = nn.PReLU(1)
|
||||
self.relu2 = nn.PReLU(out_planes)
|
||||
self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
|
||||
self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.conv0(x)
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
w = x.mean(3, True).mean(2, True)
|
||||
w = self.relu1(self.fc1(w))
|
||||
w = torch.sigmoid(self.fc2(w))
|
||||
x = self.relu2(x * w + y)
|
||||
return x
|
||||
|
||||
|
||||
class IFBlock(nn.Module):
|
||||
def __init__(self, in_planes, scale=1, c=64):
|
||||
super(IFBlock, self).__init__()
|
||||
self.scale = scale
|
||||
self.conv0 = conv(in_planes, c, 3, 1, 1)
|
||||
self.res0 = ResBlock(c, c)
|
||||
self.res1 = ResBlock(c, c)
|
||||
self.res2 = ResBlock(c, c)
|
||||
self.res3 = ResBlock(c, c)
|
||||
self.res4 = ResBlock(c, c)
|
||||
self.res5 = ResBlock(c, c)
|
||||
self.conv1 = nn.Conv2d(c, 2, 3, 1, 1)
|
||||
self.up = nn.PixelShuffle(2)
|
||||
|
||||
def forward(self, x):
|
||||
if self.scale != 1:
|
||||
x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear",
|
||||
align_corners=False)
|
||||
x = self.conv0(x)
|
||||
x = self.res0(x)
|
||||
x = self.res1(x)
|
||||
x = self.res2(x)
|
||||
x = self.res3(x)
|
||||
x = self.res4(x)
|
||||
x = self.res5(x)
|
||||
x = self.conv1(x)
|
||||
flow = x # self.up(x)
|
||||
if self.scale != 1:
|
||||
flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear",
|
||||
align_corners=False)
|
||||
return flow
|
||||
|
||||
|
||||
class IFNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(IFNet, self).__init__()
|
||||
self.block0 = IFBlock(6, scale=4, c=192)
|
||||
self.block1 = IFBlock(8, scale=2, c=128)
|
||||
self.block2 = IFBlock(8, scale=1, c=64)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.interpolate(x, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False)
|
||||
flow0 = self.block0(x)
|
||||
F1 = flow0
|
||||
warped_img0 = warp(x[:, :3], F1)
|
||||
warped_img1 = warp(x[:, 3:], -F1)
|
||||
flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1), 1))
|
||||
F2 = (flow0 + flow1)
|
||||
warped_img0 = warp(x[:, :3], F2)
|
||||
warped_img1 = warp(x[:, 3:], -F2)
|
||||
flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2), 1))
|
||||
F3 = (flow0 + flow1 + flow2)
|
||||
return F3, [F1, F2, F3]
|
||||
|
||||
if __name__ == '__main__':
|
||||
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
||||
img1 = torch.tensor(np.random.normal(
|
||||
0, 1, (3, 3, 256, 256))).float().to(device)
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
flownet = IFNet()
|
||||
flow, _ = flownet(imgs)
|
||||
print(flow.shape)
|
||||
@ -1,262 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
from torch.optim import AdamW
|
||||
import torch.optim as optim
|
||||
import itertools
|
||||
from model.warplayer import warp
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from model.IFNet import *
|
||||
import torch.nn.functional as F
|
||||
from model.loss import *
|
||||
|
||||
|
||||
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=True),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
|
||||
def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
|
||||
return nn.Sequential(
|
||||
torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes,
|
||||
kernel_size=4, stride=2, padding=1, bias=True),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
|
||||
def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=True),
|
||||
)
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes, stride=2):
|
||||
super(ResBlock, self).__init__()
|
||||
if in_planes == out_planes and stride == 1:
|
||||
self.conv0 = nn.Identity()
|
||||
else:
|
||||
self.conv0 = nn.Conv2d(in_planes, out_planes,
|
||||
3, stride, 1, bias=False)
|
||||
self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
|
||||
self.conv2 = conv_woact(out_planes, out_planes, 3, 1, 1)
|
||||
self.relu1 = nn.PReLU(1)
|
||||
self.relu2 = nn.PReLU(out_planes)
|
||||
self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
|
||||
self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.conv0(x)
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
w = x.mean(3, True).mean(2, True)
|
||||
w = self.relu1(self.fc1(w))
|
||||
w = torch.sigmoid(self.fc2(w))
|
||||
x = self.relu2(x * w + y)
|
||||
return x
|
||||
|
||||
|
||||
c = 16
|
||||
|
||||
|
||||
class ContextNet(nn.Module):
|
||||
def __init__(self, cFlag):
|
||||
super(ContextNet, self).__init__()
|
||||
self.conv1 = ResBlock(3, c)
|
||||
self.conv2 = ResBlock(c, 2 * c)
|
||||
self.conv3 = ResBlock(2 * c, 4 * c)
|
||||
self.conv4 = ResBlock(4 * c, 8 * c)
|
||||
self.cFlag = cFlag
|
||||
|
||||
def forward(self, x, flow):
|
||||
x = self.conv1(x)
|
||||
f1 = warp(x, flow, self.cFlag)
|
||||
x = self.conv2(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f2 = warp(x, flow, self.cFlag)
|
||||
x = self.conv3(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f3 = warp(x, flow, self.cFlag)
|
||||
x = self.conv4(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f4 = warp(x, flow, self.cFlag)
|
||||
return [f1, f2, f3, f4]
|
||||
|
||||
|
||||
class FusionNet(nn.Module):
|
||||
def __init__(self, cFlag):
|
||||
super(FusionNet, self).__init__()
|
||||
self.down0 = ResBlock(8, 2 * c)
|
||||
self.down1 = ResBlock(4 * c, 4 * c)
|
||||
self.down2 = ResBlock(8 * c, 8 * c)
|
||||
self.down3 = ResBlock(16 * c, 16 * c)
|
||||
self.up0 = deconv(32 * c, 8 * c)
|
||||
self.up1 = deconv(16 * c, 4 * c)
|
||||
self.up2 = deconv(8 * c, 2 * c)
|
||||
self.up3 = deconv(4 * c, c)
|
||||
self.conv = nn.Conv2d(c, 4, 3, 1, 1)
|
||||
self.cFlag = cFlag
|
||||
|
||||
def forward(self, img0, img1, flow, c0, c1, flow_gt):
|
||||
warped_img0 = warp(img0, flow, self.cFlag)
|
||||
warped_img1 = warp(img1, -flow, self.cFlag)
|
||||
if flow_gt == None:
|
||||
warped_img0_gt, warped_img1_gt = None, None
|
||||
else:
|
||||
warped_img0_gt = warp(img0, flow_gt[:, :2])
|
||||
warped_img1_gt = warp(img1, flow_gt[:, 2:4])
|
||||
s0 = self.down0(torch.cat((warped_img0, warped_img1, flow), 1))
|
||||
s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
|
||||
s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
|
||||
s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
|
||||
x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
|
||||
x = self.up1(torch.cat((x, s2), 1))
|
||||
x = self.up2(torch.cat((x, s1), 1))
|
||||
x = self.up3(torch.cat((x, s0), 1))
|
||||
x = self.conv(x)
|
||||
return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
|
||||
|
||||
|
||||
class Model:
|
||||
def __init__(self, c_flag, local_rank=-1):
|
||||
self.flownet = IFNet(c_flag)
|
||||
self.contextnet = ContextNet(c_flag)
|
||||
self.fusionnet = FusionNet(c_flag)
|
||||
self.device(c_flag)
|
||||
self.optimG = AdamW(itertools.chain(
|
||||
self.flownet.parameters(),
|
||||
self.contextnet.parameters(),
|
||||
self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-5)
|
||||
self.schedulerG = optim.lr_scheduler.CyclicLR(
|
||||
self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False)
|
||||
self.epe = EPE()
|
||||
self.ter = Ternary(c_flag)
|
||||
self.sobel = SOBEL(c_flag)
|
||||
if local_rank != -1:
|
||||
self.flownet = DDP(self.flownet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
self.contextnet = DDP(self.contextnet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
self.fusionnet = DDP(self.fusionnet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
|
||||
def train(self):
|
||||
self.flownet.train()
|
||||
self.contextnet.train()
|
||||
self.fusionnet.train()
|
||||
|
||||
def eval(self):
|
||||
self.flownet.eval()
|
||||
self.contextnet.eval()
|
||||
self.fusionnet.eval()
|
||||
|
||||
def device(self, c_flag):
|
||||
if torch.cuda.is_available() and not c_flag:
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
self.flownet.to(device)
|
||||
self.contextnet.to(device)
|
||||
self.fusionnet.to(device)
|
||||
|
||||
def load_model(self, path, rank=0):
|
||||
def convert(param):
|
||||
return {
|
||||
k.replace("module.", ""): v
|
||||
for k, v in param.items()
|
||||
if "module." in k
|
||||
}
|
||||
|
||||
if rank == 0:
|
||||
self.flownet.load_state_dict(
|
||||
convert(torch.load('{}/flownet.pkl'.format(path), map_location=torch.device("cpu"))))
|
||||
self.contextnet.load_state_dict(
|
||||
convert(torch.load('{}/contextnet.pkl'.format(path), map_location=torch.device("cpu"))))
|
||||
self.fusionnet.load_state_dict(
|
||||
convert(torch.load('{}/unet.pkl'.format(path), map_location=torch.device("cpu"))))
|
||||
|
||||
def save_model(self, path, rank=0):
|
||||
if rank == 0:
|
||||
torch.save(self.flownet.state_dict(),
|
||||
'{}/flownet.pkl'.format(path))
|
||||
torch.save(self.contextnet.state_dict(),
|
||||
'{}/contextnet.pkl'.format(path))
|
||||
torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
|
||||
|
||||
def predict(self, imgs, flow, training=True, flow_gt=None):
|
||||
img0 = imgs[:, :3]
|
||||
img1 = imgs[:, 3:]
|
||||
c0 = self.contextnet(img0, flow)
|
||||
c1 = self.contextnet(img1, -flow)
|
||||
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
|
||||
align_corners=False) * 2.0
|
||||
refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet(
|
||||
img0, img1, flow, c0, c1, flow_gt)
|
||||
res = torch.sigmoid(refine_output[:, :3]) * 2 - 1
|
||||
mask = torch.sigmoid(refine_output[:, 3:4])
|
||||
merged_img = warped_img0 * mask + warped_img1 * (1 - mask)
|
||||
pred = merged_img + res
|
||||
pred = torch.clamp(pred, 0, 1)
|
||||
if training:
|
||||
return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
|
||||
else:
|
||||
return pred
|
||||
|
||||
def inference(self, img0, img1):
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
flow, _ = self.flownet(imgs)
|
||||
return self.predict(imgs, flow, training=False).detach()
|
||||
|
||||
def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
|
||||
for param_group in self.optimG.param_groups:
|
||||
param_group['lr'] = learning_rate
|
||||
if training:
|
||||
self.train()
|
||||
else:
|
||||
self.eval()
|
||||
flow, flow_list = self.flownet(imgs)
|
||||
pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict(
|
||||
imgs, flow, flow_gt=flow_gt)
|
||||
loss_ter = self.ter(pred, gt).mean()
|
||||
if training:
|
||||
with torch.no_grad():
|
||||
loss_flow = torch.abs(warped_img0_gt - gt).mean()
|
||||
loss_mask = torch.abs(
|
||||
merged_img - gt).sum(1, True).float().detach()
|
||||
loss_mask = F.interpolate(loss_mask, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False).detach()
|
||||
flow_gt = (F.interpolate(flow_gt, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5).detach()
|
||||
loss_cons = 0
|
||||
for i in range(3):
|
||||
loss_cons += self.epe(flow_list[i], flow_gt[:, :2], 1)
|
||||
loss_cons += self.epe(-flow_list[i], flow_gt[:, 2:4], 1)
|
||||
loss_cons = loss_cons.mean() * 0.01
|
||||
else:
|
||||
loss_cons = torch.tensor([0])
|
||||
loss_flow = torch.abs(warped_img0 - gt).mean()
|
||||
loss_mask = 1
|
||||
loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean()
|
||||
if training:
|
||||
self.optimG.zero_grad()
|
||||
loss_G = loss_l1 + loss_cons + loss_ter
|
||||
loss_G.backward()
|
||||
self.optimG.step()
|
||||
return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
||||
img1 = torch.tensor(np.random.normal(
|
||||
0, 1, (3, 3, 256, 256))).float().to(device)
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
model = Model()
|
||||
model.eval()
|
||||
print(model.inference(imgs).shape)
|
||||
@ -1,250 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
from torch.optim import AdamW
|
||||
import torch.optim as optim
|
||||
import itertools
|
||||
from model.warplayer import warp
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from model.IFNet2F import *
|
||||
import torch.nn.functional as F
|
||||
from model.loss import *
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=True),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
|
||||
def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
|
||||
return nn.Sequential(
|
||||
torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes,
|
||||
kernel_size=4, stride=2, padding=1, bias=True),
|
||||
nn.PReLU(out_planes)
|
||||
)
|
||||
|
||||
def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, bias=True),
|
||||
)
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes, stride=2):
|
||||
super(ResBlock, self).__init__()
|
||||
if in_planes == out_planes and stride == 1:
|
||||
self.conv0 = nn.Identity()
|
||||
else:
|
||||
self.conv0 = nn.Conv2d(in_planes, out_planes,
|
||||
3, stride, 1, bias=False)
|
||||
self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
|
||||
self.conv2 = conv_woact(out_planes, out_planes, 3, 1, 1)
|
||||
self.relu1 = nn.PReLU(1)
|
||||
self.relu2 = nn.PReLU(out_planes)
|
||||
self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
|
||||
self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.conv0(x)
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
w = x.mean(3, True).mean(2, True)
|
||||
w = self.relu1(self.fc1(w))
|
||||
w = torch.sigmoid(self.fc2(w))
|
||||
x = self.relu2(x * w + y)
|
||||
return x
|
||||
|
||||
c = 16
|
||||
|
||||
class ContextNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(ContextNet, self).__init__()
|
||||
self.conv1 = ResBlock(3, c, 1)
|
||||
self.conv2 = ResBlock(c, 2*c)
|
||||
self.conv3 = ResBlock(2*c, 4*c)
|
||||
self.conv4 = ResBlock(4*c, 8*c)
|
||||
|
||||
def forward(self, x, flow):
|
||||
x = self.conv1(x)
|
||||
f1 = warp(x, flow)
|
||||
x = self.conv2(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f2 = warp(x, flow)
|
||||
x = self.conv3(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f3 = warp(x, flow)
|
||||
x = self.conv4(x)
|
||||
flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
|
||||
align_corners=False) * 0.5
|
||||
f4 = warp(x, flow)
|
||||
return [f1, f2, f3, f4]
|
||||
|
||||
|
||||
class FusionNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(FusionNet, self).__init__()
|
||||
self.down0 = ResBlock(8, 2*c, 1)
|
||||
self.down1 = ResBlock(4*c, 4*c)
|
||||
self.down2 = ResBlock(8*c, 8*c)
|
||||
self.down3 = ResBlock(16*c, 16*c)
|
||||
self.up0 = deconv(32*c, 8*c)
|
||||
self.up1 = deconv(16*c, 4*c)
|
||||
self.up2 = deconv(8*c, 2*c)
|
||||
self.up3 = deconv(4*c, c)
|
||||
self.conv = nn.Conv2d(c, 4, 3, 2, 1)
|
||||
|
||||
def forward(self, img0, img1, flow, c0, c1, flow_gt):
|
||||
warped_img0 = warp(img0, flow)
|
||||
warped_img1 = warp(img1, -flow)
|
||||
if flow_gt == None:
|
||||
warped_img0_gt, warped_img1_gt = None, None
|
||||
else:
|
||||
warped_img0_gt = warp(img0, flow_gt[:, :2])
|
||||
warped_img1_gt = warp(img1, flow_gt[:, 2:4])
|
||||
s0 = self.down0(torch.cat((warped_img0, warped_img1, flow), 1))
|
||||
s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
|
||||
s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
|
||||
s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
|
||||
x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
|
||||
x = self.up1(torch.cat((x, s2), 1))
|
||||
x = self.up2(torch.cat((x, s1), 1))
|
||||
x = self.up3(torch.cat((x, s0), 1))
|
||||
x = self.conv(x)
|
||||
return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
|
||||
|
||||
|
||||
class Model:
|
||||
def __init__(self, local_rank=-1):
|
||||
self.flownet = IFNet()
|
||||
self.contextnet = ContextNet()
|
||||
self.fusionnet = FusionNet()
|
||||
self.device()
|
||||
self.optimG = AdamW(itertools.chain(
|
||||
self.flownet.parameters(),
|
||||
self.contextnet.parameters(),
|
||||
self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-5)
|
||||
self.schedulerG = optim.lr_scheduler.CyclicLR(
|
||||
self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False)
|
||||
self.epe = EPE()
|
||||
self.ter = Ternary()
|
||||
self.sobel = SOBEL()
|
||||
if local_rank != -1:
|
||||
self.flownet = DDP(self.flownet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
self.contextnet = DDP(self.contextnet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
self.fusionnet = DDP(self.fusionnet, device_ids=[
|
||||
local_rank], output_device=local_rank)
|
||||
|
||||
def train(self):
|
||||
self.flownet.train()
|
||||
self.contextnet.train()
|
||||
self.fusionnet.train()
|
||||
|
||||
def eval(self):
|
||||
self.flownet.eval()
|
||||
self.contextnet.eval()
|
||||
self.fusionnet.eval()
|
||||
|
||||
def device(self):
|
||||
self.flownet.to(device)
|
||||
self.contextnet.to(device)
|
||||
self.fusionnet.to(device)
|
||||
|
||||
def load_model(self, path, rank=0):
|
||||
def convert(param):
|
||||
return {
|
||||
k.replace("module.", ""): v
|
||||
for k, v in param.items()
|
||||
if "module." in k
|
||||
}
|
||||
if rank == 0:
|
||||
self.flownet.load_state_dict(
|
||||
convert(torch.load('{}/flownet.pkl'.format(path), map_location=device)))
|
||||
self.contextnet.load_state_dict(
|
||||
convert(torch.load('{}/contextnet.pkl'.format(path), map_location=device)))
|
||||
self.fusionnet.load_state_dict(
|
||||
convert(torch.load('{}/unet.pkl'.format(path), map_location=device)))
|
||||
|
||||
def save_model(self, path, rank=0):
|
||||
if rank == 0:
|
||||
torch.save(self.flownet.state_dict(),
|
||||
'{}/flownet.pkl'.format(path))
|
||||
torch.save(self.contextnet.state_dict(),
|
||||
'{}/contextnet.pkl'.format(path))
|
||||
torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
|
||||
|
||||
def predict(self, imgs, flow, training=True, flow_gt=None):
|
||||
img0 = imgs[:, :3]
|
||||
img1 = imgs[:, 3:]
|
||||
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
|
||||
align_corners=False) * 2.0
|
||||
c0 = self.contextnet(img0, flow)
|
||||
c1 = self.contextnet(img1, -flow)
|
||||
refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet(
|
||||
img0, img1, flow, c0, c1, flow_gt)
|
||||
res = torch.sigmoid(refine_output[:, :3]) * 2 - 1
|
||||
mask = torch.sigmoid(refine_output[:, 3:4])
|
||||
merged_img = warped_img0 * mask + warped_img1 * (1 - mask)
|
||||
pred = merged_img + res
|
||||
pred = torch.clamp(pred, 0, 1)
|
||||
if training:
|
||||
return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
|
||||
else:
|
||||
return pred
|
||||
|
||||
def inference(self, img0, img1):
|
||||
with torch.no_grad():
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
flow, _ = self.flownet(imgs)
|
||||
return self.predict(imgs, flow, training=False).detach()
|
||||
|
||||
def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
|
||||
for param_group in self.optimG.param_groups:
|
||||
param_group['lr'] = learning_rate
|
||||
if training:
|
||||
self.train()
|
||||
else:
|
||||
self.eval()
|
||||
flow, flow_list = self.flownet(imgs)
|
||||
pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict(
|
||||
imgs, flow, flow_gt=flow_gt)
|
||||
loss_ter = self.ter(pred, gt).mean()
|
||||
if training:
|
||||
with torch.no_grad():
|
||||
loss_flow = torch.abs(warped_img0_gt - gt).mean()
|
||||
loss_mask = torch.abs(
|
||||
merged_img - gt).sum(1, True).float().detach()
|
||||
loss_cons = 0
|
||||
for i in range(3):
|
||||
loss_cons += self.epe(flow_list[i], flow_gt[:, :2], 1)
|
||||
loss_cons += self.epe(-flow_list[i], flow_gt[:, 2:4], 1)
|
||||
loss_cons = loss_cons.mean() * 0.01
|
||||
else:
|
||||
loss_cons = torch.tensor([0])
|
||||
loss_flow = torch.abs(warped_img0 - gt).mean()
|
||||
loss_mask = 1
|
||||
loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean()
|
||||
if training:
|
||||
self.optimG.zero_grad()
|
||||
loss_G = loss_l1 + loss_cons + loss_ter
|
||||
loss_G.backward()
|
||||
self.optimG.step()
|
||||
return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
||||
img1 = torch.tensor(np.random.normal(
|
||||
0, 1, (3, 3, 256, 256))).float().to(device)
|
||||
imgs = torch.cat((img0, img1), 1)
|
||||
model = Model()
|
||||
model.eval()
|
||||
print(model.inference(imgs).shape)
|
||||
@ -1,90 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class EPE(nn.Module):
|
||||
def __init__(self):
|
||||
super(EPE, self).__init__()
|
||||
|
||||
def forward(self, flow, gt, loss_mask):
|
||||
loss_map = (flow - gt.detach()) ** 2
|
||||
loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
|
||||
return (loss_map * loss_mask)
|
||||
|
||||
|
||||
class Ternary(nn.Module):
|
||||
def __init__(self, cFlag):
|
||||
super(Ternary, self).__init__()
|
||||
patch_size = 7
|
||||
out_channels = patch_size * patch_size
|
||||
self.w = np.eye(out_channels).reshape(
|
||||
(patch_size, patch_size, 1, out_channels))
|
||||
self.w = np.transpose(self.w, (3, 2, 0, 1))
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() and not cFlag else "cpu")
|
||||
self.w = torch.tensor(self.w).float().to(self.device)
|
||||
|
||||
def transform(self, img):
|
||||
patches = F.conv2d(img, self.w, padding=3, bias=None)
|
||||
transf = patches - img
|
||||
transf_norm = transf / torch.sqrt(0.81 + transf**2)
|
||||
return transf_norm
|
||||
|
||||
def rgb2gray(self, rgb):
|
||||
r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
|
||||
gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
|
||||
return gray
|
||||
|
||||
def hamming(self, t1, t2):
|
||||
dist = (t1 - t2) ** 2
|
||||
dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
|
||||
return dist_norm
|
||||
|
||||
def valid_mask(self, t, padding):
|
||||
n, _, h, w = t.size()
|
||||
inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
|
||||
mask = F.pad(inner, [padding] * 4)
|
||||
return mask
|
||||
|
||||
def forward(self, img0, img1):
|
||||
img0 = self.transform(self.rgb2gray(img0))
|
||||
img1 = self.transform(self.rgb2gray(img1))
|
||||
return self.hamming(img0, img1) * self.valid_mask(img0, 1)
|
||||
|
||||
|
||||
class SOBEL(nn.Module):
|
||||
def __init__(self, cFlag):
|
||||
super(SOBEL, self).__init__()
|
||||
self.kernelX = torch.tensor([
|
||||
[1, 0, -1],
|
||||
[2, 0, -2],
|
||||
[1, 0, -1],
|
||||
]).float()
|
||||
self.kernelY = self.kernelX.clone().T
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() and not cFlag else "cpu")
|
||||
self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(self.device)
|
||||
self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(self.device)
|
||||
|
||||
def forward(self, pred, gt):
|
||||
N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
|
||||
img_stack = torch.cat(
|
||||
[pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0)
|
||||
sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
|
||||
sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
|
||||
pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:]
|
||||
pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:]
|
||||
|
||||
L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y)
|
||||
loss = (L1X+L1Y)
|
||||
return loss
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
||||
img1 = torch.tensor(np.random.normal(
|
||||
0, 1, (3, 3, 256, 256))).float().to(device)
|
||||
ternary_loss = Ternary()
|
||||
print(ternary_loss(img0, img1).shape)
|
||||
@ -1,23 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
backwarp_tenGrid = {}
|
||||
|
||||
|
||||
def warp(tenInput, tenFlow, cFlag):
|
||||
device = torch.device("cuda" if torch.cuda.is_available() and not cFlag else "cpu")
|
||||
k = (str(tenFlow.device), str(tenFlow.size()))
|
||||
if k not in backwarp_tenGrid:
|
||||
tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3]).view(
|
||||
1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
|
||||
tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2]).view(
|
||||
1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
|
||||
backwarp_tenGrid[k] = torch.cat(
|
||||
[tenHorizontal, tenVertical], 1).to(device)
|
||||
|
||||
tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
|
||||
tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1)
|
||||
|
||||
g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
|
||||
return torch.nn.functional.grid_sample(input=tenInput, grid=torch.clamp(g, -1, 1), mode='bilinear',
|
||||
padding_mode='zeros', align_corners=True)
|
||||
@ -1,58 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
from gimpfu import *
|
||||
import sys
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools'])
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, image.active_layer.type, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def genNewImg(name, layer_np):
|
||||
h, w, d = layer_np.shape
|
||||
img = pdb.gimp_image_new(w, h, RGB)
|
||||
display = pdb.gimp_display_new(img)
|
||||
|
||||
rlBytes = np.uint8(layer_np).tobytes();
|
||||
rl = gimp.Layer(img, name, img.width, img.height, RGB, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
|
||||
pdb.gimp_image_insert_layer(img, rl, None, 0)
|
||||
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def colorpalette():
|
||||
cpy = cv2.cvtColor(cv2.imread(baseLoc+'color_palette.png'),cv2.COLOR_BGR2RGB)
|
||||
genNewImg('palette', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"colorpalette",
|
||||
"colorpalette",
|
||||
"colorpalette.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"colorpalette...",
|
||||
"", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[],
|
||||
[],
|
||||
colorpalette, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,75 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'DeblurGANv2'])
|
||||
|
||||
import cv2
|
||||
from predictorClass import Predictor
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def getdeblur(img, flag):
|
||||
predictor = Predictor(weights_path=baseLoc + 'weights/deblur/' + 'best_fpn.h5', cf=flag)
|
||||
if img.shape[2] == 4: # get rid of alpha channel
|
||||
img = img[:, :, 0:3]
|
||||
pred = predictor(img, None, cf=flag)
|
||||
return pred
|
||||
|
||||
|
||||
def deblur(img, layer, flag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not flag:
|
||||
gimp.progress_init("(Using GPU) Deblurring " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Deblurring " + layer.name + "...")
|
||||
|
||||
imgmat = channelData(layer)
|
||||
pred = getdeblur(imgmat, flag)
|
||||
createResultLayer(img, 'deblur_' + layer.name, pred)
|
||||
|
||||
|
||||
register(
|
||||
"deblur",
|
||||
"deblur",
|
||||
"Running deblurring.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"deblur...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
# (PF_LAYER, "drawinglayer", "Original Image", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
|
||||
],
|
||||
[],
|
||||
deblur, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,86 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
from gimpfu import *
|
||||
import sys
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools',baseLoc+'ideepcolor'])
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
from data import colorize_image as CI
|
||||
|
||||
def createResultLayer(image,name,result):
|
||||
rlBytes=np.uint8(result).tobytes();
|
||||
rl=gimp.Layer(image,name,image.width,image.height)
|
||||
# ,image.active_layer.type,100,NORMAL_MODE)
|
||||
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
||||
region[:,:]=rlBytes
|
||||
image.add_layer(rl,0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def channelData(layer):#convert gimp image to numpy
|
||||
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
||||
pixChars=region[:,:] # Take whole layer
|
||||
bpp=region.bpp
|
||||
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
||||
|
||||
|
||||
def deepcolor(tmp1, tmp2, ilayerimg,ilayerc,cflag) :
|
||||
layerimg = channelData(ilayerimg)
|
||||
layerc = channelData(ilayerc)
|
||||
|
||||
if ilayerimg.name == ilayerc.name: # if local color hints are not provided by user
|
||||
mask = np.zeros((1, 256, 256)) # giving no user points, so mask is all 0's
|
||||
input_ab = np.zeros((2, 256, 256)) # ab values of user points, default to 0 for no input
|
||||
else:
|
||||
if layerc.shape[2] == 3: # error
|
||||
pdb.gimp_message("Alpha channel missing in " + ilayerc.name + " !")
|
||||
return
|
||||
else:
|
||||
input_ab = cv2.cvtColor(layerc[:,:,0:3].astype(np.float32)/255, cv2.COLOR_RGB2LAB)
|
||||
mask = layerc[:,:,3]>0
|
||||
mask = mask.astype(np.uint8)
|
||||
input_ab = cv2.resize(input_ab,(256,256))
|
||||
mask = cv2.resize(mask, (256, 256))
|
||||
mask = mask[np.newaxis, :, :]
|
||||
input_ab = input_ab[:,:, 1:3].transpose((2, 0, 1))
|
||||
|
||||
if layerimg.shape[2] == 4: #remove alpha channel in image if present
|
||||
layerimg = layerimg[:,:,0:3]
|
||||
|
||||
if torch.cuda.is_available() and not cflag:
|
||||
gimp.progress_init("(Using GPU) Running deepcolor for " + ilayerimg.name + "...")
|
||||
gpu_id = 0
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running deepcolor for " + ilayerimg.name + "...")
|
||||
gpu_id = None
|
||||
|
||||
colorModel = CI.ColorizeImageTorch(Xd=256)
|
||||
colorModel.prep_net(gpu_id, baseLoc + 'weights/colorize/caffemodel.pth')
|
||||
colorModel.load_image(layerimg) # load an image
|
||||
|
||||
img_out = colorModel.net_forward(input_ab, mask,f=cflag) # run model, returns 256x256 image
|
||||
img_out_fullres = colorModel.get_img_fullres() # get image at full resolution
|
||||
|
||||
createResultLayer(tmp1, 'new_' + ilayerimg.name, img_out_fullres)
|
||||
|
||||
|
||||
|
||||
register(
|
||||
"deepcolor",
|
||||
"deepcolor",
|
||||
"Running deepcolor.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"deepcolor...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_LAYER, "drawinglayer", "Original Image:", None),
|
||||
(PF_LAYER, "drawinglayer", "Color Mask:", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
deepcolor, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,87 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'PyTorch-Image-Dehazing'])
|
||||
|
||||
|
||||
import torch
|
||||
import net
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def clrImg(data_hazy,cFlag):
|
||||
data_hazy = (data_hazy / 255.0)
|
||||
data_hazy = torch.from_numpy(data_hazy).float()
|
||||
data_hazy = data_hazy.permute(2, 0, 1)
|
||||
dehaze_net = net.dehaze_net()
|
||||
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
dehaze_net = dehaze_net.cuda()
|
||||
dehaze_net.load_state_dict(torch.load(baseLoc+'weights/deepdehaze/dehazer.pth'))
|
||||
data_hazy = data_hazy.cuda()
|
||||
else:
|
||||
dehaze_net.load_state_dict(torch.load(baseLoc+'weights/deepdehaze/dehazer.pth',map_location=torch.device("cpu")))
|
||||
|
||||
gimp.progress_update(float(0.005))
|
||||
gimp.displays_flush()
|
||||
data_hazy = data_hazy.unsqueeze(0)
|
||||
clean_image = dehaze_net(data_hazy)
|
||||
out = clean_image.detach().cpu().numpy()[0,:,:,:]*255
|
||||
out = np.clip(np.transpose(out,(1,2,0)),0,255).astype(np.uint8)
|
||||
return out
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def deepdehazing(img, layer, cFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Dehazing " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Dehazing " + layer.name + "...")
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:,:,0:3]
|
||||
cpy = clrImg(imgmat,cFlag)
|
||||
createResultLayer(img, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"deep-dehazing",
|
||||
"deep-dehazing",
|
||||
"Dehaze image based on deep learning.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"deep-dehazing...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
deepdehazing, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,132 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'PD-Denoising-pytorch'])
|
||||
|
||||
|
||||
from denoiser import *
|
||||
from argparse import Namespace
|
||||
|
||||
def clrImg(Img,cFlag):
|
||||
w, h, _ = Img.shape
|
||||
opt = Namespace(color=1, cond=1, delog='logsdc', ext_test_noise_level=None,
|
||||
k=0, keep_ind=None, mode='MC', num_of_layers=20, out_dir='results_bc',
|
||||
output_map=0, ps=2, ps_scale=2, real_n=1, refine=0, refine_opt=1,
|
||||
rescale=1, scale=1, spat_n=0, test_data='real_night', test_data_gnd='Set12',
|
||||
test_noise_level=None, wbin=512, zeroout=0)
|
||||
c = 1 if opt.color == 0 else 3
|
||||
model = DnCNN_c(channels=c, num_of_layers=opt.num_of_layers, num_of_est=2 * c)
|
||||
model_est = Estimation_direct(c, 2 * c)
|
||||
# device_ids = [0]
|
||||
# model = nn.DataParallel(net, device_ids=device_ids)
|
||||
# model_est = nn.DataParallel(est_net, device_ids=device_ids)# Estimator Model
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
ckpt_est = torch.load(baseLoc+'weights/deepdenoise/est_net.pth')
|
||||
ckpt = torch.load(baseLoc+'weights/deepdenoise/net.pth')
|
||||
model = model.cuda()
|
||||
model_est = model_est.cuda()
|
||||
else:
|
||||
ckpt = torch.load(baseLoc+'weights/deepdenoise/net.pth',map_location=torch.device("cpu"))
|
||||
ckpt_est = torch.load(baseLoc+'weights/deepdenoise/est_net.pth',map_location=torch.device("cpu"))
|
||||
|
||||
ckpt = {key.replace("module.",""):value for key,value in ckpt.items()}
|
||||
ckpt_est = {key.replace("module.",""):value for key,value in ckpt_est.items()}
|
||||
|
||||
model.load_state_dict(ckpt)
|
||||
model.eval()
|
||||
model_est.load_state_dict(ckpt_est)
|
||||
model_est.eval()
|
||||
gimp.progress_update(float(0.005))
|
||||
gimp.displays_flush()
|
||||
|
||||
Img = Img[:, :, ::-1] # change it to RGB
|
||||
Img = cv2.resize(Img, (0, 0), fx=opt.scale, fy=opt.scale)
|
||||
if opt.color == 0:
|
||||
Img = Img[:, :, 0] # For gray images
|
||||
Img = np.expand_dims(Img, 2)
|
||||
pss = 1
|
||||
if opt.ps == 1:
|
||||
pss = decide_scale_factor(Img / 255., model_est, color=opt.color, thre=0.008, plot_flag=1, stopping=4,
|
||||
mark=opt.out_dir + '/' + file_name)[0]
|
||||
# print(pss)
|
||||
Img = pixelshuffle(Img, pss)
|
||||
elif opt.ps == 2:
|
||||
pss = opt.ps_scale
|
||||
|
||||
merge_out = np.zeros([w, h, 3])
|
||||
wbin = opt.wbin
|
||||
i = 0
|
||||
idx=0
|
||||
t=float(w*h)/float(wbin*wbin)
|
||||
while i < w:
|
||||
i_end = min(i + wbin, w)
|
||||
j = 0
|
||||
while j < h:
|
||||
j_end = min(j + wbin, h)
|
||||
patch = Img[i:i_end, j:j_end, :]
|
||||
patch_merge_out_numpy = denoiser(patch, c, pss, model, model_est, opt, cFlag)
|
||||
merge_out[i:i_end, j:j_end, :] = patch_merge_out_numpy
|
||||
j = j_end
|
||||
idx=idx+1
|
||||
gimp.progress_update(float(idx)/float(t))
|
||||
gimp.displays_flush()
|
||||
i = i_end
|
||||
|
||||
|
||||
return merge_out[:, :, ::-1]
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def deepdenoise(img, layer,cFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Denoising " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Denoising " + layer.name + "...")
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:,:,0:3]
|
||||
cpy = clrImg(imgmat,cFlag)
|
||||
createResultLayer(img, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"deep-denoising",
|
||||
"deep-denoising",
|
||||
"Denoise image based on deep learning.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"deep-denoising...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
deepdenoise, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,112 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'pytorch-deep-image-matting'])
|
||||
|
||||
import torch
|
||||
from argparse import Namespace
|
||||
import net
|
||||
import cv2
|
||||
import os
|
||||
import numpy as np
|
||||
from deploy import inference_img_whole
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 1, 100,
|
||||
NORMAL_MODE) # image.active_layer.type or RGB_IMAGE
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def getnewalpha(image, mask, cFlag):
|
||||
if image.shape[2] == 4: # get rid of alpha channel
|
||||
image = image[:, :, 0:3]
|
||||
if mask.shape[2] == 4: # get rid of alpha channel
|
||||
mask = mask[:, :, 0:3]
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
trimap = mask[:, :, 0]
|
||||
|
||||
cudaFlag = False
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
cudaFlag = True
|
||||
|
||||
args = Namespace(crop_or_resize='whole', cuda=cudaFlag, max_size=1600,
|
||||
resume=baseLoc + 'weights/deepmatting/stage1_sad_57.1.pth', stage=1)
|
||||
model = net.VGG16(args)
|
||||
|
||||
if cudaFlag:
|
||||
ckpt = torch.load(args.resume)
|
||||
else:
|
||||
ckpt = torch.load(args.resume, map_location=torch.device("cpu"))
|
||||
model.load_state_dict(ckpt['state_dict'], strict=True)
|
||||
if cudaFlag:
|
||||
model = model.cuda()
|
||||
|
||||
# ckpt = torch.load(args.resume)
|
||||
# model.load_state_dict(ckpt['state_dict'], strict=True)
|
||||
# model = model.cuda()
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
with torch.no_grad():
|
||||
pred_mattes = inference_img_whole(args, model, image, trimap)
|
||||
pred_mattes = (pred_mattes * 255).astype(np.uint8)
|
||||
pred_mattes[trimap == 255] = 255
|
||||
pred_mattes[trimap == 0] = 0
|
||||
# pred_mattes = np.repeat(pred_mattes[:, :, np.newaxis], 3, axis=2)
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
pred_mattes = np.dstack((image, pred_mattes))
|
||||
return pred_mattes
|
||||
|
||||
|
||||
def deepmatting(imggimp, curlayer, layeri, layerm, cFlag):
|
||||
img = channelData(layeri)
|
||||
mask = channelData(layerm)
|
||||
if img.shape[0] != imggimp.height or img.shape[1] != imggimp.width or mask.shape[0] != imggimp.height or mask.shape[1] != imggimp.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) for both layers and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Running deep-matting for " + layeri.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running deep-matting for " + layeri.name + "...")
|
||||
cpy = getnewalpha(img, mask, cFlag)
|
||||
createResultLayer(imggimp, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"deep-matting",
|
||||
"deep-matting",
|
||||
"Running image matting.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"deepmatting...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_LAYER, "drawinglayer", "Original Image:", None),
|
||||
(PF_LAYER, "drawinglayer", "Trimap Mask:", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
deepmatting, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,111 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'EnlightenGAN'])
|
||||
|
||||
from argparse import Namespace
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from models.models import create_model
|
||||
from data.base_dataset import get_transform
|
||||
|
||||
|
||||
|
||||
|
||||
def getEnlighten(input_image,cFlag):
|
||||
|
||||
opt = Namespace(D_P_times2=False, IN_vgg=False, aspect_ratio=1.0, batchSize=1,
|
||||
checkpoints_dir=baseLoc+'weights/', dataroot='test_dataset',
|
||||
dataset_mode='unaligned', display_id=1, display_port=8097,
|
||||
display_single_pane_ncols=0, display_winsize=256, fcn=0,
|
||||
fineSize=256, gpu_ids=[0], high_times=400, how_many=50,
|
||||
hybrid_loss=False, identity=0.0, input_linear=False, input_nc=3,
|
||||
instance_norm=0.0, isTrain=False, l1=10.0, lambda_A=10.0,
|
||||
lambda_B=10.0, latent_norm=False, latent_threshold=False,
|
||||
lighten=False, linear=False, linear_add=False, loadSize=286,
|
||||
low_times=200, max_dataset_size='inf', model='single',
|
||||
multiply=False, nThreads=1, n_layers_D=3, n_layers_patchD=3,
|
||||
name='enlightening', ndf=64, new_lr=False, ngf=64, no_dropout=True,
|
||||
no_flip=True, no_vgg_instance=False, noise=0, norm='instance',
|
||||
norm_attention=False, ntest='inf', output_nc=3, patchD=False,
|
||||
patchD_3=0, patchSize=64, patch_vgg=False, phase='test',
|
||||
resize_or_crop='no', results_dir='./results/', self_attention=True,
|
||||
serial_batches=True, skip=1.0, syn_norm=False, tanh=False,
|
||||
times_residual=True, use_avgpool=0, use_mse=False, use_norm=1.0,
|
||||
use_ragan=False, use_wgan=0.0, vary=1, vgg=0, vgg_choose='relu5_3',
|
||||
vgg_maxpooling=False, vgg_mean=False, which_direction='AtoB',
|
||||
which_epoch='200', which_model_netD='basic', which_model_netG='sid_unet_resize', cFlag=cFlag)
|
||||
|
||||
im = cv2.cvtColor(input_image,cv2.COLOR_RGB2BGR)
|
||||
transform = get_transform(opt)
|
||||
A_img = transform(im)
|
||||
r, g, b = A_img[0] + 1, A_img[1] + 1, A_img[2] + 1
|
||||
A_gray = 1. - (0.299 * r + 0.587 * g + 0.114 * b) / 2.
|
||||
A_gray = torch.unsqueeze(A_gray, 0)
|
||||
data = {'A': A_img.unsqueeze(0), 'B': A_img.unsqueeze(0), 'A_gray': A_gray.unsqueeze(0), 'input_img': A_img.unsqueeze(0), 'A_paths': 'A_path', 'B_paths': 'B_path'}
|
||||
|
||||
model = create_model(opt)
|
||||
model.set_input(data)
|
||||
visuals = model.predict()
|
||||
out = visuals['fake_B'].astype(np.uint8)
|
||||
out = cv2.cvtColor(out,cv2.COLOR_BGR2RGB)
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def Enlighten(img, layer,cFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Enlighten " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Enlighten " + layer.name + "...")
|
||||
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:,:,0:3]
|
||||
cpy = getEnlighten(imgmat,cFlag)
|
||||
createResultLayer(img, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"enlighten",
|
||||
"enlighten",
|
||||
"Enlighten image based on deep learning.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"enlighten...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
Enlighten, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,188 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools',baseLoc+'CelebAMask-HQ/MaskGAN_demo'])
|
||||
|
||||
|
||||
import torch
|
||||
from argparse import Namespace
|
||||
from models.models import create_model
|
||||
from data.base_dataset import get_params, get_transform, normalize
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
colors = np.array([[0, 0, 0], [204, 0, 0], [76, 153, 0], \
|
||||
[204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], \
|
||||
[51, 255, 255], [102, 51, 0], [255, 0, 0], [102, 204, 0], \
|
||||
[255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], \
|
||||
[0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]])
|
||||
colors = colors.astype(np.uint8)
|
||||
|
||||
def getlabelmat(mask,idx):
|
||||
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
||||
x[mask==idx,0]=colors[idx][0]
|
||||
x[mask==idx,1]=colors[idx][1]
|
||||
x[mask==idx,2]=colors[idx][2]
|
||||
return x
|
||||
|
||||
def colorMask(mask):
|
||||
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
||||
for idx in range(19):
|
||||
x=x+getlabelmat(mask,idx)
|
||||
# mask=np.dstack((mask1,mask2,mask3))
|
||||
return np.uint8(x)
|
||||
|
||||
def labelMask(mask):
|
||||
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
||||
for idx in range(19):
|
||||
tmp=np.logical_and(mask[:,:,0]==colors[idx][0],mask[:,:,1]==colors[idx][1])
|
||||
tmp2=np.logical_and(tmp,mask[:,:,2]==colors[idx][2])
|
||||
x[tmp2]=idx
|
||||
return x
|
||||
|
||||
def getOptions():
|
||||
mydict={'aspect_ratio': 1.0,
|
||||
'batchSize': 1,
|
||||
'checkpoints_dir': baseLoc+'weights/facegen',
|
||||
'cluster_path': 'features_clustered_010.npy',
|
||||
'data_type': 32,
|
||||
'dataroot': '../Data_preprocessing/',
|
||||
'display_winsize': 512,
|
||||
'engine': None,
|
||||
'export_onnx': None,
|
||||
'fineSize': 512,
|
||||
'gpu_ids': [],
|
||||
'how_many': 1000,
|
||||
'input_nc': 3,
|
||||
'isTrain': False,
|
||||
'label_nc': 19,
|
||||
'loadSize': 512,
|
||||
'max_dataset_size': 'inf',
|
||||
'model': 'pix2pixHD',
|
||||
'nThreads': 2,
|
||||
'n_blocks_global': 4,
|
||||
'n_blocks_local': 3,
|
||||
'n_downsample_global': 4,
|
||||
'n_local_enhancers': 1,
|
||||
'name': 'label2face_512p',
|
||||
'netG': 'global',
|
||||
'ngf': 64,
|
||||
'niter_fix_global': 0,
|
||||
'no_flip': False,
|
||||
'norm': 'instance',
|
||||
'ntest': 'inf',
|
||||
'onnx': None,
|
||||
'output_nc': 3,
|
||||
'phase': 'test',
|
||||
'resize_or_crop': 'scale_width',
|
||||
'results_dir': './results/',
|
||||
'serial_batches': False,
|
||||
'tf_log': False,
|
||||
'use_dropout': False,
|
||||
'use_encoded_image': False,
|
||||
'verbose': False,
|
||||
'which_epoch': 'latest'}
|
||||
args = Namespace(**mydict)
|
||||
return args
|
||||
|
||||
def channelData(layer):#convert gimp image to numpy
|
||||
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
||||
pixChars=region[:,:] # Take whole layer
|
||||
bpp=region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
||||
|
||||
def createResultLayer(image,name,result):
|
||||
rlBytes=np.uint8(result).tobytes();
|
||||
rl=gimp.Layer(image,name,image.width,image.height,image.active_layer.type,100,NORMAL_MODE)
|
||||
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
||||
region[:,:]=rlBytes
|
||||
image.add_layer(rl,0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def getnewface(img,mask,mask_m,cFlag):
|
||||
h,w,d = img.shape
|
||||
img = Image.fromarray(img)
|
||||
lmask = labelMask(mask)
|
||||
lmask_m = labelMask(mask_m)
|
||||
|
||||
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = str(0)
|
||||
opt = getOptions()
|
||||
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
opt.gpu_ids=[0]
|
||||
|
||||
model = create_model(opt)
|
||||
|
||||
params = get_params(opt, (512,512))
|
||||
transform_mask = get_transform(opt, params, method=Image.NEAREST, normalize=False, normalize_mask=True)
|
||||
transform_image = get_transform(opt, params)
|
||||
mask = transform_mask(Image.fromarray(np.uint8(lmask)))
|
||||
mask_m = transform_mask(Image.fromarray(np.uint8(lmask_m)))
|
||||
img = transform_image(img)
|
||||
|
||||
generated = model.inference(torch.FloatTensor([mask_m.numpy()]), torch.FloatTensor([mask.numpy()]), torch.FloatTensor([img.numpy()]), cFlag)
|
||||
|
||||
result = generated.permute(0, 2, 3, 1)
|
||||
if torch.cuda.is_available():
|
||||
result = result.detach().cpu().numpy()
|
||||
else:
|
||||
result = result.detach().numpy()
|
||||
|
||||
result = (result + 1) * 127.5
|
||||
result = np.asarray(result[0,:,:,:], dtype=np.uint8)
|
||||
result = Image.fromarray(result)
|
||||
result = result.resize([w,h])
|
||||
|
||||
result = np.array(result)
|
||||
return result
|
||||
|
||||
|
||||
def facegen(imggimp, curlayer,layeri,layerm,layermm,cFlag):
|
||||
img = channelData(layeri)
|
||||
mask = channelData(layerm)
|
||||
mask_m = channelData(layermm)
|
||||
if img.shape[0] != imggimp.height or img.shape[1] != imggimp.width or mask.shape[0] != imggimp.height or mask.shape[1] != imggimp.width or mask_m.shape[0] != imggimp.height or mask_m.shape[1] != imggimp.width:
|
||||
pdb.gimp_message("Do (Layer -> Layer to Image Size) for all layers and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Running face gen for " + layeri.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running face gen for " + layeri.name + "...")
|
||||
if img.shape[2] == 4: # get rid of alpha channel
|
||||
img = img[:, :, 0:3]
|
||||
if mask.shape[2] == 4: # get rid of alpha channel
|
||||
mask = mask[:, :, 0:3]
|
||||
if mask_m.shape[2] == 4: # get rid of alpha channel
|
||||
mask_m = mask_m[:, :, 0:3]
|
||||
cpy = getnewface(img,mask,mask_m,cFlag)
|
||||
createResultLayer(imggimp,'new_output',cpy)
|
||||
|
||||
|
||||
|
||||
register(
|
||||
"facegen",
|
||||
"facegen",
|
||||
"Running face gen.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"facegen...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_LAYER, "drawinglayer", "Original Image:", None),
|
||||
(PF_LAYER, "drawinglayer2", "Original Mask:", None),
|
||||
(PF_LAYER, "drawinglayer3", "Modified Mask:", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False),
|
||||
],
|
||||
[],
|
||||
facegen, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,176 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'face-parsing-PyTorch'])
|
||||
|
||||
from model import BiSeNet
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torchvision import transforms, datasets
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
colors = np.array([[0, 0, 0],
|
||||
[204, 0, 0],
|
||||
[0, 255, 255],
|
||||
[51, 255, 255],
|
||||
[51, 51, 255],
|
||||
[204, 0, 204],
|
||||
[204, 204, 0],
|
||||
[102, 51, 0],
|
||||
[255, 0, 0],
|
||||
[0, 204, 204],
|
||||
[76, 153, 0],
|
||||
[102, 204, 0],
|
||||
[255, 255, 0],
|
||||
[0, 0, 153],
|
||||
[255, 153, 51],
|
||||
[0, 51, 0],
|
||||
[0, 204, 0],
|
||||
[0, 0, 204],
|
||||
[255, 51, 153]])
|
||||
colors = colors.astype(np.uint8)
|
||||
|
||||
|
||||
def getlabelmat(mask, idx):
|
||||
x = np.zeros((mask.shape[0], mask.shape[1], 3))
|
||||
x[mask == idx, 0] = colors[idx][0]
|
||||
x[mask == idx, 1] = colors[idx][1]
|
||||
x[mask == idx, 2] = colors[idx][2]
|
||||
return x
|
||||
|
||||
|
||||
def colorMask(mask):
|
||||
x = np.zeros((mask.shape[0], mask.shape[1], 3))
|
||||
for idx in range(19):
|
||||
x = x + getlabelmat(mask, idx)
|
||||
return np.uint8(x)
|
||||
|
||||
|
||||
def getface(input_image, cFlag):
|
||||
save_pth = baseLoc + 'weights/faceparse/79999_iter.pth'
|
||||
input_image = Image.fromarray(input_image)
|
||||
|
||||
n_classes = 19
|
||||
net = BiSeNet(n_classes=n_classes)
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
net.cuda()
|
||||
net.load_state_dict(torch.load(save_pth))
|
||||
else:
|
||||
net.load_state_dict(torch.load(save_pth, map_location=lambda storage, loc: storage))
|
||||
|
||||
net.eval()
|
||||
|
||||
to_tensor = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
||||
])
|
||||
|
||||
with torch.no_grad():
|
||||
img = input_image.resize((512, 512), Image.BILINEAR)
|
||||
img = to_tensor(img)
|
||||
img = torch.unsqueeze(img, 0)
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
img = img.cuda()
|
||||
out = net(img)[0]
|
||||
if torch.cuda.is_available():
|
||||
parsing = out.squeeze(0).cpu().numpy().argmax(0)
|
||||
else:
|
||||
parsing = out.squeeze(0).numpy().argmax(0)
|
||||
|
||||
parsing = Image.fromarray(np.uint8(parsing))
|
||||
parsing = parsing.resize(input_image.size)
|
||||
parsing = np.array(parsing)
|
||||
|
||||
return parsing
|
||||
|
||||
|
||||
def getSeg(input_image):
|
||||
model = torch.load(baseLoc + 'deeplabv3+model.pt')
|
||||
model.eval()
|
||||
preprocess = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
input_image = Image.fromarray(input_image)
|
||||
|
||||
input_tensor = preprocess(input_image)
|
||||
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
|
||||
|
||||
# move the input and model to GPU for speed if available
|
||||
if torch.cuda.is_available():
|
||||
input_batch = input_batch.to('cuda')
|
||||
model.to('cuda')
|
||||
|
||||
with torch.no_grad():
|
||||
output = model(input_batch)['out'][0]
|
||||
output_predictions = output.argmax(0)
|
||||
|
||||
# create a color pallette, selecting a color for each class
|
||||
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
|
||||
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
|
||||
colors = (colors % 255).numpy().astype("uint8")
|
||||
|
||||
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
|
||||
|
||||
tmp = np.array(r)
|
||||
tmp2 = 10 * np.repeat(tmp[:, :, np.newaxis], 3, axis=2)
|
||||
|
||||
return tmp2
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes()
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def faceparse(img, layer, cFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:, :, 0:3]
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Running face parse for " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running face parse for " + layer.name + "...")
|
||||
cpy = getface(imgmat, cFlag)
|
||||
cpy = colorMask(cpy)
|
||||
createResultLayer(img, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"faceparse",
|
||||
"faceparse",
|
||||
"Running face parse.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"faceparse...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
faceparse, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,21 +0,0 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2017 Jun-Yan Zhu and Richard Zhang
|
||||
|
||||
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.
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -1,565 +0,0 @@
|
||||
import numpy as np
|
||||
# import matplotlib.pyplot as plt
|
||||
# from skimage import color
|
||||
# from sklearn.cluster import KMeans
|
||||
import os
|
||||
import cv2
|
||||
from scipy.ndimage.interpolation import zoom
|
||||
|
||||
def create_temp_directory(path_template, N=1e8):
|
||||
print(path_template)
|
||||
cur_path = path_template % np.random.randint(0, N)
|
||||
while(os.path.exists(cur_path)):
|
||||
cur_path = path_template % np.random.randint(0, N)
|
||||
print('Creating directory: %s' % cur_path)
|
||||
os.mkdir(cur_path)
|
||||
return cur_path
|
||||
|
||||
|
||||
def lab2rgb_transpose(img_l, img_ab):
|
||||
''' INPUTS
|
||||
img_l 1xXxX [0,100]
|
||||
img_ab 2xXxX [-100,100]
|
||||
OUTPUTS
|
||||
returned value is XxXx3 '''
|
||||
pred_lab = np.concatenate((img_l, img_ab), axis=0).transpose((1, 2, 0))
|
||||
# im = color.lab2rgb(pred_lab)
|
||||
im = cv2.cvtColor(pred_lab.astype('float32'),cv2.COLOR_LAB2RGB)
|
||||
pred_rgb = (np.clip(im, 0, 1) * 255).astype('uint8')
|
||||
return pred_rgb
|
||||
|
||||
|
||||
def rgb2lab_transpose(img_rgb):
|
||||
''' INPUTS
|
||||
img_rgb XxXx3
|
||||
OUTPUTS
|
||||
returned value is 3xXxX '''
|
||||
# im=color.rgb2lab(img_rgb)
|
||||
im = cv2.cvtColor(img_rgb.astype(np.float32)/255, cv2.COLOR_RGB2LAB)
|
||||
return im.transpose((2, 0, 1))
|
||||
|
||||
|
||||
class ColorizeImageBase():
|
||||
def __init__(self, Xd=256, Xfullres_max=10000):
|
||||
self.Xd = Xd
|
||||
self.img_l_set = False
|
||||
self.net_set = False
|
||||
self.Xfullres_max = Xfullres_max # maximum size of maximum dimension
|
||||
self.img_just_set = False # this will be true whenever image is just loaded
|
||||
# net_forward can set this to False if they want
|
||||
|
||||
def prep_net(self):
|
||||
raise Exception("Should be implemented by base class")
|
||||
|
||||
# ***** Image prepping *****
|
||||
def load_image(self, im):
|
||||
# rgb image [CxXdxXd]
|
||||
self.img_rgb_fullres = im.copy()
|
||||
self._set_img_lab_fullres_()
|
||||
|
||||
im = cv2.resize(im, (self.Xd, self.Xd))
|
||||
self.img_rgb = im.copy()
|
||||
# self.img_rgb = sp.misc.imresize(plt.imread(input_path),(self.Xd,self.Xd)).transpose((2,0,1))
|
||||
|
||||
self.img_l_set = True
|
||||
|
||||
# convert into lab space
|
||||
self._set_img_lab_()
|
||||
self._set_img_lab_mc_()
|
||||
|
||||
def set_image(self, input_image):
|
||||
self.img_rgb_fullres = input_image.copy()
|
||||
self._set_img_lab_fullres_()
|
||||
|
||||
self.img_l_set = True
|
||||
|
||||
self.img_rgb = input_image
|
||||
# convert into lab space
|
||||
self._set_img_lab_()
|
||||
self._set_img_lab_mc_()
|
||||
|
||||
def net_forward(self, input_ab, input_mask):
|
||||
# INPUTS
|
||||
# ab 2xXxX input color patches (non-normalized)
|
||||
# mask 1xXxX input mask, indicating which points have been provided
|
||||
# assumes self.img_l_mc has been set
|
||||
|
||||
if(not self.img_l_set):
|
||||
print('I need to have an image!')
|
||||
return -1
|
||||
if(not self.net_set):
|
||||
print('I need to have a net!')
|
||||
return -1
|
||||
|
||||
self.input_ab = input_ab
|
||||
self.input_ab_mc = (input_ab - self.ab_mean) / self.ab_norm
|
||||
self.input_mask = input_mask
|
||||
self.input_mask_mult = input_mask * self.mask_mult
|
||||
return 0
|
||||
|
||||
def get_result_PSNR(self, result=-1, return_SE_map=False):
|
||||
if np.array((result)).flatten()[0] == -1:
|
||||
cur_result = self.get_img_forward()
|
||||
else:
|
||||
cur_result = result.copy()
|
||||
SE_map = (1. * self.img_rgb - cur_result)**2
|
||||
cur_MSE = np.mean(SE_map)
|
||||
cur_PSNR = 20 * np.log10(255. / np.sqrt(cur_MSE))
|
||||
if return_SE_map:
|
||||
return(cur_PSNR, SE_map)
|
||||
else:
|
||||
return cur_PSNR
|
||||
|
||||
def get_img_forward(self):
|
||||
# get image with point estimate
|
||||
return self.output_rgb
|
||||
|
||||
def get_img_gray(self):
|
||||
# Get black and white image
|
||||
return lab2rgb_transpose(self.img_l, np.zeros((2, self.Xd, self.Xd)))
|
||||
|
||||
def get_img_gray_fullres(self):
|
||||
# Get black and white image
|
||||
return lab2rgb_transpose(self.img_l_fullres, np.zeros((2, self.img_l_fullres.shape[1], self.img_l_fullres.shape[2])))
|
||||
|
||||
def get_img_fullres(self):
|
||||
# This assumes self.img_l_fullres, self.output_ab are set.
|
||||
# Typically, this means that set_image() and net_forward()
|
||||
# have been called.
|
||||
# bilinear upsample
|
||||
zoom_factor = (1, 1. * self.img_l_fullres.shape[1] / self.output_ab.shape[1], 1. * self.img_l_fullres.shape[2] / self.output_ab.shape[2])
|
||||
output_ab_fullres = zoom(self.output_ab, zoom_factor, order=1)
|
||||
|
||||
return lab2rgb_transpose(self.img_l_fullres, output_ab_fullres)
|
||||
|
||||
def get_input_img_fullres(self):
|
||||
zoom_factor = (1, 1. * self.img_l_fullres.shape[1] / self.input_ab.shape[1], 1. * self.img_l_fullres.shape[2] / self.input_ab.shape[2])
|
||||
input_ab_fullres = zoom(self.input_ab, zoom_factor, order=1)
|
||||
return lab2rgb_transpose(self.img_l_fullres, input_ab_fullres)
|
||||
|
||||
def get_input_img(self):
|
||||
return lab2rgb_transpose(self.img_l, self.input_ab)
|
||||
|
||||
def get_img_mask(self):
|
||||
# Get black and white image
|
||||
return lab2rgb_transpose(100. * (1 - self.input_mask), np.zeros((2, self.Xd, self.Xd)))
|
||||
|
||||
def get_img_mask_fullres(self):
|
||||
# Get black and white image
|
||||
zoom_factor = (1, 1. * self.img_l_fullres.shape[1] / self.input_ab.shape[1], 1. * self.img_l_fullres.shape[2] / self.input_ab.shape[2])
|
||||
input_mask_fullres = zoom(self.input_mask, zoom_factor, order=0)
|
||||
return lab2rgb_transpose(100. * (1 - input_mask_fullres), np.zeros((2, input_mask_fullres.shape[1], input_mask_fullres.shape[2])))
|
||||
|
||||
def get_sup_img(self):
|
||||
return lab2rgb_transpose(50 * self.input_mask, self.input_ab)
|
||||
|
||||
def get_sup_fullres(self):
|
||||
zoom_factor = (1, 1. * self.img_l_fullres.shape[1] / self.output_ab.shape[1], 1. * self.img_l_fullres.shape[2] / self.output_ab.shape[2])
|
||||
input_mask_fullres = zoom(self.input_mask, zoom_factor, order=0)
|
||||
input_ab_fullres = zoom(self.input_ab, zoom_factor, order=0)
|
||||
return lab2rgb_transpose(50 * input_mask_fullres, input_ab_fullres)
|
||||
|
||||
# ***** Private functions *****
|
||||
def _set_img_lab_fullres_(self):
|
||||
# adjust full resolution image to be within maximum dimension is within Xfullres_max
|
||||
Xfullres = self.img_rgb_fullres.shape[0]
|
||||
Yfullres = self.img_rgb_fullres.shape[1]
|
||||
if Xfullres > self.Xfullres_max or Yfullres > self.Xfullres_max:
|
||||
if Xfullres > Yfullres:
|
||||
zoom_factor = 1. * self.Xfullres_max / Xfullres
|
||||
else:
|
||||
zoom_factor = 1. * self.Xfullres_max / Yfullres
|
||||
self.img_rgb_fullres = zoom(self.img_rgb_fullres, (zoom_factor, zoom_factor, 1), order=1)
|
||||
|
||||
self.img_lab_fullres = cv2.cvtColor(self.img_rgb_fullres.astype(np.float32) / 255, cv2.COLOR_RGB2LAB).transpose((2, 0, 1))
|
||||
# self.img_lab_fullres = color.rgb2lab(self.img_rgb_fullres).transpose((2, 0, 1))
|
||||
self.img_l_fullres = self.img_lab_fullres[[0], :, :]
|
||||
self.img_ab_fullres = self.img_lab_fullres[1:, :, :]
|
||||
|
||||
def _set_img_lab_(self):
|
||||
# set self.img_lab from self.im_rgb
|
||||
self.img_lab = cv2.cvtColor(self.img_rgb.astype(np.float32) / 255, cv2.COLOR_RGB2LAB).transpose((2, 0, 1))
|
||||
# self.img_lab = color.rgb2lab(self.img_rgb).transpose((2, 0, 1))
|
||||
self.img_l = self.img_lab[[0], :, :]
|
||||
self.img_ab = self.img_lab[1:, :, :]
|
||||
|
||||
def _set_img_lab_mc_(self):
|
||||
# set self.img_lab_mc from self.img_lab
|
||||
# lab image, mean centered [XxYxX]
|
||||
self.img_lab_mc = self.img_lab / np.array((self.l_norm, self.ab_norm, self.ab_norm))[:, np.newaxis, np.newaxis] - np.array(
|
||||
(self.l_mean / self.l_norm, self.ab_mean / self.ab_norm, self.ab_mean / self.ab_norm))[:, np.newaxis, np.newaxis]
|
||||
self._set_img_l_()
|
||||
|
||||
def _set_img_l_(self):
|
||||
self.img_l_mc = self.img_lab_mc[[0], :, :]
|
||||
self.img_l_set = True
|
||||
|
||||
def _set_img_ab_(self):
|
||||
self.img_ab_mc = self.img_lab_mc[[1, 2], :, :]
|
||||
|
||||
def _set_out_ab_(self):
|
||||
self.output_lab = rgb2lab_transpose(self.output_rgb)
|
||||
self.output_ab = self.output_lab[1:, :, :]
|
||||
|
||||
|
||||
class ColorizeImageTorch(ColorizeImageBase):
|
||||
def __init__(self, Xd=256, maskcent=False):
|
||||
print('ColorizeImageTorch instantiated')
|
||||
ColorizeImageBase.__init__(self, Xd)
|
||||
self.l_norm = 1.
|
||||
self.ab_norm = 1.
|
||||
self.l_mean = 50.
|
||||
self.ab_mean = 0.
|
||||
self.mask_mult = 1.
|
||||
self.mask_cent = .5 if maskcent else 0
|
||||
|
||||
# Load grid properties
|
||||
self.pts_in_hull = np.array(np.meshgrid(np.arange(-110, 120, 10), np.arange(-110, 120, 10))).reshape((2, 529)).T
|
||||
|
||||
# ***** Net preparation *****
|
||||
def prep_net(self, gpu_id=None, path='', dist=False):
|
||||
import torch
|
||||
import models.pytorch.model as model
|
||||
print('path = %s' % path)
|
||||
print('Model set! dist mode? ', dist)
|
||||
self.net = model.SIGGRAPHGenerator(dist=dist)
|
||||
state_dict = torch.load(path)
|
||||
if hasattr(state_dict, '_metadata'):
|
||||
del state_dict._metadata
|
||||
|
||||
# patch InstanceNorm checkpoints prior to 0.4
|
||||
for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop
|
||||
self.__patch_instance_norm_state_dict(state_dict, self.net, key.split('.'))
|
||||
self.net.load_state_dict(state_dict)
|
||||
if gpu_id != None:
|
||||
self.net.cuda()
|
||||
self.net.eval()
|
||||
self.net_set = True
|
||||
|
||||
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
|
||||
key = keys[i]
|
||||
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
|
||||
if module.__class__.__name__.startswith('InstanceNorm') and \
|
||||
(key == 'running_mean' or key == 'running_var'):
|
||||
if getattr(module, key) is None:
|
||||
state_dict.pop('.'.join(keys))
|
||||
if module.__class__.__name__.startswith('InstanceNorm') and \
|
||||
(key == 'num_batches_tracked'):
|
||||
state_dict.pop('.'.join(keys))
|
||||
else:
|
||||
self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
|
||||
|
||||
# ***** Call forward *****
|
||||
def net_forward(self, input_ab, input_mask, f):
|
||||
# INPUTS
|
||||
# ab 2xXxX input color patches (non-normalized)
|
||||
# mask 1xXxX input mask, indicating which points have been provided
|
||||
# assumes self.img_l_mc has been set
|
||||
|
||||
if ColorizeImageBase.net_forward(self, input_ab, input_mask) == -1:
|
||||
return -1
|
||||
|
||||
# net_input_prepped = np.concatenate((self.img_l_mc, self.input_ab_mc, self.input_mask_mult), axis=0)
|
||||
|
||||
# return prediction
|
||||
# self.net.blobs['data_l_ab_mask'].data[...] = net_input_prepped
|
||||
# embed()
|
||||
output_ab = self.net.forward(self.img_l_mc, self.input_ab_mc, self.input_mask_mult, self.mask_cent,f)[0, :, :, :].cpu().data.numpy()
|
||||
self.output_rgb = lab2rgb_transpose(self.img_l, output_ab)
|
||||
# self.output_rgb = lab2rgb_transpose(self.img_l, self.net.blobs[self.pred_ab_layer].data[0, :, :, :])
|
||||
|
||||
self._set_out_ab_()
|
||||
return self.output_rgb
|
||||
|
||||
def get_img_forward(self):
|
||||
# get image with point estimate
|
||||
return self.output_rgb
|
||||
|
||||
def get_img_gray(self):
|
||||
# Get black and white image
|
||||
return lab2rgb_transpose(self.img_l, np.zeros((2, self.Xd, self.Xd)))
|
||||
|
||||
|
||||
class ColorizeImageTorchDist(ColorizeImageTorch):
|
||||
def __init__(self, Xd=256, maskcent=False):
|
||||
ColorizeImageTorch.__init__(self, Xd)
|
||||
self.dist_ab_set = False
|
||||
self.pts_grid = np.array(np.meshgrid(np.arange(-110, 120, 10), np.arange(-110, 120, 10))).reshape((2, 529)).T
|
||||
self.in_hull = np.ones(529, dtype=bool)
|
||||
self.AB = self.pts_grid.shape[0] # 529
|
||||
self.A = int(np.sqrt(self.AB)) # 23
|
||||
self.B = int(np.sqrt(self.AB)) # 23
|
||||
self.dist_ab_full = np.zeros((self.AB, self.Xd, self.Xd))
|
||||
self.dist_ab_grid = np.zeros((self.A, self.B, self.Xd, self.Xd))
|
||||
self.dist_entropy = np.zeros((self.Xd, self.Xd))
|
||||
self.mask_cent = .5 if maskcent else 0
|
||||
|
||||
def prep_net(self, gpu_id=None, path='', dist=True, S=.2):
|
||||
ColorizeImageTorch.prep_net(self, gpu_id=gpu_id, path=path, dist=dist)
|
||||
# set S somehow
|
||||
|
||||
def net_forward(self, input_ab, input_mask):
|
||||
# INPUTS
|
||||
# ab 2xXxX input color patches (non-normalized)
|
||||
# mask 1xXxX input mask, indicating which points have been provided
|
||||
# assumes self.img_l_mc has been set
|
||||
|
||||
# embed()
|
||||
if ColorizeImageBase.net_forward(self, input_ab, input_mask) == -1:
|
||||
return -1
|
||||
|
||||
# set distribution
|
||||
(function_return, self.dist_ab) = self.net.forward(self.img_l_mc, self.input_ab_mc, self.input_mask_mult, self.mask_cent)
|
||||
function_return = function_return[0, :, :, :].cpu().data.numpy()
|
||||
self.dist_ab = self.dist_ab[0, :, :, :].cpu().data.numpy()
|
||||
self.dist_ab_set = True
|
||||
|
||||
# full grid, ABxXxX, AB = 529
|
||||
self.dist_ab_full[self.in_hull, :, :] = self.dist_ab
|
||||
|
||||
# gridded, AxBxXxX, A = 23
|
||||
self.dist_ab_grid = self.dist_ab_full.reshape((self.A, self.B, self.Xd, self.Xd))
|
||||
|
||||
# return
|
||||
return function_return
|
||||
|
||||
# def get_ab_reccs(self, h, w, K=5, N=25000, return_conf=False):
|
||||
# ''' Recommended colors at point (h,w)
|
||||
# Call this after calling net_forward
|
||||
# '''
|
||||
# if not self.dist_ab_set:
|
||||
# print('Need to set prediction first')
|
||||
# return 0
|
||||
#
|
||||
# # randomly sample from pdf
|
||||
# cmf = np.cumsum(self.dist_ab[:, h, w]) # CMF
|
||||
# cmf = cmf / cmf[-1]
|
||||
# cmf_bins = cmf
|
||||
#
|
||||
# # randomly sample N points
|
||||
# rnd_pts = np.random.uniform(low=0, high=1.0, size=N)
|
||||
# inds = np.digitize(rnd_pts, bins=cmf_bins)
|
||||
# rnd_pts_ab = self.pts_in_hull[inds, :]
|
||||
#
|
||||
# # run k-means
|
||||
# kmeans = KMeans(n_clusters=K).fit(rnd_pts_ab)
|
||||
#
|
||||
# # sort by cluster occupancy
|
||||
# k_label_cnt = np.histogram(kmeans.labels_, np.arange(0, K + 1))[0]
|
||||
# k_inds = np.argsort(k_label_cnt, axis=0)[::-1]
|
||||
#
|
||||
# cluster_per = 1. * k_label_cnt[k_inds] / N # percentage of points within cluster
|
||||
# cluster_centers = kmeans.cluster_centers_[k_inds, :] # cluster centers
|
||||
#
|
||||
# # cluster_centers = np.random.uniform(low=-100,high=100,size=(N,2))
|
||||
# if return_conf:
|
||||
# return cluster_centers, cluster_per
|
||||
# else:
|
||||
# return cluster_centers
|
||||
|
||||
def compute_entropy(self):
|
||||
# compute the distribution entropy (really slow right now)
|
||||
self.dist_entropy = np.sum(self.dist_ab * np.log(self.dist_ab), axis=0)
|
||||
|
||||
# def plot_dist_grid(self, h, w):
|
||||
# # Plots distribution at a given point
|
||||
# plt.figure()
|
||||
# plt.imshow(self.dist_ab_grid[:, :, h, w], extent=[-110, 110, 110, -110], interpolation='nearest')
|
||||
# plt.colorbar()
|
||||
# plt.ylabel('a')
|
||||
# plt.xlabel('b')
|
||||
|
||||
# def plot_dist_entropy(self):
|
||||
# # Plots distribution at a given point
|
||||
# plt.figure()
|
||||
# plt.imshow(-self.dist_entropy, interpolation='nearest')
|
||||
# plt.colorbar()
|
||||
|
||||
|
||||
class ColorizeImageCaffe(ColorizeImageBase):
|
||||
def __init__(self, Xd=256):
|
||||
print('ColorizeImageCaffe instantiated')
|
||||
ColorizeImageBase.__init__(self, Xd)
|
||||
self.l_norm = 1.
|
||||
self.ab_norm = 1.
|
||||
self.l_mean = 50.
|
||||
self.ab_mean = 0.
|
||||
self.mask_mult = 110.
|
||||
|
||||
self.pred_ab_layer = 'pred_ab' # predicted ab layer
|
||||
|
||||
# Load grid properties
|
||||
self.pts_in_hull_path = './data/color_bins/pts_in_hull.npy'
|
||||
self.pts_in_hull = np.load(self.pts_in_hull_path) # 313x2, in-gamut
|
||||
|
||||
# ***** Net preparation *****
|
||||
def prep_net(self, gpu_id, prototxt_path='', caffemodel_path=''):
|
||||
import caffe
|
||||
print('gpu_id = %d, net_path = %s, model_path = %s' % (gpu_id, prototxt_path, caffemodel_path))
|
||||
if gpu_id == -1:
|
||||
caffe.set_mode_cpu()
|
||||
else:
|
||||
caffe.set_device(gpu_id)
|
||||
caffe.set_mode_gpu()
|
||||
self.gpu_id = gpu_id
|
||||
self.net = caffe.Net(prototxt_path, caffemodel_path, caffe.TEST)
|
||||
self.net_set = True
|
||||
|
||||
# automatically set cluster centers
|
||||
if len(self.net.params[self.pred_ab_layer][0].data[...].shape) == 4 and self.net.params[self.pred_ab_layer][0].data[...].shape[1] == 313:
|
||||
print('Setting ab cluster centers in layer: %s' % self.pred_ab_layer)
|
||||
self.net.params[self.pred_ab_layer][0].data[:, :, 0, 0] = self.pts_in_hull.T
|
||||
|
||||
# automatically set upsampling kernel
|
||||
for layer in self.net._layer_names:
|
||||
if layer[-3:] == '_us':
|
||||
print('Setting upsampling layer kernel: %s' % layer)
|
||||
self.net.params[layer][0].data[:, 0, :, :] = np.array(((.25, .5, .25, 0), (.5, 1., .5, 0), (.25, .5, .25, 0), (0, 0, 0, 0)))[np.newaxis, :, :]
|
||||
|
||||
# ***** Call forward *****
|
||||
def net_forward(self, input_ab, input_mask):
|
||||
# INPUTS
|
||||
# ab 2xXxX input color patches (non-normalized)
|
||||
# mask 1xXxX input mask, indicating which points have been provided
|
||||
# assumes self.img_l_mc has been set
|
||||
|
||||
if ColorizeImageBase.net_forward(self, input_ab, input_mask) == -1:
|
||||
return -1
|
||||
|
||||
net_input_prepped = np.concatenate((self.img_l_mc, self.input_ab_mc, self.input_mask_mult), axis=0)
|
||||
|
||||
self.net.blobs['data_l_ab_mask'].data[...] = net_input_prepped
|
||||
self.net.forward()
|
||||
|
||||
# return prediction
|
||||
self.output_rgb = lab2rgb_transpose(self.img_l, self.net.blobs[self.pred_ab_layer].data[0, :, :, :])
|
||||
|
||||
self._set_out_ab_()
|
||||
return self.output_rgb
|
||||
|
||||
def get_img_forward(self):
|
||||
# get image with point estimate
|
||||
return self.output_rgb
|
||||
|
||||
def get_img_gray(self):
|
||||
# Get black and white image
|
||||
return lab2rgb_transpose(self.img_l, np.zeros((2, self.Xd, self.Xd)))
|
||||
|
||||
|
||||
class ColorizeImageCaffeGlobDist(ColorizeImageCaffe):
|
||||
# Caffe colorization, with additional global histogram as input
|
||||
def __init__(self, Xd=256):
|
||||
ColorizeImageCaffe.__init__(self, Xd)
|
||||
self.glob_mask_mult = 1.
|
||||
self.glob_layer = 'glob_ab_313_mask'
|
||||
|
||||
def net_forward(self, input_ab, input_mask, glob_dist=-1):
|
||||
# glob_dist is 313 array, or -1
|
||||
if np.array(glob_dist).flatten()[0] == -1: # run without this, zero it out
|
||||
self.net.blobs[self.glob_layer].data[0, :-1, 0, 0] = 0.
|
||||
self.net.blobs[self.glob_layer].data[0, -1, 0, 0] = 0.
|
||||
else: # run conditioned on global histogram
|
||||
self.net.blobs[self.glob_layer].data[0, :-1, 0, 0] = glob_dist
|
||||
self.net.blobs[self.glob_layer].data[0, -1, 0, 0] = self.glob_mask_mult
|
||||
|
||||
self.output_rgb = ColorizeImageCaffe.net_forward(self, input_ab, input_mask)
|
||||
self._set_out_ab_()
|
||||
return self.output_rgb
|
||||
|
||||
|
||||
class ColorizeImageCaffeDist(ColorizeImageCaffe):
|
||||
# caffe model which includes distribution prediction
|
||||
def __init__(self, Xd=256):
|
||||
ColorizeImageCaffe.__init__(self, Xd)
|
||||
self.dist_ab_set = False
|
||||
self.scale_S_layer = 'scale_S'
|
||||
self.dist_ab_S_layer = 'dist_ab_S' # softened distribution layer
|
||||
self.pts_grid = np.load('./data/color_bins/pts_grid.npy') # 529x2, all points
|
||||
self.in_hull = np.load('./data/color_bins/in_hull.npy') # 529 bool
|
||||
self.AB = self.pts_grid.shape[0] # 529
|
||||
self.A = int(np.sqrt(self.AB)) # 23
|
||||
self.B = int(np.sqrt(self.AB)) # 23
|
||||
self.dist_ab_full = np.zeros((self.AB, self.Xd, self.Xd))
|
||||
self.dist_ab_grid = np.zeros((self.A, self.B, self.Xd, self.Xd))
|
||||
self.dist_entropy = np.zeros((self.Xd, self.Xd))
|
||||
|
||||
def prep_net(self, gpu_id, prototxt_path='', caffemodel_path='', S=.2):
|
||||
ColorizeImageCaffe.prep_net(self, gpu_id, prototxt_path=prototxt_path, caffemodel_path=caffemodel_path)
|
||||
self.S = S
|
||||
self.net.params[self.scale_S_layer][0].data[...] = S
|
||||
|
||||
def net_forward(self, input_ab, input_mask):
|
||||
# INPUTS
|
||||
# ab 2xXxX input color patches (non-normalized)
|
||||
# mask 1xXxX input mask, indicating which points have been provided
|
||||
# assumes self.img_l_mc has been set
|
||||
|
||||
function_return = ColorizeImageCaffe.net_forward(self, input_ab, input_mask)
|
||||
if np.array(function_return).flatten()[0] == -1: # errored out
|
||||
return -1
|
||||
|
||||
# set distribution
|
||||
# in-gamut, CxXxX, C = 313
|
||||
self.dist_ab = self.net.blobs[self.dist_ab_S_layer].data[0, :, :, :]
|
||||
self.dist_ab_set = True
|
||||
|
||||
# full grid, ABxXxX, AB = 529
|
||||
self.dist_ab_full[self.in_hull, :, :] = self.dist_ab
|
||||
|
||||
# gridded, AxBxXxX, A = 23
|
||||
self.dist_ab_grid = self.dist_ab_full.reshape((self.A, self.B, self.Xd, self.Xd))
|
||||
|
||||
# return
|
||||
return function_return
|
||||
|
||||
# def get_ab_reccs(self, h, w, K=5, N=25000, return_conf=False):
|
||||
# ''' Recommended colors at point (h,w)
|
||||
# Call this after calling net_forward
|
||||
# '''
|
||||
# if not self.dist_ab_set:
|
||||
# print('Need to set prediction first')
|
||||
# return 0
|
||||
#
|
||||
# # randomly sample from pdf
|
||||
# cmf = np.cumsum(self.dist_ab[:, h, w]) # CMF
|
||||
# cmf = cmf / cmf[-1]
|
||||
# cmf_bins = cmf
|
||||
#
|
||||
# # randomly sample N points
|
||||
# rnd_pts = np.random.uniform(low=0, high=1.0, size=N)
|
||||
# inds = np.digitize(rnd_pts, bins=cmf_bins)
|
||||
# rnd_pts_ab = self.pts_in_hull[inds, :]
|
||||
#
|
||||
# # run k-means
|
||||
# kmeans = KMeans(n_clusters=K).fit(rnd_pts_ab)
|
||||
#
|
||||
# # sort by cluster occupancy
|
||||
# k_label_cnt = np.histogram(kmeans.labels_, np.arange(0, K + 1))[0]
|
||||
# k_inds = np.argsort(k_label_cnt, axis=0)[::-1]
|
||||
#
|
||||
# cluster_per = 1. * k_label_cnt[k_inds] / N # percentage of points within cluster
|
||||
# cluster_centers = kmeans.cluster_centers_[k_inds, :] # cluster centers
|
||||
#
|
||||
# # cluster_centers = np.random.uniform(low=-100,high=100,size=(N,2))
|
||||
# if return_conf:
|
||||
# return cluster_centers, cluster_per
|
||||
# else:
|
||||
# return cluster_centers
|
||||
|
||||
def compute_entropy(self):
|
||||
# compute the distribution entropy (really slow right now)
|
||||
self.dist_entropy = np.sum(self.dist_ab * np.log(self.dist_ab), axis=0)
|
||||
|
||||
# def plot_dist_grid(self, h, w):
|
||||
# Plots distribution at a given point
|
||||
# plt.figure()
|
||||
# plt.imshow(self.dist_ab_grid[:, :, h, w], extent=[-110, 110, 110, -110], interpolation='nearest')
|
||||
# plt.colorbar()
|
||||
# plt.ylabel('a')
|
||||
# plt.xlabel('b')
|
||||
|
||||
# def plot_dist_entropy(self):
|
||||
# Plots distribution at a given point
|
||||
# plt.figure()
|
||||
# plt.imshow(-self.dist_entropy, interpolation='nearest')
|
||||
# plt.colorbar()
|
||||
@ -1,90 +0,0 @@
|
||||
import numpy as np
|
||||
from skimage import color
|
||||
import warnings
|
||||
|
||||
|
||||
def qcolor2lab_1d(qc):
|
||||
# take 1d numpy array and do color conversion
|
||||
c = np.array([qc.red(), qc.green(), qc.blue()], np.uint8)
|
||||
return rgb2lab_1d(c)
|
||||
|
||||
|
||||
def rgb2lab_1d(in_rgb):
|
||||
# take 1d numpy array and do color conversion
|
||||
# print('in_rgb', in_rgb)
|
||||
return color.rgb2lab(in_rgb[np.newaxis, np.newaxis, :]).flatten()
|
||||
|
||||
|
||||
def lab2rgb_1d(in_lab, clip=True, dtype='uint8'):
|
||||
warnings.filterwarnings("ignore")
|
||||
tmp_rgb = color.lab2rgb(in_lab[np.newaxis, np.newaxis, :]).flatten()
|
||||
if clip:
|
||||
tmp_rgb = np.clip(tmp_rgb, 0, 1)
|
||||
if dtype == 'uint8':
|
||||
tmp_rgb = np.round(tmp_rgb * 255).astype('uint8')
|
||||
return tmp_rgb
|
||||
|
||||
|
||||
def snap_ab(input_l, input_rgb, return_type='rgb'):
|
||||
''' given an input lightness and rgb, snap the color into a region where l,a,b is in-gamut
|
||||
'''
|
||||
T = 20
|
||||
warnings.filterwarnings("ignore")
|
||||
input_lab = rgb2lab_1d(np.array(input_rgb)) # convert input to lab
|
||||
conv_lab = input_lab.copy() # keep ab from input
|
||||
for t in range(T):
|
||||
conv_lab[0] = input_l # overwrite input l with input ab
|
||||
old_lab = conv_lab
|
||||
tmp_rgb = color.lab2rgb(conv_lab[np.newaxis, np.newaxis, :]).flatten()
|
||||
tmp_rgb = np.clip(tmp_rgb, 0, 1)
|
||||
conv_lab = color.rgb2lab(tmp_rgb[np.newaxis, np.newaxis, :]).flatten()
|
||||
dif_lab = np.sum(np.abs(conv_lab - old_lab))
|
||||
if dif_lab < 1:
|
||||
break
|
||||
# print(conv_lab)
|
||||
|
||||
conv_rgb_ingamut = lab2rgb_1d(conv_lab, clip=True, dtype='uint8')
|
||||
if (return_type == 'rgb'):
|
||||
return conv_rgb_ingamut
|
||||
|
||||
elif(return_type == 'lab'):
|
||||
conv_lab_ingamut = rgb2lab_1d(conv_rgb_ingamut)
|
||||
return conv_lab_ingamut
|
||||
|
||||
|
||||
class abGrid():
|
||||
def __init__(self, gamut_size=110, D=1):
|
||||
self.D = D
|
||||
self.vals_b, self.vals_a = np.meshgrid(np.arange(-gamut_size, gamut_size + D, D),
|
||||
np.arange(-gamut_size, gamut_size + D, D))
|
||||
self.pts_full_grid = np.concatenate((self.vals_a[:, :, np.newaxis], self.vals_b[:, :, np.newaxis]), axis=2)
|
||||
self.A = self.pts_full_grid.shape[0]
|
||||
self.B = self.pts_full_grid.shape[1]
|
||||
self.AB = self.A * self.B
|
||||
self.gamut_size = gamut_size
|
||||
|
||||
def update_gamut(self, l_in):
|
||||
warnings.filterwarnings("ignore")
|
||||
thresh = 1.0
|
||||
pts_lab = np.concatenate((l_in + np.zeros((self.A, self.B, 1)), self.pts_full_grid), axis=2)
|
||||
self.pts_rgb = (255 * np.clip(color.lab2rgb(pts_lab), 0, 1)).astype('uint8')
|
||||
pts_lab_back = color.rgb2lab(self.pts_rgb)
|
||||
pts_lab_diff = np.linalg.norm(pts_lab - pts_lab_back, axis=2)
|
||||
|
||||
self.mask = pts_lab_diff < thresh
|
||||
mask3 = np.tile(self.mask[..., np.newaxis], [1, 1, 3])
|
||||
self.masked_rgb = self.pts_rgb.copy()
|
||||
self.masked_rgb[np.invert(mask3)] = 255
|
||||
return self.masked_rgb, self.mask
|
||||
|
||||
def ab2xy(self, a, b):
|
||||
y = self.gamut_size + a
|
||||
x = self.gamut_size + b
|
||||
# print('ab2xy (%d, %d) -> (%d, %d)' % (a, b, x, y))
|
||||
return x, y
|
||||
|
||||
def xy2ab(self, x, y):
|
||||
a = y - self.gamut_size
|
||||
b = x - self.gamut_size
|
||||
# print('xy2ab (%d, %d) -> (%d, %d)' % (x, y, a, b))
|
||||
return a, b
|
||||
@ -1,177 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class SIGGRAPHGenerator(nn.Module):
|
||||
def __init__(self, dist=False):
|
||||
super(SIGGRAPHGenerator, self).__init__()
|
||||
self.dist = dist
|
||||
use_bias = True
|
||||
norm_layer = nn.BatchNorm2d
|
||||
|
||||
# Conv1
|
||||
model1 = [nn.Conv2d(4, 64, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model1 += [nn.ReLU(True), ]
|
||||
model1 += [nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model1 += [nn.ReLU(True), ]
|
||||
model1 += [norm_layer(64), ]
|
||||
# add a subsampling operation
|
||||
|
||||
# Conv2
|
||||
model2 = [nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model2 += [nn.ReLU(True), ]
|
||||
model2 += [nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model2 += [nn.ReLU(True), ]
|
||||
model2 += [norm_layer(128), ]
|
||||
# add a subsampling layer operation
|
||||
|
||||
# Conv3
|
||||
model3 = [nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model3 += [nn.ReLU(True), ]
|
||||
model3 += [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model3 += [nn.ReLU(True), ]
|
||||
model3 += [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model3 += [nn.ReLU(True), ]
|
||||
model3 += [norm_layer(256), ]
|
||||
# add a subsampling layer operation
|
||||
|
||||
# Conv4
|
||||
model4 = [nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model4 += [nn.ReLU(True), ]
|
||||
model4 += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model4 += [nn.ReLU(True), ]
|
||||
model4 += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model4 += [nn.ReLU(True), ]
|
||||
model4 += [norm_layer(512), ]
|
||||
|
||||
# Conv5
|
||||
model5 = [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model5 += [nn.ReLU(True), ]
|
||||
model5 += [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model5 += [nn.ReLU(True), ]
|
||||
model5 += [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model5 += [nn.ReLU(True), ]
|
||||
model5 += [norm_layer(512), ]
|
||||
|
||||
# Conv6
|
||||
model6 = [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model6 += [nn.ReLU(True), ]
|
||||
model6 += [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model6 += [nn.ReLU(True), ]
|
||||
model6 += [nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=use_bias), ]
|
||||
model6 += [nn.ReLU(True), ]
|
||||
model6 += [norm_layer(512), ]
|
||||
|
||||
# Conv7
|
||||
model7 = [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model7 += [nn.ReLU(True), ]
|
||||
model7 += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model7 += [nn.ReLU(True), ]
|
||||
model7 += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model7 += [nn.ReLU(True), ]
|
||||
model7 += [norm_layer(512), ]
|
||||
|
||||
# Conv7
|
||||
model8up = [nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=use_bias)]
|
||||
model3short8 = [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
|
||||
model8 = [nn.ReLU(True), ]
|
||||
model8 += [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model8 += [nn.ReLU(True), ]
|
||||
model8 += [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model8 += [nn.ReLU(True), ]
|
||||
model8 += [norm_layer(256), ]
|
||||
|
||||
# Conv9
|
||||
model9up = [nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=use_bias), ]
|
||||
model2short9 = [nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
# add the two feature maps above
|
||||
|
||||
model9 = [nn.ReLU(True), ]
|
||||
model9 += [nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
model9 += [nn.ReLU(True), ]
|
||||
model9 += [norm_layer(128), ]
|
||||
|
||||
# Conv10
|
||||
model10up = [nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=use_bias), ]
|
||||
model1short10 = [nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=use_bias), ]
|
||||
# add the two feature maps above
|
||||
|
||||
model10 = [nn.ReLU(True), ]
|
||||
model10 += [nn.Conv2d(128, 128, kernel_size=3, dilation=1, stride=1, padding=1, bias=use_bias), ]
|
||||
model10 += [nn.LeakyReLU(negative_slope=.2), ]
|
||||
|
||||
# classification output
|
||||
model_class = [nn.Conv2d(256, 529, kernel_size=1, padding=0, dilation=1, stride=1, bias=use_bias), ]
|
||||
|
||||
# regression output
|
||||
model_out = [nn.Conv2d(128, 2, kernel_size=1, padding=0, dilation=1, stride=1, bias=use_bias), ]
|
||||
model_out += [nn.Tanh()]
|
||||
|
||||
self.model1 = nn.Sequential(*model1)
|
||||
self.model2 = nn.Sequential(*model2)
|
||||
self.model3 = nn.Sequential(*model3)
|
||||
self.model4 = nn.Sequential(*model4)
|
||||
self.model5 = nn.Sequential(*model5)
|
||||
self.model6 = nn.Sequential(*model6)
|
||||
self.model7 = nn.Sequential(*model7)
|
||||
self.model8up = nn.Sequential(*model8up)
|
||||
self.model8 = nn.Sequential(*model8)
|
||||
self.model9up = nn.Sequential(*model9up)
|
||||
self.model9 = nn.Sequential(*model9)
|
||||
self.model10up = nn.Sequential(*model10up)
|
||||
self.model10 = nn.Sequential(*model10)
|
||||
self.model3short8 = nn.Sequential(*model3short8)
|
||||
self.model2short9 = nn.Sequential(*model2short9)
|
||||
self.model1short10 = nn.Sequential(*model1short10)
|
||||
|
||||
self.model_class = nn.Sequential(*model_class)
|
||||
self.model_out = nn.Sequential(*model_out)
|
||||
|
||||
self.upsample4 = nn.Sequential(*[nn.Upsample(scale_factor=4, mode='nearest'), ])
|
||||
self.softmax = nn.Sequential(*[nn.Softmax(dim=1), ])
|
||||
|
||||
def forward(self, input_A, input_B, mask_B, maskcent=0,f=False):
|
||||
# input_A \in [-50,+50]
|
||||
# input_B \in [-110, +110]
|
||||
# mask_B \in [0, +1.0]
|
||||
|
||||
input_A = torch.Tensor(input_A)[None, :, :, :]
|
||||
input_B = torch.Tensor(input_B)[None, :, :, :]
|
||||
mask_B = torch.Tensor(mask_B)[None, :, :, :]
|
||||
|
||||
if torch.cuda.is_available() and not f:
|
||||
input_A = input_A.cuda()
|
||||
input_B = input_B.cuda()
|
||||
mask_B = mask_B.cuda()
|
||||
|
||||
mask_B = mask_B - maskcent
|
||||
|
||||
conv1_2 = self.model1(torch.cat((input_A / 100., input_B / 110., mask_B), dim=1))
|
||||
conv2_2 = self.model2(conv1_2[:, :, ::2, ::2])
|
||||
conv3_3 = self.model3(conv2_2[:, :, ::2, ::2])
|
||||
conv4_3 = self.model4(conv3_3[:, :, ::2, ::2])
|
||||
conv5_3 = self.model5(conv4_3)
|
||||
conv6_3 = self.model6(conv5_3)
|
||||
conv7_3 = self.model7(conv6_3)
|
||||
|
||||
conv8_up = self.model8up(conv7_3) + self.model3short8(conv3_3)
|
||||
conv8_3 = self.model8(conv8_up)
|
||||
|
||||
if(self.dist):
|
||||
out_cl = self.upsample4(self.softmax(self.model_class(conv8_3) * .2))
|
||||
|
||||
conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
|
||||
conv9_3 = self.model9(conv9_up)
|
||||
conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
|
||||
conv10_2 = self.model10(conv10_up)
|
||||
out_reg = self.model_out(conv10_2) * 110
|
||||
|
||||
return (out_reg * 110, out_cl)
|
||||
else:
|
||||
conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
|
||||
conv9_3 = self.model9(conv9_up)
|
||||
conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
|
||||
conv10_2 = self.model10(conv10_up)
|
||||
out_reg = self.model_out(conv10_2)
|
||||
return out_reg * 110
|
||||
@ -1,191 +0,0 @@
|
||||
from __future__ import division
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools',baseLoc+'Inpainting'])
|
||||
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch import nn
|
||||
import scipy.ndimage
|
||||
import cv2
|
||||
from DFNet_core import DFNet
|
||||
from RefinementNet_core import RefinementNet
|
||||
|
||||
|
||||
|
||||
def channelData(layer):#convert gimp image to numpy
|
||||
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
||||
pixChars=region[:,:] # Take whole layer
|
||||
bpp=region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
||||
|
||||
def createResultLayer(image,name,result):
|
||||
rlBytes=np.uint8(result).tobytes();
|
||||
rl=gimp.Layer(image,name,image.width,image.height,0,100,NORMAL_MODE)
|
||||
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
||||
region[:,:]=rlBytes
|
||||
image.add_layer(rl,0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def to_numpy(tensor):
|
||||
tensor = tensor.mul(255).byte().data.cpu().numpy()
|
||||
tensor = np.transpose(tensor, [0, 2, 3, 1])
|
||||
return tensor
|
||||
|
||||
def padding(img, height=512, width=512, channels=3):
|
||||
channels = img.shape[2] if len(img.shape) > 2 else 1
|
||||
interpolation=cv2.INTER_NEAREST
|
||||
|
||||
if channels == 1:
|
||||
img_padded = np.zeros((height, width), dtype=img.dtype)
|
||||
else:
|
||||
img_padded = np.zeros((height, width, channels), dtype=img.dtype)
|
||||
|
||||
original_shape = img.shape
|
||||
rows_rate = original_shape[0] / height
|
||||
cols_rate = original_shape[1] / width
|
||||
new_cols = width
|
||||
new_rows = height
|
||||
if rows_rate > cols_rate:
|
||||
new_cols = (original_shape[1] * height) // original_shape[0]
|
||||
img = cv2.resize(img, (new_cols, height), interpolation=interpolation)
|
||||
if new_cols > width:
|
||||
new_cols = width
|
||||
img_padded[:, ((img_padded.shape[1] - new_cols) // 2):((img_padded.shape[1] - new_cols) // 2 + new_cols)] = img
|
||||
else:
|
||||
new_rows = (original_shape[0] * width) // original_shape[1]
|
||||
img = cv2.resize(img, (width, new_rows), interpolation=interpolation)
|
||||
if new_rows > height:
|
||||
new_rows = height
|
||||
img_padded[((img_padded.shape[0] - new_rows) // 2):((img_padded.shape[0] - new_rows) // 2 + new_rows), :] = img
|
||||
return img_padded, new_cols, new_rows
|
||||
|
||||
|
||||
|
||||
def preprocess_image_dfnet(image, mask, model,device):
|
||||
image, new_cols, new_rows = padding(image, 512, 512)
|
||||
mask, _, _ = padding(mask, 512, 512)
|
||||
image = np.ascontiguousarray(image.transpose(2, 0, 1)).astype(np.uint8)
|
||||
mask = np.ascontiguousarray(np.expand_dims(mask, 0)).astype(np.uint8)
|
||||
|
||||
image = torch.from_numpy(image).to(device).float().div(255)
|
||||
mask = 1 - torch.from_numpy(mask).to(device).float().div(255)
|
||||
image_miss = image * mask
|
||||
DFNET_output = model(image_miss.unsqueeze(0), mask.unsqueeze(0))[0]
|
||||
DFNET_output = image * mask + DFNET_output * (1 - mask)
|
||||
DFNET_output = to_numpy(DFNET_output)[0]
|
||||
DFNET_output = cv2.cvtColor(DFNET_output, cv2.COLOR_BGR2RGB)
|
||||
DFNET_output = DFNET_output[(DFNET_output.shape[0] - new_rows) // 2: (DFNET_output.shape[0] - new_rows) // 2 + new_rows,
|
||||
(DFNET_output.shape[1] - new_cols) // 2: (DFNET_output.shape[1] - new_cols) // 2 + new_cols, ...]
|
||||
|
||||
return DFNET_output
|
||||
|
||||
|
||||
|
||||
def preprocess_image(image, mask, image_before_resize, model,device):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
|
||||
shift_val = (100 / 512) * image.shape[0]
|
||||
|
||||
image_resized = cv2.resize(image_before_resize, (image.shape[1], image.shape[0]))
|
||||
|
||||
mask = mask // 255
|
||||
image_matched = image * (1 - mask) + image_resized * mask
|
||||
mask = mask * 255
|
||||
|
||||
img_1 = scipy.ndimage.shift(image_matched, (-shift_val, 0, 0), order=0, mode='constant', cval=1)
|
||||
mask_1 = scipy.ndimage.shift(mask, (-shift_val, 0, 0), order=0, mode='constant', cval=255)
|
||||
img_2 = scipy.ndimage.shift(image_matched, (shift_val, 0, 0), order=0, mode='constant', cval=1)
|
||||
mask_2 = scipy.ndimage.shift(mask, (shift_val, 0, 0), order=0, mode='constant', cval=255)
|
||||
img_3 = scipy.ndimage.shift(image_matched, (0, shift_val, 0), order=0, mode='constant', cval=1)
|
||||
mask_3 = scipy.ndimage.shift(mask, (0, shift_val, 0), order=0, mode='constant', cval=255)
|
||||
img_4 = scipy.ndimage.shift(image_matched, (0, -shift_val, 0), order=0, mode='constant', cval=1)
|
||||
mask_4 = scipy.ndimage.shift(mask, (0, -shift_val, 0), order=0, mode='constant', cval=255)
|
||||
image_cat = np.dstack((mask, image_matched, img_1, mask_1, img_2, mask_2, img_3, mask_3, img_4, mask_4))
|
||||
|
||||
mask_patch = torch.from_numpy(image_cat).to(device).float().div(255).unsqueeze(0)
|
||||
mask_patch = mask_patch.permute(0, -1, 1, 2)
|
||||
inputs = mask_patch[:, 1:, ...]
|
||||
mask = mask_patch[:, 0:1, ...]
|
||||
out = model(inputs, mask)
|
||||
out = out.mul(255).byte().data.cpu().numpy()
|
||||
out = np.transpose(out, [0, 2, 3, 1])[0]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def pad_image(image):
|
||||
x = ((image.shape[0] // 256) + (1 if image.shape[0] % 256 != 0 else 0)) * 256
|
||||
y = ((image.shape[1] // 256) + (1 if image.shape[1] % 256 != 0 else 0)) * 256
|
||||
padded = np.zeros((x, y, image.shape[2]), dtype='uint8')
|
||||
padded[:image.shape[0], :image.shape[1], ...] = image
|
||||
return padded
|
||||
|
||||
|
||||
def inpaint(imggimp, curlayer,layeri,layerm,cFlag) :
|
||||
|
||||
img = channelData(layeri)[..., :3]
|
||||
mask = channelData(layerm)[..., :3]
|
||||
|
||||
if img.shape[0] != imggimp.height or img.shape[1] != imggimp.width or mask.shape[0] != imggimp.height or mask.shape[1] != imggimp.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Running inpainting for " + layeri.name + "...")
|
||||
device = torch.device('cuda')
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running inpainting for " + layeri.name + "...")
|
||||
device = torch.device('cpu')
|
||||
|
||||
assert img.shape[:2] == mask.shape[:2]
|
||||
|
||||
mask = mask[..., :1]
|
||||
|
||||
|
||||
image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||||
shape = image.shape
|
||||
|
||||
image = pad_image(image)
|
||||
mask = pad_image(mask)
|
||||
|
||||
DFNet_model = DFNet().to(device)
|
||||
DFNet_model.load_state_dict(torch.load(baseLoc + '/weights/inpainting/model_places2.pth', map_location=device))
|
||||
DFNet_model.eval()
|
||||
DFNET_output = preprocess_image_dfnet(image, mask, DFNet_model,device)
|
||||
del DFNet_model
|
||||
Refinement_model = RefinementNet().to(device)
|
||||
Refinement_model.load_state_dict(torch.load(baseLoc+'/weights/inpainting/refinement.pth', map_location=device)['state_dict'])
|
||||
Refinement_model.eval()
|
||||
out = preprocess_image(image, mask, DFNET_output, Refinement_model,device)
|
||||
out = out[:shape[0], :shape[1], ...]
|
||||
del Refinement_model
|
||||
createResultLayer(imggimp,'output',out)
|
||||
|
||||
|
||||
|
||||
register(
|
||||
"inpainting",
|
||||
"inpainting",
|
||||
"Running inpainting.",
|
||||
"Andrey Moskalenko",
|
||||
"Your",
|
||||
"2020",
|
||||
"inpainting...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_LAYER, "drawinglayer", "Image:", None),
|
||||
(PF_LAYER, "drawinglayer", "Mask:", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
inpaint, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,46 +0,0 @@
|
||||
if [ ! -d "gimpenv" ]; then
|
||||
|
||||
echo "\n-----------Installing GIMP-ML-----------\n"
|
||||
|
||||
if [ "$(uname)" == "Linux" ]; then
|
||||
if [[ $(lsb_release -rs) == "18.04" ]]; then #for ubuntu 18.04
|
||||
sudo apt install python-minimal
|
||||
alias python='python2'
|
||||
|
||||
elif [[ $(lsb_release -rs) == "20.04" ]]; then #for ubuntu 20.04
|
||||
sudo apt install python2-minimal
|
||||
wget https://bootstrap.pypa.io/get-pip.py
|
||||
alias python='python2'
|
||||
python get-pip.py
|
||||
sudo apt-get install libpython2.7
|
||||
|
||||
elif [[ $(lsb_release -rs) == "10" ]]; then #for debian 10
|
||||
wget https://bootstrap.pypa.io/get-pip.py
|
||||
python get-pip.py
|
||||
fi
|
||||
fi
|
||||
|
||||
python -m pip install --user virtualenv
|
||||
python -m virtualenv gimpenv
|
||||
source gimpenv/bin/activate
|
||||
python -m pip install torchvision
|
||||
python -m pip install "opencv-python<=4.3"
|
||||
python -m pip install numpy
|
||||
python -m pip install future
|
||||
python -m pip install torch
|
||||
python -m pip install scipy
|
||||
python -m pip install typing
|
||||
python -m pip install enum
|
||||
python -m pip install pretrainedmodels
|
||||
python -m pip install requests
|
||||
deactivate
|
||||
|
||||
echo "\n-----------Installed GIMP-ML------------\n"
|
||||
|
||||
else
|
||||
|
||||
echo "Environment already setup!"
|
||||
|
||||
fi
|
||||
|
||||
|
||||
@ -1,130 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
savePath = '/'.join(baseLoc.split('/')[:-2]) + '/output/interpolateframes'
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'RIFE'])
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from model import RIFE
|
||||
import numpy as np
|
||||
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def getinter(img_s, img_e, c_flag, string_path):
|
||||
exp = 4
|
||||
out_path = string_path
|
||||
|
||||
model = RIFE.Model(c_flag)
|
||||
model.load_model(baseLoc + 'weights' + '/interpolateframes')
|
||||
model.eval()
|
||||
model.device(c_flag)
|
||||
|
||||
img0 = img_s
|
||||
img1 = img_e
|
||||
|
||||
img0 = (torch.tensor(img0.transpose(2, 0, 1)) / 255.).unsqueeze(0)
|
||||
img1 = (torch.tensor(img1.transpose(2, 0, 1))/ 255.).unsqueeze(0)
|
||||
if torch.cuda.is_available() and not c_flag:
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
img0 = img0.to(device)
|
||||
img1 = img1.to(device)
|
||||
|
||||
n, c, h, w = img0.shape
|
||||
ph = ((h - 1) // 32 + 1) * 32
|
||||
pw = ((w - 1) // 32 + 1) * 32
|
||||
padding = (0, pw - w, 0, ph - h)
|
||||
img0 = F.pad(img0, padding)
|
||||
img1 = F.pad(img1, padding)
|
||||
|
||||
img_list = [img0, img1]
|
||||
idx=0
|
||||
t=exp * (len(img_list) - 1)
|
||||
for i in range(exp):
|
||||
tmp = []
|
||||
for j in range(len(img_list) - 1):
|
||||
mid = model.inference(img_list[j], img_list[j + 1])
|
||||
tmp.append(img_list[j])
|
||||
tmp.append(mid)
|
||||
idx=idx+1
|
||||
gimp.progress_update(float(idx)/float(t))
|
||||
gimp.displays_flush()
|
||||
tmp.append(img1)
|
||||
img_list = tmp
|
||||
|
||||
if not os.path.exists(out_path):
|
||||
os.makedirs(out_path)
|
||||
for i in range(len(img_list)):
|
||||
cv2.imwrite(out_path + '/img{}.png'.format(i),
|
||||
(img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w, ::-1])
|
||||
|
||||
|
||||
|
||||
|
||||
def interpolateframes(imggimp, curlayer, string_path, layer_s, layer_e, c_flag):
|
||||
layer_1 = channelData(layer_s)
|
||||
layer_2 = channelData(layer_e)
|
||||
|
||||
if layer_1.shape[0] != imggimp.height or layer_1.shape[1] != imggimp.width or layer_2.shape[0] != imggimp.height or layer_2.shape[1] != imggimp.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) for both layers and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not c_flag:
|
||||
gimp.progress_init("(Using GPU) Running slomo and saving frames in "+string_path)
|
||||
# device = torch.device("cuda")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running slomo and saving frames in "+string_path)
|
||||
# device = torch.device("cpu")
|
||||
|
||||
if layer_1.shape[2] == 4: # get rid of alpha channel
|
||||
layer_1 = layer_1[:, :, 0:3]
|
||||
if layer_2.shape[2] == 4: # get rid of alpha channel
|
||||
layer_2 = layer_2[:, :, 0:3]
|
||||
getinter(layer_1, layer_2, c_flag, string_path)
|
||||
# pdb.gimp_message("Saved")
|
||||
|
||||
|
||||
register(
|
||||
"interpolate-frames",
|
||||
"interpolate-frames",
|
||||
"Running slomo...",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"interpolate-frames...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_STRING, "string", "Output Location", savePath),
|
||||
(PF_LAYER, "drawinglayer", "Start Frame:", None),
|
||||
(PF_LAYER, "drawinglayer", "End Frame:", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
interpolateframes, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,25 +0,0 @@
|
||||
|
||||
from gimpfu import *
|
||||
|
||||
def invert(img, layer) :
|
||||
gimp.progress_init("Inverting " + layer.name + "...")
|
||||
pdb.gimp_undo_push_group_start(img)
|
||||
pdb.gimp_invert(layer)
|
||||
pdb.gimp_undo_push_group_end(img)
|
||||
|
||||
register(
|
||||
"Invert",
|
||||
"Invert",
|
||||
"Invert",
|
||||
"Kritik Soman",
|
||||
"Your Name",
|
||||
"2020",
|
||||
"Invert...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
],
|
||||
[],
|
||||
invert, menu="<Image>/Filters/Enhance")
|
||||
|
||||
main()
|
||||
@ -1,79 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools'])
|
||||
|
||||
import numpy as np
|
||||
from scipy.cluster.vq import kmeans2
|
||||
|
||||
|
||||
def channelData(layer):#convert gimp image to numpy
|
||||
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
||||
pixChars=region[:,:] # Take whole layer
|
||||
bpp=region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
||||
|
||||
|
||||
def createResultLayer(image,name,result):
|
||||
rlBytes=np.uint8(result).tobytes();
|
||||
rl=gimp.Layer(image,name,image.width,image.height,0,100,NORMAL_MODE)#1 is for RGB with alpha
|
||||
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
||||
region[:,:]=rlBytes
|
||||
image.add_layer(rl,0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def kmeans(imggimp, curlayer,layeri,n_clusters,locflag) :
|
||||
image = channelData(layeri)
|
||||
if image.shape[0] != imggimp.height or image.shape[1] != imggimp.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if image.shape[2] == 4: # get rid of alpha channel
|
||||
image = image[:,:,0:3]
|
||||
h,w,d = image.shape
|
||||
# reshape the image to a 2D array of pixels and 3 color values (RGB)
|
||||
pixel_values = image.reshape((-1, 3))
|
||||
|
||||
if locflag:
|
||||
xx,yy = np.meshgrid(range(w),range(h))
|
||||
x = xx.reshape(-1,1)
|
||||
y = yy.reshape(-1,1)
|
||||
pixel_values = np.concatenate((pixel_values,x,y),axis=1)
|
||||
|
||||
pixel_values = np.float32(pixel_values)
|
||||
c,out = kmeans2(pixel_values,n_clusters)
|
||||
|
||||
if locflag:
|
||||
c = np.uint8(c[:,0:3])
|
||||
else:
|
||||
c = np.uint8(c)
|
||||
|
||||
segmented_image = c[out.flatten()]
|
||||
segmented_image = segmented_image.reshape((h,w,d))
|
||||
createResultLayer(imggimp,'new_output',segmented_image)
|
||||
|
||||
|
||||
register(
|
||||
"kmeans",
|
||||
"kmeans clustering",
|
||||
"Running kmeans clustering.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"kmeans...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_LAYER, "drawinglayer", "Original Image", None),
|
||||
(PF_INT, "depth", "Number of clusters", 3),
|
||||
(PF_BOOL, "position", "Use position", False)
|
||||
|
||||
],
|
||||
[],
|
||||
kmeans, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,77 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'MiDaS'])
|
||||
|
||||
from run import run_depth
|
||||
from monodepth_net import MonoDepthNet
|
||||
import MiDaS_utils as MiDaS_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
def getMonoDepth(input_image,cFlag):
|
||||
image = input_image / 255.0
|
||||
out = run_depth(image, baseLoc+'weights/MiDaS/model.pt', MonoDepthNet, MiDaS_utils, target_w=640,f=cFlag)
|
||||
out = np.repeat(out[:, :, np.newaxis], 3, axis=2)
|
||||
d1,d2 = input_image.shape[:2]
|
||||
out = cv2.resize(out,(d2,d1))
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
|
||||
def MonoDepth(img, layer,cFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Generating disparity map for " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Generating disparity map for " + layer.name + "...")
|
||||
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:,:,0:3]
|
||||
cpy = getMonoDepth(imgmat,cFlag)
|
||||
createResultLayer(img, 'new_output', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"MonoDepth",
|
||||
"MonoDepth",
|
||||
"Generate monocular disparity map based on deep learning.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"MonoDepth...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
MonoDepth, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,194 +0,0 @@
|
||||
import torch
|
||||
import cv2
|
||||
import os
|
||||
import random
|
||||
import numpy as np
|
||||
from torchvision import transforms
|
||||
import logging
|
||||
|
||||
def gen_trimap(alpha):
|
||||
k_size = random.choice(range(2, 5))
|
||||
iterations = np.random.randint(5, 15)
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k_size, k_size))
|
||||
dilated = cv2.dilate(alpha, kernel, iterations=iterations)
|
||||
eroded = cv2.erode(alpha, kernel, iterations=iterations)
|
||||
trimap = np.zeros(alpha.shape)
|
||||
trimap.fill(128)
|
||||
#trimap[alpha >= 255] = 255
|
||||
trimap[eroded >= 255] = 255
|
||||
trimap[dilated <= 0] = 0
|
||||
|
||||
'''
|
||||
alpha_unknown = alpha[trimap == 128]
|
||||
num_all = alpha_unknown.size
|
||||
num_0 = (alpha_unknown == 0).sum()
|
||||
num_1 = (alpha_unknown == 255).sum()
|
||||
print("Debug: 0 : {}/{} {:.3f}".format(num_0, num_all, float(num_0)/num_all))
|
||||
print("Debug: 255: {}/{} {:.3f}".format(num_1, num_all, float(num_1)/num_all))
|
||||
'''
|
||||
|
||||
return trimap
|
||||
|
||||
|
||||
def compute_gradient(img):
|
||||
x = cv2.Sobel(img, cv2.CV_16S, 1, 0)
|
||||
y = cv2.Sobel(img, cv2.CV_16S, 0, 1)
|
||||
absX = cv2.convertScaleAbs(x)
|
||||
absY = cv2.convertScaleAbs(y)
|
||||
grad = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
|
||||
grad=cv2.cvtColor(grad, cv2.COLOR_BGR2GRAY)
|
||||
return grad
|
||||
|
||||
|
||||
class MatTransform(object):
|
||||
def __init__(self, flip=False):
|
||||
self.flip = flip
|
||||
|
||||
def __call__(self, img, alpha, fg, bg, crop_h, crop_w):
|
||||
h, w = alpha.shape
|
||||
|
||||
# trimap is dilated maybe choose some bg region(0)
|
||||
# random crop in the unknown region center
|
||||
target = np.where((alpha > 0) & (alpha < 255))
|
||||
delta_h = center_h = crop_h / 2
|
||||
delta_w = center_w = crop_w / 2
|
||||
|
||||
if len(target[0]) > 0:
|
||||
rand_ind = np.random.randint(len(target[0]))
|
||||
center_h = min(max(target[0][rand_ind], delta_h), h - delta_h)
|
||||
center_w = min(max(target[1][rand_ind], delta_w), w - delta_w)
|
||||
|
||||
# choose unknown point as center not as left-top
|
||||
start_h = int(center_h - delta_h)
|
||||
start_w = int(center_w - delta_w)
|
||||
end_h = int(center_h + delta_h)
|
||||
end_w = int(center_w + delta_w)
|
||||
|
||||
#print("Debug: center({},{}) start({},{}) end({},{}) alpha:{} alpha-len:{} unknown-len:{}".format(center_h, center_w, start_h, start_w, end_h, end_w, alpha[int(center_h), int(center_w)], alpha.size, len(target[0])))
|
||||
|
||||
img = img [start_h : end_h, start_w : end_w]
|
||||
fg = fg [start_h : end_h, start_w : end_w]
|
||||
bg = bg [start_h : end_h, start_w : end_w]
|
||||
alpha = alpha [start_h : end_h, start_w : end_w]
|
||||
|
||||
# random flip
|
||||
if self.flip and random.random() < 0.5:
|
||||
img = cv2.flip(img, 1)
|
||||
alpha = cv2.flip(alpha, 1)
|
||||
fg = cv2.flip(fg, 1)
|
||||
bg = cv2.flip(bg, 1)
|
||||
|
||||
return img, alpha, fg, bg
|
||||
|
||||
|
||||
def get_files(mydir):
|
||||
res = []
|
||||
for root, dirs, files in os.walk(mydir, followlinks=True):
|
||||
for f in files:
|
||||
if f.endswith(".jpg") or f.endswith(".png") or f.endswith(".jpeg") or f.endswith(".JPG"):
|
||||
res.append(os.path.join(root, f))
|
||||
return res
|
||||
|
||||
|
||||
# Dataset not composite online
|
||||
class MatDatasetOffline(torch.utils.data.Dataset):
|
||||
def __init__(self, args, transform=None, normalize=None):
|
||||
self.samples=[]
|
||||
self.transform = transform
|
||||
self.normalize = normalize
|
||||
self.args = args
|
||||
self.size_h = args.size_h
|
||||
self.size_w = args.size_w
|
||||
self.crop_h = args.crop_h
|
||||
self.crop_w = args.crop_w
|
||||
self.logger = logging.getLogger("DeepImageMatting")
|
||||
assert(len(self.crop_h) == len(self.crop_w))
|
||||
|
||||
fg_paths = get_files(self.args.fgDir)
|
||||
|
||||
self.cnt = len(fg_paths)
|
||||
|
||||
for fg_path in fg_paths:
|
||||
alpha_path = fg_path.replace(self.args.fgDir, self.args.alphaDir)
|
||||
img_path = fg_path.replace(self.args.fgDir, self.args.imgDir)
|
||||
bg_path = fg_path.replace(self.args.fgDir, self.args.bgDir)
|
||||
assert(os.path.exists(alpha_path))
|
||||
assert(os.path.exists(fg_path))
|
||||
assert(os.path.exists(bg_path))
|
||||
assert(os.path.exists(img_path))
|
||||
self.samples.append((alpha_path, fg_path, bg_path, img_path))
|
||||
self.logger.info("MatDatasetOffline Samples: {}".format(self.cnt))
|
||||
assert(self.cnt > 0)
|
||||
|
||||
def __getitem__(self,index):
|
||||
alpha_path, fg_path, bg_path, img_path = self.samples[index]
|
||||
|
||||
img_info = [fg_path, alpha_path, bg_path, img_path]
|
||||
|
||||
# read fg, alpha
|
||||
fg = cv2.imread(fg_path)[:, :, :3]
|
||||
bg = cv2.imread(bg_path)[:, :, :3]
|
||||
img = cv2.imread(img_path)[:, :, :3]
|
||||
alpha = cv2.imread(alpha_path)[:, :, 0]
|
||||
|
||||
assert(bg.shape == fg.shape and bg.shape == img.shape)
|
||||
img_info.append(fg.shape)
|
||||
bh, bw, bc, = fg.shape
|
||||
|
||||
rand_ind = random.randint(0, len(self.crop_h) - 1)
|
||||
cur_crop_h = self.crop_h[rand_ind]
|
||||
cur_crop_w = self.crop_w[rand_ind]
|
||||
|
||||
# if ratio!=1: make the img (h==croph and w>=cropw)or(w==cropw and h>=croph)
|
||||
wratio = float(cur_crop_w) / bw
|
||||
hratio = float(cur_crop_h) / bh
|
||||
ratio = wratio if wratio > hratio else hratio
|
||||
if ratio > 1:
|
||||
nbw = int(bw * ratio + 1.0)
|
||||
nbh = int(bh * ratio + 1.0)
|
||||
fg = cv2.resize(fg, (nbw, nbh), interpolation=cv2.INTER_LINEAR)
|
||||
bg = cv2.resize(bg, (nbw, nbh), interpolation=cv2.INTER_LINEAR)
|
||||
img = cv2.resize(img, (nbw, nbh), interpolation=cv2.INTER_LINEAR)
|
||||
alpha = cv2.resize(alpha, (nbw, nbh), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# random crop(crop_h, crop_w) and flip
|
||||
if self.transform:
|
||||
img, alpha, fg, bg = self.transform(img, alpha, fg, bg, cur_crop_h, cur_crop_w)
|
||||
|
||||
# resize to (size_h, size_w)
|
||||
if self.size_h != img.shape[0] or self.size_w != img.shape[1]:
|
||||
# resize
|
||||
img =cv2.resize(img, (self.size_w, self.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
fg =cv2.resize(fg, (self.size_w, self.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
bg =cv2.resize(bg, (self.size_w, self.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
alpha =cv2.resize(alpha, (self.size_w, self.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
trimap = gen_trimap(alpha)
|
||||
grad = compute_gradient(img)
|
||||
|
||||
if self.normalize:
|
||||
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
# first, 0-255 to 0-1
|
||||
# second, x-mean/std and HWC to CHW
|
||||
img_norm = self.normalize(img_rgb)
|
||||
else:
|
||||
img_norm = None
|
||||
|
||||
#img_id = img_info[0].split('/')[-1]
|
||||
#cv2.imwrite("result/debug/{}_img.png".format(img_id), img)
|
||||
#cv2.imwrite("result/debug/{}_alpha.png".format(img_id), alpha)
|
||||
#cv2.imwrite("result/debug/{}_fg.png".format(img_id), fg)
|
||||
#cv2.imwrite("result/debug/{}_bg.png".format(img_id), bg)
|
||||
#cv2.imwrite("result/debug/{}_trimap.png".format(img_id), trimap)
|
||||
#cv2.imwrite("result/debug/{}_grad.png".format(img_id), grad)
|
||||
alpha = torch.from_numpy(alpha.astype(np.float32)[np.newaxis, :, :])
|
||||
trimap = torch.from_numpy(trimap.astype(np.float32)[np.newaxis, :, :])
|
||||
grad = torch.from_numpy(grad.astype(np.float32)[np.newaxis, :, :])
|
||||
img = torch.from_numpy(img.astype(np.float32)).permute(2, 0, 1)
|
||||
fg = torch.from_numpy(fg.astype(np.float32)).permute(2, 0, 1)
|
||||
bg = torch.from_numpy(bg.astype(np.float32)).permute(2, 0, 1)
|
||||
|
||||
return img, alpha, fg, bg, trimap, grad, img_norm, img_info
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
@ -1,36 +0,0 @@
|
||||
import torch
|
||||
from argparse import Namespace
|
||||
import net
|
||||
import cv2
|
||||
import os
|
||||
import numpy as np
|
||||
from deploy import inference_img_whole
|
||||
|
||||
# input file list
|
||||
image_path = "boy-1518482_1920_12_img.png"
|
||||
trimap_path = "boy-1518482_1920_12.png"
|
||||
image = cv2.imread(image_path)
|
||||
trimap = cv2.imread(trimap_path)
|
||||
# print(trimap.shape)
|
||||
trimap = trimap[:, :, 0]
|
||||
# init model
|
||||
args = Namespace(crop_or_resize='whole', cuda=True, max_size=1600, resume='model/stage1_sad_57.1.pth', stage=1)
|
||||
model = net.VGG16(args)
|
||||
ckpt = torch.load(args.resume)
|
||||
model.load_state_dict(ckpt['state_dict'], strict=True)
|
||||
model = model.cuda()
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
with torch.no_grad():
|
||||
pred_mattes = inference_img_whole(args, model, image, trimap)
|
||||
pred_mattes = (pred_mattes * 255).astype(np.uint8)
|
||||
pred_mattes[trimap == 255] = 255
|
||||
pred_mattes[trimap == 0] = 0
|
||||
# print(pred_mattes)
|
||||
# cv2.imwrite('out.png', pred_mattes)
|
||||
|
||||
# import matplotlib.pyplot as plt
|
||||
# plt.imshow(image)
|
||||
# plt.show()
|
||||
|
||||
|
||||
@ -1,281 +0,0 @@
|
||||
import torch
|
||||
import argparse
|
||||
import torch.nn as nn
|
||||
import net
|
||||
import cv2
|
||||
import os
|
||||
from torchvision import transforms
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
def get_args():
|
||||
# Training settings
|
||||
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
|
||||
parser.add_argument('--size_h', type=int, default=320, help="height size of input image")
|
||||
parser.add_argument('--size_w', type=int, default=320, help="width size of input image")
|
||||
parser.add_argument('--imgDir', type=str, required=True, help="directory of image")
|
||||
parser.add_argument('--trimapDir', type=str, required=True, help="directory of trimap")
|
||||
parser.add_argument('--cuda', action='store_true', help='use cuda?')
|
||||
parser.add_argument('--resume', type=str, required=True, help="checkpoint that model resume from")
|
||||
parser.add_argument('--saveDir', type=str, required=True, help="where prediction result save to")
|
||||
parser.add_argument('--alphaDir', type=str, default='', help="directory of gt")
|
||||
parser.add_argument('--stage', type=int, required=True, choices=[0,1,2,3], help="backbone stage")
|
||||
parser.add_argument('--not_strict', action='store_true', help='not copy ckpt strict?')
|
||||
parser.add_argument('--crop_or_resize', type=str, default="whole", choices=["resize", "crop", "whole"], help="how manipulate image before test")
|
||||
parser.add_argument('--max_size', type=int, default=1600, help="max size of test image")
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
return args
|
||||
|
||||
def gen_dataset(imgdir, trimapdir):
|
||||
sample_set = []
|
||||
img_ids = os.listdir(imgdir)
|
||||
img_ids.sort()
|
||||
cnt = len(img_ids)
|
||||
cur = 1
|
||||
for img_id in img_ids:
|
||||
img_name = os.path.join(imgdir, img_id)
|
||||
trimap_name = os.path.join(trimapdir, img_id)
|
||||
|
||||
assert(os.path.exists(img_name))
|
||||
assert(os.path.exists(trimap_name))
|
||||
|
||||
sample_set.append((img_name, trimap_name))
|
||||
|
||||
return sample_set
|
||||
|
||||
def compute_gradient(img):
|
||||
x = cv2.Sobel(img, cv2.CV_16S, 1, 0)
|
||||
y = cv2.Sobel(img, cv2.CV_16S, 0, 1)
|
||||
absX = cv2.convertScaleAbs(x)
|
||||
absY = cv2.convertScaleAbs(y)
|
||||
grad = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
|
||||
grad=cv2.cvtColor(grad, cv2.COLOR_BGR2GRAY)
|
||||
return grad
|
||||
|
||||
|
||||
# inference once for image, return numpy
|
||||
def inference_once(args, model, scale_img, scale_trimap, aligned=True):
|
||||
|
||||
if aligned:
|
||||
assert(scale_img.shape[0] == args.size_h)
|
||||
assert(scale_img.shape[1] == args.size_w)
|
||||
|
||||
normalize = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean = [0.485, 0.456, 0.406],std = [0.229, 0.224, 0.225])
|
||||
])
|
||||
|
||||
scale_img_rgb = cv2.cvtColor(scale_img, cv2.COLOR_BGR2RGB)
|
||||
# first, 0-255 to 0-1
|
||||
# second, x-mean/std and HWC to CHW
|
||||
tensor_img = normalize(scale_img_rgb).unsqueeze(0)
|
||||
|
||||
scale_grad = compute_gradient(scale_img)
|
||||
#tensor_img = torch.from_numpy(scale_img.astype(np.float32)[np.newaxis, :, :, :]).permute(0, 3, 1, 2)
|
||||
tensor_trimap = torch.from_numpy(scale_trimap.astype(np.float32)[np.newaxis, np.newaxis, :, :])
|
||||
tensor_grad = torch.from_numpy(scale_grad.astype(np.float32)[np.newaxis, np.newaxis, :, :])
|
||||
|
||||
if args.cuda:
|
||||
tensor_img = tensor_img.cuda()
|
||||
tensor_trimap = tensor_trimap.cuda()
|
||||
tensor_grad = tensor_grad.cuda()
|
||||
#print('Img Shape:{} Trimap Shape:{}'.format(img.shape, trimap.shape))
|
||||
|
||||
input_t = torch.cat((tensor_img, tensor_trimap / 255.), 1)
|
||||
|
||||
# forward
|
||||
if args.stage <= 1:
|
||||
# stage 1
|
||||
pred_mattes, _ = model(input_t)
|
||||
else:
|
||||
# stage 2, 3
|
||||
_, pred_mattes = model(input_t)
|
||||
pred_mattes = pred_mattes.data
|
||||
if args.cuda:
|
||||
pred_mattes = pred_mattes.cpu()
|
||||
pred_mattes = pred_mattes.numpy()[0, 0, :, :]
|
||||
return pred_mattes
|
||||
|
||||
|
||||
|
||||
# forward for a full image by crop method
|
||||
def inference_img_by_crop(args, model, img, trimap):
|
||||
# crop the pictures, and forward one by one
|
||||
h, w, c = img.shape
|
||||
origin_pred_mattes = np.zeros((h, w), dtype=np.float32)
|
||||
marks = np.zeros((h, w), dtype=np.float32)
|
||||
|
||||
for start_h in range(0, h, args.size_h):
|
||||
end_h = start_h + args.size_h
|
||||
for start_w in range(0, w, args.size_w):
|
||||
|
||||
end_w = start_w + args.size_w
|
||||
crop_img = img[start_h: end_h, start_w: end_w, :]
|
||||
crop_trimap = trimap[start_h: end_h, start_w: end_w]
|
||||
|
||||
crop_origin_h = crop_img.shape[0]
|
||||
crop_origin_w = crop_img.shape[1]
|
||||
|
||||
#print("startH:{} startW:{} H:{} W:{}".format(start_h, start_w, crop_origin_h, crop_origin_w))
|
||||
|
||||
if len(np.where(crop_trimap == 128)[0]) <= 0:
|
||||
continue
|
||||
|
||||
# egde patch in the right or bottom
|
||||
if crop_origin_h != args.size_h or crop_origin_w != args.size_w:
|
||||
crop_img = cv2.resize(crop_img, (args.size_w, args.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
crop_trimap = cv2.resize(crop_trimap, (args.size_w, args.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# inference for each crop image patch
|
||||
pred_mattes = inference_once(args, model, crop_img, crop_trimap)
|
||||
|
||||
if crop_origin_h != args.size_h or crop_origin_w != args.size_w:
|
||||
pred_mattes = cv2.resize(pred_mattes, (crop_origin_w, crop_origin_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
origin_pred_mattes[start_h: end_h, start_w: end_w] += pred_mattes
|
||||
marks[start_h: end_h, start_w: end_w] += 1
|
||||
|
||||
# smooth for overlap part
|
||||
marks[marks <= 0] = 1.
|
||||
origin_pred_mattes /= marks
|
||||
return origin_pred_mattes
|
||||
|
||||
|
||||
# forward for a full image by resize method
|
||||
def inference_img_by_resize(args, model, img, trimap):
|
||||
h, w, c = img.shape
|
||||
# resize for network input, to Tensor
|
||||
scale_img = cv2.resize(img, (args.size_w, args.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
scale_trimap = cv2.resize(trimap, (args.size_w, args.size_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
pred_mattes = inference_once(args, model, scale_img, scale_trimap)
|
||||
|
||||
# resize to origin size
|
||||
origin_pred_mattes = cv2.resize(pred_mattes, (w, h), interpolation = cv2.INTER_LINEAR)
|
||||
assert(origin_pred_mattes.shape == trimap.shape)
|
||||
return origin_pred_mattes
|
||||
|
||||
|
||||
# forward a whole image
|
||||
def inference_img_whole(args, model, img, trimap):
|
||||
h, w, c = img.shape
|
||||
new_h = min(args.max_size, h - (h % 32))
|
||||
new_w = min(args.max_size, w - (w % 32))
|
||||
|
||||
# resize for network input, to Tensor
|
||||
scale_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
scale_trimap = cv2.resize(trimap, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
pred_mattes = inference_once(args, model, scale_img, scale_trimap, aligned=False)
|
||||
|
||||
# resize to origin size
|
||||
origin_pred_mattes = cv2.resize(pred_mattes, (w, h), interpolation = cv2.INTER_LINEAR)
|
||||
assert(origin_pred_mattes.shape == trimap.shape)
|
||||
return origin_pred_mattes
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
print("===> Loading args")
|
||||
args = get_args()
|
||||
|
||||
print("===> Environment init")
|
||||
#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
if args.cuda and not torch.cuda.is_available():
|
||||
raise Exception("No GPU found, please run without --cuda")
|
||||
|
||||
model = net.VGG16(args)
|
||||
ckpt = torch.load(args.resume)
|
||||
if args.not_strict:
|
||||
model.load_state_dict(ckpt['state_dict'], strict=False)
|
||||
else:
|
||||
model.load_state_dict(ckpt['state_dict'], strict=True)
|
||||
|
||||
if args.cuda:
|
||||
model = model.cuda()
|
||||
|
||||
print("===> Load dataset")
|
||||
dataset = gen_dataset(args.imgDir, args.trimapDir)
|
||||
|
||||
mse_diffs = 0.
|
||||
sad_diffs = 0.
|
||||
cnt = len(dataset)
|
||||
cur = 0
|
||||
t0 = time.time()
|
||||
for img_path, trimap_path in dataset:
|
||||
img = cv2.imread(img_path)
|
||||
trimap = cv2.imread(trimap_path)[:, :, 0]
|
||||
|
||||
assert(img.shape[:2] == trimap.shape[:2])
|
||||
|
||||
img_info = (img_path.split('/')[-1], img.shape[0], img.shape[1])
|
||||
|
||||
cur += 1
|
||||
print('[{}/{}] {}'.format(cur, cnt, img_info[0]))
|
||||
|
||||
with torch.no_grad():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if args.crop_or_resize == "whole":
|
||||
origin_pred_mattes = inference_img_whole(args, model, img, trimap)
|
||||
elif args.crop_or_resize == "crop":
|
||||
origin_pred_mattes = inference_img_by_crop(args, model, img, trimap)
|
||||
else:
|
||||
origin_pred_mattes = inference_img_by_resize(args, model, img, trimap)
|
||||
|
||||
# only attention unknown region
|
||||
origin_pred_mattes[trimap == 255] = 1.
|
||||
origin_pred_mattes[trimap == 0 ] = 0.
|
||||
|
||||
# origin trimap
|
||||
pixel = float((trimap == 128).sum())
|
||||
|
||||
# eval if gt alpha is given
|
||||
if args.alphaDir != '':
|
||||
alpha_name = os.path.join(args.alphaDir, img_info[0])
|
||||
assert(os.path.exists(alpha_name))
|
||||
alpha = cv2.imread(alpha_name)[:, :, 0] / 255.
|
||||
assert(alpha.shape == origin_pred_mattes.shape)
|
||||
|
||||
#x1 = (alpha[trimap == 255] == 1.0).sum() # x3
|
||||
#x2 = (alpha[trimap == 0] == 0.0).sum() # x5
|
||||
#x3 = (trimap == 255).sum()
|
||||
#x4 = (trimap == 128).sum()
|
||||
#x5 = (trimap == 0).sum()
|
||||
#x6 = trimap.size # sum(x3,x4,x5)
|
||||
#x7 = (alpha[trimap == 255] < 1.0).sum() # 0
|
||||
#x8 = (alpha[trimap == 0] > 0).sum() #
|
||||
|
||||
#print(x1, x2, x3, x4, x5, x6, x7, x8)
|
||||
#assert(x1 == x3)
|
||||
#assert(x2 == x5)
|
||||
#assert(x6 == x3 + x4 + x5)
|
||||
#assert(x7 == 0)
|
||||
#assert(x8 == 0)
|
||||
|
||||
mse_diff = ((origin_pred_mattes - alpha) ** 2).sum() / pixel
|
||||
sad_diff = np.abs(origin_pred_mattes - alpha).sum()
|
||||
mse_diffs += mse_diff
|
||||
sad_diffs += sad_diff
|
||||
print("sad:{} mse:{}".format(sad_diff, mse_diff))
|
||||
|
||||
origin_pred_mattes = (origin_pred_mattes * 255).astype(np.uint8)
|
||||
res = origin_pred_mattes.copy()
|
||||
|
||||
# only attention unknown region
|
||||
res[trimap == 255] = 255
|
||||
res[trimap == 0 ] = 0
|
||||
|
||||
if not os.path.exists(args.saveDir):
|
||||
os.makedirs(args.saveDir)
|
||||
cv2.imwrite(os.path.join(args.saveDir, img_info[0]), res)
|
||||
|
||||
print("Avg-Cost: {} s/image".format((time.time() - t0) / cnt))
|
||||
if args.alphaDir != '':
|
||||
print("Eval-MSE: {}".format(mse_diffs / cur))
|
||||
print("Eval-SAD: {}".format(sad_diffs / cur))
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,126 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
import cv2
|
||||
import torch.nn.functional as F
|
||||
|
||||
class VGG16(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(VGG16, self).__init__()
|
||||
self.stage = args.stage
|
||||
|
||||
self.conv1_1 = nn.Conv2d(4, 64, kernel_size=3,stride = 1, padding=1,bias=True)
|
||||
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3,stride = 1, padding=1,bias=True)
|
||||
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1,bias=True)
|
||||
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1,bias=True)
|
||||
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1,bias=True)
|
||||
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1,bias=True)
|
||||
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1,bias=True)
|
||||
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1,bias=True)
|
||||
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
|
||||
# model released before 2019.09.09 should use kernel_size=1 & padding=0
|
||||
#self.conv6_1 = nn.Conv2d(512, 512, kernel_size=1, padding=0,bias=True)
|
||||
self.conv6_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1,bias=True)
|
||||
|
||||
self.deconv6_1 = nn.Conv2d(512, 512, kernel_size=1,bias=True)
|
||||
self.deconv5_1 = nn.Conv2d(512, 512, kernel_size=5, padding=2,bias=True)
|
||||
self.deconv4_1 = nn.Conv2d(512, 256, kernel_size=5, padding=2,bias=True)
|
||||
self.deconv3_1 = nn.Conv2d(256, 128, kernel_size=5, padding=2,bias=True)
|
||||
self.deconv2_1 = nn.Conv2d(128, 64, kernel_size=5, padding=2,bias=True)
|
||||
self.deconv1_1 = nn.Conv2d(64, 64, kernel_size=5, padding=2,bias=True)
|
||||
|
||||
self.deconv1 = nn.Conv2d(64, 1, kernel_size=5, padding=2,bias=True)
|
||||
|
||||
if args.stage == 2:
|
||||
# for stage2 training
|
||||
for p in self.parameters():
|
||||
p.requires_grad=False
|
||||
|
||||
if self.stage == 2 or self.stage == 3:
|
||||
self.refine_conv1 = nn.Conv2d(4, 64, kernel_size=3, padding=1, bias=True)
|
||||
self.refine_conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=True)
|
||||
self.refine_conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=True)
|
||||
self.refine_pred = nn.Conv2d(64, 1, kernel_size=3, padding=1, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
# Stage 1
|
||||
x11 = F.relu(self.conv1_1(x))
|
||||
x12 = F.relu(self.conv1_2(x11))
|
||||
x1p, id1 = F.max_pool2d(x12,kernel_size=(2,2), stride=(2,2),return_indices=True)
|
||||
|
||||
# Stage 2
|
||||
x21 = F.relu(self.conv2_1(x1p))
|
||||
x22 = F.relu(self.conv2_2(x21))
|
||||
x2p, id2 = F.max_pool2d(x22,kernel_size=(2,2), stride=(2,2),return_indices=True)
|
||||
|
||||
# Stage 3
|
||||
x31 = F.relu(self.conv3_1(x2p))
|
||||
x32 = F.relu(self.conv3_2(x31))
|
||||
x33 = F.relu(self.conv3_3(x32))
|
||||
x3p, id3 = F.max_pool2d(x33,kernel_size=(2,2), stride=(2,2),return_indices=True)
|
||||
|
||||
# Stage 4
|
||||
x41 = F.relu(self.conv4_1(x3p))
|
||||
x42 = F.relu(self.conv4_2(x41))
|
||||
x43 = F.relu(self.conv4_3(x42))
|
||||
x4p, id4 = F.max_pool2d(x43,kernel_size=(2,2), stride=(2,2),return_indices=True)
|
||||
|
||||
# Stage 5
|
||||
x51 = F.relu(self.conv5_1(x4p))
|
||||
x52 = F.relu(self.conv5_2(x51))
|
||||
x53 = F.relu(self.conv5_3(x52))
|
||||
x5p, id5 = F.max_pool2d(x53,kernel_size=(2,2), stride=(2,2),return_indices=True)
|
||||
|
||||
# Stage 6
|
||||
x61 = F.relu(self.conv6_1(x5p))
|
||||
|
||||
# Stage 6d
|
||||
x61d = F.relu(self.deconv6_1(x61))
|
||||
|
||||
# Stage 5d
|
||||
x5d = F.max_unpool2d(x61d,id5, kernel_size=2, stride=2)
|
||||
x51d = F.relu(self.deconv5_1(x5d))
|
||||
|
||||
# Stage 4d
|
||||
x4d = F.max_unpool2d(x51d, id4, kernel_size=2, stride=2)
|
||||
x41d = F.relu(self.deconv4_1(x4d))
|
||||
|
||||
# Stage 3d
|
||||
x3d = F.max_unpool2d(x41d, id3, kernel_size=2, stride=2)
|
||||
x31d = F.relu(self.deconv3_1(x3d))
|
||||
|
||||
# Stage 2d
|
||||
x2d = F.max_unpool2d(x31d, id2, kernel_size=2, stride=2)
|
||||
x21d = F.relu(self.deconv2_1(x2d))
|
||||
|
||||
# Stage 1d
|
||||
x1d = F.max_unpool2d(x21d, id1, kernel_size=2, stride=2)
|
||||
x12d = F.relu(self.deconv1_1(x1d))
|
||||
|
||||
# Should add sigmoid? github repo add so.
|
||||
raw_alpha = self.deconv1(x12d)
|
||||
pred_mattes = F.sigmoid(raw_alpha)
|
||||
|
||||
if self.stage <= 1:
|
||||
return pred_mattes, 0
|
||||
|
||||
# Stage2 refine conv1
|
||||
refine0 = torch.cat((x[:, :3, :, :], pred_mattes), 1)
|
||||
refine1 = F.relu(self.refine_conv1(refine0))
|
||||
refine2 = F.relu(self.refine_conv2(refine1))
|
||||
refine3 = F.relu(self.refine_conv3(refine2))
|
||||
# Should add sigmoid?
|
||||
# sigmoid lead to refine result all converge to 0...
|
||||
#pred_refine = F.sigmoid(self.refine_pred(refine3))
|
||||
pred_refine = self.refine_pred(refine3)
|
||||
|
||||
pred_alpha = F.sigmoid(raw_alpha + pred_refine)
|
||||
|
||||
#print(pred_mattes.mean(), pred_alpha.mean(), pred_refine.sum())
|
||||
|
||||
return pred_mattes, pred_alpha
|
||||
@ -1,63 +0,0 @@
|
||||
import torch
|
||||
import torchvision
|
||||
import collections
|
||||
import os
|
||||
|
||||
HOME = os.environ['HOME']
|
||||
model_path = "{}/.torch/models/vgg16-397923af.pth".format(HOME)
|
||||
#model_path = "/data/liuliang/deep_image_matting/train/vgg16-397923af.pth"
|
||||
if not os.path.exists(model_path):
|
||||
model = torchvision.models.vgg16(pretrained=True)
|
||||
assert(os.path.exists(model_path))
|
||||
|
||||
x = torch.load(model_path)
|
||||
|
||||
val = collections.OrderedDict()
|
||||
val['conv1_1.weight'] = torch.cat((x['features.0.weight'], torch.zeros(64, 1, 3, 3)), 1)
|
||||
|
||||
replace = { u'features.0.bias' : 'conv1_1.bias',
|
||||
u'features.2.weight' : 'conv1_2.weight',
|
||||
u'features.2.bias' : 'conv1_2.bias',
|
||||
u'features.5.weight' : 'conv2_1.weight',
|
||||
u'features.5.bias' : 'conv2_1.bias',
|
||||
u'features.7.weight' : 'conv2_2.weight',
|
||||
u'features.7.bias' : 'conv2_2.bias',
|
||||
u'features.10.weight': 'conv3_1.weight',
|
||||
u'features.10.bias' : 'conv3_1.bias',
|
||||
u'features.12.weight': 'conv3_2.weight',
|
||||
u'features.12.bias' : 'conv3_2.bias',
|
||||
u'features.14.weight': 'conv3_3.weight',
|
||||
u'features.14.bias' : 'conv3_3.bias',
|
||||
u'features.17.weight': 'conv4_1.weight',
|
||||
u'features.17.bias' : 'conv4_1.bias',
|
||||
u'features.19.weight': 'conv4_2.weight',
|
||||
u'features.19.bias' : 'conv4_2.bias',
|
||||
u'features.21.weight': 'conv4_3.weight',
|
||||
u'features.21.bias' : 'conv4_3.bias',
|
||||
u'features.24.weight': 'conv5_1.weight',
|
||||
u'features.24.bias' : 'conv5_1.bias',
|
||||
u'features.26.weight': 'conv5_2.weight',
|
||||
u'features.26.bias' : 'conv5_2.bias',
|
||||
u'features.28.weight': 'conv5_3.weight',
|
||||
u'features.28.bias' : 'conv5_3.bias'
|
||||
}
|
||||
|
||||
#print(x['classifier.0.weight'].shape)
|
||||
#print(x['classifier.0.bias'].shape)
|
||||
|
||||
#tmp1 = x['classifier.0.weight'].reshape(4096, 512, 7, 7)
|
||||
#print(tmp1.shape)
|
||||
|
||||
#val['conv6_1.weight'] = tmp1[:512, :, :, :]
|
||||
#val['conv6_1.bias'] = x['classifier.0.bias']
|
||||
|
||||
for key in replace.keys():
|
||||
print(key, replace[key])
|
||||
val[replace[key]] = x[key]
|
||||
|
||||
y = {}
|
||||
y['state_dict'] = val
|
||||
y['epoch'] = 0
|
||||
if not os.path.exists('./model'):
|
||||
os.makedirs('./model')
|
||||
torch.save(y, './model/vgg_state_dict.pth')
|
||||
@ -1,137 +0,0 @@
|
||||
# composite image with dataset from "deep image matting"
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import math
|
||||
import time
|
||||
import shutil
|
||||
|
||||
root_dir = "/home/liuliang/Downloads/Combined_Dataset"
|
||||
test_bg_dir = '/home/liuliang/Desktop/dataset/matting/VOCdevkit/VOC2012/JPEGImages'
|
||||
train_bg_dir = '/home/liuliang/Desktop/dataset/matting/mscoco/train2017'
|
||||
|
||||
def my_composite(fg_names, bg_names, fg_dir, alpha_dir, bg_dir, num_bg, comp_dir):
|
||||
fg_ids = open(fg_names).readlines()
|
||||
bg_ids = open(bg_names).readlines()
|
||||
|
||||
fg_cnt = len(fg_ids)
|
||||
bg_cnt = len(bg_ids)
|
||||
print(fg_cnt, bg_cnt)
|
||||
assert(fg_cnt * num_bg == bg_cnt)
|
||||
|
||||
for i in range(fg_cnt):
|
||||
im_name = fg_ids[i].strip("\n").strip("\r")
|
||||
fg_path = os.path.join(fg_dir, im_name)
|
||||
alpha_path = os.path.join(alpha_dir, im_name)
|
||||
|
||||
#print(fg_path, alpha_path)
|
||||
assert(os.path.exists(fg_path))
|
||||
assert(os.path.exists(alpha_path))
|
||||
|
||||
fg = cv2.imread(fg_path)
|
||||
alpha = cv2.imread(alpha_path)
|
||||
|
||||
#print("alpha shape:", alpha.shape, "image shape:", fg.shape)
|
||||
assert(alpha.shape == fg.shape)
|
||||
|
||||
h, w ,c = fg.shape
|
||||
base = i * num_bg
|
||||
for bcount in range(num_bg):
|
||||
bg_path = os.path.join(bg_dir, bg_ids[base + bcount].strip("\n").strip("\r"))
|
||||
print(base + bcount, fg_path, bg_path)
|
||||
assert(os.path.exists(bg_path))
|
||||
bg = cv2.imread(bg_path)
|
||||
bh, bw, bc = bg.shape
|
||||
|
||||
wratio = float(w) / bw
|
||||
hratio = float(h) / bh
|
||||
ratio = wratio if wratio > hratio else hratio
|
||||
if ratio > 1:
|
||||
new_bw = int(bw * ratio + 1.0)
|
||||
new_bh = int(bh * ratio + 1.0)
|
||||
bg = cv2.resize(bg, (new_bw, new_bh), interpolation=cv2.INTER_LINEAR)
|
||||
bg = bg[0 : h, 0 : w, :]
|
||||
#print(bg.shape)
|
||||
assert(bg.shape == fg.shape)
|
||||
alpha_f = alpha / 255.
|
||||
comp = fg * alpha_f + bg * (1. - alpha_f)
|
||||
|
||||
img_save_id = im_name[:len(im_name)-4] + '_' + str(bcount) + '.png'
|
||||
comp_save_path = os.path.join(comp_dir, "image/" + img_save_id)
|
||||
fg_save_path = os.path.join(comp_dir, "fg/" + img_save_id)
|
||||
bg_save_path = os.path.join(comp_dir, "bg/" + img_save_id)
|
||||
alpha_save_path = os.path.join(comp_dir, "alpha/" + img_save_id)
|
||||
|
||||
cv2.imwrite(comp_save_path, comp)
|
||||
cv2.imwrite(fg_save_path, fg)
|
||||
cv2.imwrite(bg_save_path, bg)
|
||||
cv2.imwrite(alpha_save_path, alpha)
|
||||
|
||||
|
||||
def copy_dir2dir(src_dir, des_dir):
|
||||
for img_id in os.listdir(src_dir):
|
||||
shutil.copyfile(os.path.join(src_dir, img_id), os.path.join(des_dir, img_id))
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
test_num_bg = 20
|
||||
test_fg_names = os.path.join(root_dir, "Test_set/test_fg_names.txt")
|
||||
test_bg_names = os.path.join(root_dir, "Test_set/test_bg_names.txt")
|
||||
test_fg_dir = os.path.join(root_dir, "Test_set/Adobe-licensed images/fg")
|
||||
test_alpha_dir = os.path.join(root_dir, "Test_set/Adobe-licensed images/alpha")
|
||||
test_trimap_dir = os.path.join(root_dir, "Test_set/Adobe-licensed images/trimaps")
|
||||
test_comp_dir = os.path.join(root_dir, "Test_set/comp")
|
||||
|
||||
train_num_bg = 100
|
||||
train_fg_names = os.path.join(root_dir, "Training_set/training_fg_names.txt")
|
||||
train_bg_names_coco2014 = os.path.join(root_dir, "Training_set/training_bg_names.txt")
|
||||
train_bg_names_coco2017 = os.path.join(root_dir, "Training_set/training_bg_names_coco2017.txt")
|
||||
train_fg_dir = os.path.join(root_dir, "Training_set/all/fg")
|
||||
train_alpha_dir = os.path.join(root_dir, "Training_set/all/alpha")
|
||||
train_comp_dir = os.path.join(root_dir, "Training_set/comp")
|
||||
|
||||
# change the bg names formate if is coco 2017
|
||||
fin = open(train_bg_names_coco2014, 'r')
|
||||
fout = open(train_bg_names_coco2017, 'w')
|
||||
lls = fin.readlines()
|
||||
for l in lls:
|
||||
fout.write(l[15:])
|
||||
fin.close()
|
||||
fout.close()
|
||||
|
||||
if not os.path.exists(test_comp_dir):
|
||||
os.makedirs(test_comp_dir + '/image')
|
||||
os.makedirs(test_comp_dir + '/fg')
|
||||
os.makedirs(test_comp_dir + '/bg')
|
||||
os.makedirs(test_comp_dir + '/alpha')
|
||||
os.makedirs(test_comp_dir + '/trimap')
|
||||
|
||||
if not os.path.exists(train_comp_dir):
|
||||
os.makedirs(train_comp_dir + '/image')
|
||||
os.makedirs(train_comp_dir + '/fg')
|
||||
os.makedirs(train_comp_dir + '/bg')
|
||||
os.makedirs(train_comp_dir + '/alpha')
|
||||
|
||||
if not os.path.exists(train_alpha_dir):
|
||||
os.makedirs(train_alpha_dir)
|
||||
|
||||
if not os.path.exists(train_fg_dir):
|
||||
os.makedirs(train_fg_dir)
|
||||
|
||||
# copy test trimaps
|
||||
copy_dir2dir(test_trimap_dir, test_comp_dir + '/trimap')
|
||||
# copy train images together
|
||||
copy_dir2dir(os.path.join(root_dir, "Training_set/Adobe-licensed images/alpha"), train_alpha_dir)
|
||||
copy_dir2dir(os.path.join(root_dir, "Training_set/Adobe-licensed images/fg"), train_fg_dir)
|
||||
copy_dir2dir(os.path.join(root_dir, "Training_set/Other/alpha"), train_alpha_dir)
|
||||
copy_dir2dir(os.path.join(root_dir, "Training_set/Other/fg"), train_fg_dir)
|
||||
|
||||
# composite test image
|
||||
my_composite(test_fg_names, test_bg_names, test_fg_dir, test_alpha_dir, test_bg_dir, test_num_bg, test_comp_dir)
|
||||
# composite train image
|
||||
my_composite(train_fg_names, train_bg_names_coco2017, train_fg_dir, train_alpha_dir, train_bg_dir, train_num_bg, train_comp_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@ -1,50 +0,0 @@
|
||||
import re
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pylab import *
|
||||
|
||||
# args: log_name, match_rule, self_log_interval, smooth_log_interation
|
||||
loss_file_name = "simple_loss"
|
||||
title = "{}_Loss".format(loss_file_name)
|
||||
f = open("../log/{}.log".format(loss_file_name))
|
||||
pattern = re.compile(r'Loss:[ ]*\d+\.\d+')
|
||||
self_inter = 10
|
||||
smooth = 20
|
||||
|
||||
# read log file
|
||||
lines = f.readlines()
|
||||
print("Line: {}".format(len(lines)))
|
||||
ys = []
|
||||
k = 0
|
||||
cnt = 0
|
||||
sum_y = 0.
|
||||
|
||||
# read one by one
|
||||
for line in lines:
|
||||
obj = re.search(pattern, line)
|
||||
if obj:
|
||||
val = float(obj.group().split(':')[-1])
|
||||
sum_y += val
|
||||
k += 1
|
||||
if k >= smooth:
|
||||
ys.append(sum_y / k)
|
||||
sum_y = 0.
|
||||
k = 0
|
||||
cnt += 1
|
||||
if cnt % 10 == 0:
|
||||
print("ys cnt: {}".format(cnt))
|
||||
if k > 0:
|
||||
ys.append(sum_y / k)
|
||||
|
||||
ys = np.array(ys)
|
||||
xs = np.arange(len(ys)) * self_inter * smooth
|
||||
|
||||
print(xs)
|
||||
print(ys)
|
||||
|
||||
plt.plot(xs, ys)
|
||||
plt.title(title)
|
||||
plt.xlabel("Iter")
|
||||
plt.ylabel("Loss")
|
||||
plt.savefig("../log/{}.png".format(title))
|
||||
plt.show()
|
||||
@ -1,96 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools'])
|
||||
|
||||
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torchvision import transforms, datasets
|
||||
import numpy as np
|
||||
|
||||
def getSeg(input_image,f):
|
||||
model = torch.load(baseLoc+'weights/deeplabv3/deeplabv3+model.pt')
|
||||
model.eval()
|
||||
preprocess = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
input_image = Image.fromarray(input_image)
|
||||
input_tensor = preprocess(input_image)
|
||||
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
|
||||
if torch.cuda.is_available() and not f:
|
||||
input_batch = input_batch.to('cuda')
|
||||
model.to('cuda')
|
||||
|
||||
with torch.no_grad():
|
||||
output = model(input_batch)['out'][0]
|
||||
output_predictions = output.argmax(0)
|
||||
|
||||
|
||||
# create a color pallette, selecting a color for each class
|
||||
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
|
||||
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
|
||||
colors = (colors % 255).numpy().astype("uint8")
|
||||
|
||||
# plot the semantic segmentation predictions of 21 classes in each color
|
||||
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
|
||||
|
||||
tmp = np.array(r)
|
||||
tmp2 = 10*np.repeat(tmp[:, :, np.newaxis], 3, axis=2)
|
||||
|
||||
|
||||
return tmp2
|
||||
|
||||
|
||||
def channelData(layer):#convert gimp image to numpy
|
||||
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
||||
pixChars=region[:,:] # Take whole layer
|
||||
bpp=region.bpp
|
||||
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
||||
|
||||
def createResultLayer(image,name,result):
|
||||
rlBytes=np.uint8(result).tobytes();
|
||||
rl=gimp.Layer(image,name,image.width,image.height,0,100,NORMAL_MODE)
|
||||
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
||||
region[:,:]=rlBytes
|
||||
image.add_layer(rl,0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def deeplabv3(img, layer,cFlag) :
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Generating semantic segmentation map for " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Generating semantic segmentation map for " + layer.name + "...")
|
||||
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:,:,0:3]
|
||||
cpy=getSeg(imgmat,cFlag)
|
||||
createResultLayer(img,'new_output',cpy)
|
||||
|
||||
|
||||
register(
|
||||
"semantic-segmentation",
|
||||
"semantic-segmentation",
|
||||
"Generate semantic segmentation map based on deep learning.",
|
||||
"Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2020",
|
||||
"semantic-segmentation...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[ (PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False)
|
||||
],
|
||||
[],
|
||||
deeplabv3, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,157 +0,0 @@
|
||||
import os
|
||||
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
||||
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'pytorch-SRResNet'])
|
||||
|
||||
from argparse import Namespace
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import cv2
|
||||
|
||||
|
||||
def getlabelmat(mask, idx):
|
||||
x = np.zeros((mask.shape[0], mask.shape[1], 3))
|
||||
x[mask == idx, 0] = colors[idx][0]
|
||||
x[mask == idx, 1] = colors[idx][1]
|
||||
x[mask == idx, 2] = colors[idx][2]
|
||||
return x
|
||||
|
||||
|
||||
def colorMask(mask):
|
||||
x = np.zeros((mask.shape[0], mask.shape[1], 3))
|
||||
for idx in range(19):
|
||||
x = x + getlabelmat(mask, idx)
|
||||
return np.uint8(x)
|
||||
|
||||
|
||||
def getnewimg(input_image, s, cFlag, fFlag):
|
||||
opt = Namespace(cuda=torch.cuda.is_available() and not cFlag,
|
||||
model=baseLoc + 'weights/super_resolution/model_srresnet.pth',
|
||||
dataset='Set5', scale=s, gpus=0)
|
||||
|
||||
w, h = input_image.shape[0:2]
|
||||
cuda = opt.cuda
|
||||
|
||||
if cuda:
|
||||
model = torch.load(opt.model)["model"]
|
||||
else:
|
||||
model = torch.load(opt.model, map_location=torch.device('cpu'))["model"]
|
||||
|
||||
im_input = input_image.astype(np.float32).transpose(2, 0, 1)
|
||||
im_input = im_input.reshape(1, im_input.shape[0], im_input.shape[1], im_input.shape[2])
|
||||
im_input = Variable(torch.from_numpy(im_input / 255.).float())
|
||||
|
||||
if cuda and not cFlag:
|
||||
model = model.cuda()
|
||||
im_input = im_input.cuda()
|
||||
else:
|
||||
model = model.cpu()
|
||||
|
||||
if fFlag:
|
||||
im_h = np.zeros([4 * w, 4 * h, 3])
|
||||
wbin = 300
|
||||
i = 0
|
||||
idx = 0
|
||||
t = float(w * h) / float(wbin * wbin)
|
||||
while i < w:
|
||||
i_end = min(i + wbin, w)
|
||||
j = 0
|
||||
while j < h:
|
||||
j_end = min(j + wbin, h)
|
||||
patch = im_input[:, :, i:i_end, j:j_end]
|
||||
# patch_merge_out_numpy = denoiser(patch, c, pss, model, model_est, opt, cFlag)
|
||||
HR_4x = model(patch)
|
||||
HR_4x = HR_4x.cpu().data[0].numpy().astype(np.float32) * 255.
|
||||
HR_4x = np.clip(HR_4x, 0., 255.).transpose(1, 2, 0).astype(np.uint8)
|
||||
|
||||
im_h[4 * i:4 * i_end, 4 * j:4 * j_end, :] = HR_4x
|
||||
j = j_end
|
||||
idx = idx + 1
|
||||
gimp.progress_update(float(idx) / float(t))
|
||||
gimp.displays_flush()
|
||||
i = i_end
|
||||
else:
|
||||
HR_4x = model(im_input)
|
||||
HR_4x = HR_4x.cpu()
|
||||
im_h = HR_4x.data[0].numpy().astype(np.float32)
|
||||
im_h = im_h * 255.
|
||||
im_h = np.clip(im_h, 0., 255.)
|
||||
im_h = im_h.transpose(1, 2, 0).astype(np.uint8)
|
||||
return im_h
|
||||
|
||||
|
||||
def channelData(layer): # convert gimp image to numpy
|
||||
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
||||
pixChars = region[:, :] # Take whole layer
|
||||
bpp = region.bpp
|
||||
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
||||
|
||||
|
||||
def createResultFile(name, layer_np):
|
||||
h, w, d = layer_np.shape
|
||||
img = pdb.gimp_image_new(w, h, RGB)
|
||||
display = pdb.gimp_display_new(img)
|
||||
|
||||
rlBytes = np.uint8(layer_np).tobytes();
|
||||
rl = gimp.Layer(img, name, img.width, img.height, RGB, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
|
||||
pdb.gimp_image_insert_layer(img, rl, None, 0)
|
||||
|
||||
gimp.displays_flush()
|
||||
|
||||
def createResultLayer(image, name, result):
|
||||
rlBytes = np.uint8(result).tobytes();
|
||||
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
||||
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
||||
region[:, :] = rlBytes
|
||||
image.add_layer(rl, 0)
|
||||
gimp.displays_flush()
|
||||
|
||||
def super_resolution(img, layer, scale, cFlag, fFlag):
|
||||
imgmat = channelData(layer)
|
||||
if imgmat.shape[0] != img.height or imgmat.shape[1] != img.width:
|
||||
pdb.gimp_message(" Do (Layer -> Layer to Image Size) first and try again.")
|
||||
else:
|
||||
if torch.cuda.is_available() and not cFlag:
|
||||
gimp.progress_init("(Using GPU) Running super-resolution for " + layer.name + "...")
|
||||
else:
|
||||
gimp.progress_init("(Using CPU) Running super-resolution for " + layer.name + "...")
|
||||
|
||||
if imgmat.shape[2] == 4: # get rid of alpha channel
|
||||
imgmat = imgmat[:, :, 0:3]
|
||||
cpy = getnewimg(imgmat, scale, cFlag, fFlag)
|
||||
cpy = cv2.resize(cpy, (0, 0), fx=scale / 4, fy=scale / 4)
|
||||
if scale==1:
|
||||
createResultLayer(img, layer.name + '_super', cpy)
|
||||
else:
|
||||
createResultFile(layer.name + '_super', cpy)
|
||||
|
||||
|
||||
register(
|
||||
"super-resolution",
|
||||
"super-resolution",
|
||||
"Running super-resolution.",
|
||||
"Kritik Soman",
|
||||
"Your",
|
||||
"2020",
|
||||
"super-resolution...",
|
||||
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
||||
[(PF_IMAGE, "image", "Input image", None),
|
||||
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
||||
(PF_SLIDER, "Scale", "Scale", 4, (1, 4, 0.5)),
|
||||
(PF_BOOL, "fcpu", "Force CPU", False),
|
||||
(PF_BOOL, "ffilter", "Use as filter", True)
|
||||
],
|
||||
[],
|
||||
super_resolution, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
@ -1,262 +0,0 @@
|
||||
import requests
|
||||
import os
|
||||
from gimpfu import *
|
||||
|
||||
def download_file_from_google_drive(id, destination,fileSize):
|
||||
URL = "https://docs.google.com/uc?export=download"
|
||||
session = requests.Session()
|
||||
response = session.get(URL, params = { 'id' : id }, stream = True)
|
||||
token = get_confirm_token(response)
|
||||
if token:
|
||||
params = { 'id' : id, 'confirm' : token }
|
||||
response = session.get(URL, params = params, stream = True)
|
||||
save_response_content(response, destination,fileSize)
|
||||
|
||||
def get_confirm_token(response):
|
||||
for key, value in response.cookies.items():
|
||||
if key.startswith('download_warning'):
|
||||
return value
|
||||
return None
|
||||
|
||||
def save_response_content(response, destination,fileSize):
|
||||
CHUNK_SIZE = 1 * 1024 * 1024
|
||||
n = 0
|
||||
with open(destination, "wb") as f:
|
||||
for chunk in response.iter_content(CHUNK_SIZE):
|
||||
if chunk: # filter out keep-alive new chunks
|
||||
f.write(chunk)
|
||||
n = n + 1
|
||||
gimp.progress_update(float(n)/float(fileSize))
|
||||
gimp.displays_flush()
|
||||
|
||||
def syncGit(baseLoc):
|
||||
#git
|
||||
file_id = '17hdADAEWzHJCtSQrDyQs8v_MHTq1H04z'
|
||||
fileSize = 0.5 #in MB
|
||||
mFName = 'tmp.zip'
|
||||
destination = baseLoc + '/' + mFName
|
||||
gimp.progress_init("Downloading plugins (~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
import zipfile
|
||||
with zipfile.ZipFile(destination, 'r') as zip_ref:
|
||||
zip_ref.extractall(baseLoc)
|
||||
|
||||
import shutil
|
||||
root_src_dir = baseLoc + 'GIMP-ML-master/gimp-plugins/'
|
||||
root_dst_dir = baseLoc
|
||||
|
||||
for src_dir, dirs, files in os.walk(root_src_dir):
|
||||
dst_dir = src_dir.replace(root_src_dir, root_dst_dir, 1)
|
||||
if not os.path.exists(dst_dir):
|
||||
os.makedirs(dst_dir)
|
||||
for file_ in files:
|
||||
src_file = os.path.join(src_dir, file_)
|
||||
dst_file = os.path.join(dst_dir, file_)
|
||||
if os.path.exists(dst_file):
|
||||
if os.path.samefile(src_file, dst_file):
|
||||
continue
|
||||
os.remove(dst_file)
|
||||
shutil.move(src_file, dst_dir)
|
||||
shutil.rmtree(baseLoc+'GIMP-ML-master')
|
||||
os.remove(baseLoc+'tmp.zip')
|
||||
|
||||
def sync(path,flag):
|
||||
if not os.path.isdir(path):
|
||||
os.mkdir(path)
|
||||
|
||||
#deepmatting
|
||||
model = 'deepmatting'
|
||||
file_id = '11dxJKH8p7xkcGtMtvzMUw-ua6pZ0vrfw'
|
||||
fileSize = 108 #in MB
|
||||
mFName = 'stage1_sad_57.1.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#MiDaS
|
||||
model = 'MiDaS'
|
||||
file_id = '11eap5jc-4SCX_sMMxYE6Bi5q_BKw894a'
|
||||
fileSize = 143 #in MB
|
||||
mFName = 'model.pt'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#colorize
|
||||
model = 'colorize'
|
||||
file_id = '12tKfNIDewgJPbW3FiITV_AMbOtZWP0Eg'
|
||||
fileSize = 130 #in MB
|
||||
mFName = 'caffemodel.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#super_resolution
|
||||
model = 'super_resolution'
|
||||
file_id = '11GwnqKsYo2jujACD_GMB9uMTQfsuk2RY'
|
||||
fileSize = 6 #in MB
|
||||
mFName = 'model_srresnet.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#faceparse
|
||||
model = 'faceparse'
|
||||
file_id = '115nnWD0FoDkplTJYBY7lTQu1VNXFbCA_'
|
||||
fileSize = 51 #in MB
|
||||
mFName = '79999_iter.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#deblur
|
||||
model = 'deblur'
|
||||
file_id = '11Tt4a_URCer4ZxZA2l3dLMRVeSwoBFYP'
|
||||
fileSize = 233 #in MB
|
||||
mFName = 'mymodel.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
file_id = '11MCHMVhs4aaMGSusqiu0rtAo97xuC1GA'
|
||||
fileSize = 234 #in MB
|
||||
mFName = 'best_fpn.h5'
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#deeplabv3
|
||||
model = 'deeplabv3'
|
||||
file_id = '11rX1MHjhmtaoFTQ7ao4p6b31Oz300i0G'
|
||||
fileSize = 233 #in MB
|
||||
mFName = 'deeplabv3+model.pt'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#facegen
|
||||
model = 'facegen'
|
||||
ifolder = 'label2face_512p'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
if not os.path.isdir(path + '/' + model + '/' + ifolder):
|
||||
os.mkdir(path + '/' + model + '/' + ifolder)
|
||||
file_id = '122dREA3R0vsSWbrzBwhF5oqSEJ7yrRbL'
|
||||
fileSize = 342 #in MB
|
||||
mFName = 'latest_net_G.pth'
|
||||
|
||||
destination = path + '/' + model + '/' + ifolder + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#deepdehaze
|
||||
model = 'deepdehaze'
|
||||
file_id = '1hrd310nYCbh6ui_ZsZci7Zna2AFP1sMS'
|
||||
fileSize = 0.008 #in MB
|
||||
mFName = 'dehazer.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#deepdenoise
|
||||
model = 'deepdenoise'
|
||||
file_id = '1acZ1FTNMuAQaYtE3RYLA8fs8cQrW2tZ_'
|
||||
fileSize = 0.166 #in MB
|
||||
mFName = 'est_net.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
file_id = '1tBoyDxYJ92pvopBJeK9PmG_jMA_Ut38_'
|
||||
fileSize = 3 #in MB
|
||||
mFName = 'net.pth'
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#enlighten
|
||||
model = 'enlightening'
|
||||
file_id = '1V8ARc2tDgUUpc11xiT5Y9HFQgC6Ug2T6'
|
||||
fileSize = 35 #in MB
|
||||
mFName = '200_net_G_A.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#interpolateframes
|
||||
model = 'interpolateframes'
|
||||
file_id = '1bHmO9-_ENTYoN1-BNwSk3nLN9-NDUnRg'
|
||||
fileSize = 1.6 #in MB
|
||||
mFName = 'contextnet.pkl'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
file_id = '1cQvDPBKsz3TAi0Q5bJXsu6A-Z7lpk_cE'
|
||||
fileSize = 25.4 #in MB
|
||||
mFName = 'flownet.pkl'
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
file_id = '1mlA8VtxIcvJfz51OsQMvWX24oqxZ429r'
|
||||
fileSize = 15 #in MB
|
||||
mFName = 'unet.pkl'
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
#inpainting
|
||||
model = 'inpainting'
|
||||
file_id = '1WmPevEnRcdUynVHL8pZNzHPmLQVCFjuE'
|
||||
fileSize = 132 #in MB
|
||||
mFName = 'model_places2.pth'
|
||||
if not os.path.isdir(path + '/' + model):
|
||||
os.mkdir(path + '/' + model)
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
file_id = '1hIcPqDp8JjzR5kmt275DaVgX2PEtahWS'
|
||||
fileSize = 148 #in MB
|
||||
mFName = 'refinement.pth'
|
||||
destination = path + '/' + model + '/' + mFName
|
||||
if not os.path.isfile(destination):
|
||||
gimp.progress_init("Downloading " + model +"(~" + str(fileSize) + "MB)...")
|
||||
download_file_from_google_drive(file_id, destination,fileSize)
|
||||
|
||||
@ -1,48 +0,0 @@
|
||||
import os
|
||||
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
||||
|
||||
|
||||
from gimpfu import *
|
||||
import sys
|
||||
|
||||
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools',baseLoc])
|
||||
|
||||
import shutil
|
||||
import syncWeights
|
||||
|
||||
def update(flag) :
|
||||
gimp.progress_init("Updating plugins...")
|
||||
for filename in os.listdir(baseLoc):
|
||||
file_path = os.path.join(baseLoc, filename)
|
||||
try:
|
||||
if os.path.isfile(file_path) and not file_path.endswith('update.py'):
|
||||
os.unlink(file_path)
|
||||
elif os.path.isdir(file_path) and not (file_path.endswith('weights') or file_path.endswith('gimpenv')) :
|
||||
shutil.rmtree(file_path)
|
||||
except Exception as e:
|
||||
print('Failed to delete %s. Reason: %s' % (file_path, e))
|
||||
# os.system("cd "+baseLoc+";git fetch;git checkout .")
|
||||
syncWeights.syncGit(baseLoc)
|
||||
os.system("cd "+baseLoc+";chmod +x *.py")
|
||||
if flag:
|
||||
syncWeights.sync(baseLoc+'weights',flag)
|
||||
# pdb.gimp_message("Update Completed Successfully!")
|
||||
return
|
||||
|
||||
register(
|
||||
"update",
|
||||
"update",
|
||||
"update",
|
||||
"Kritik Soman",
|
||||
"Your Name",
|
||||
"2020",
|
||||
"update...",
|
||||
"",
|
||||
[(PF_BOOL, "wUpdate", "Update weights", True)],
|
||||
[],
|
||||
update, menu="<Image>/Layer/GIML-ML")
|
||||
|
||||
main()
|
||||
|
||||
|
||||
|
||||
@ -0,0 +1,75 @@
|
||||
Metadata-Version: 2.1
|
||||
Name: gimpml
|
||||
Version: 0.0.5
|
||||
Summary: A.I. for GIMP
|
||||
Home-page: https://github.com/kritiksoman/GIMP-ML
|
||||
Author: Kritik Soman
|
||||
Author-email: kritiksoman@ieee.org
|
||||
License: UNKNOWN
|
||||
Description: <img src="https://github.com/kritiksoman/tmp/blob/master/cover.png" width="1280" height="180"> <br>
|
||||
# Semantics for GNU Image Manipulation Program
|
||||
### [<img src="https://github.com/kritiksoman/tmp/blob/master/yt.png" width="70" height="50">](https://www.youtube.com/channel/UCzZn99R6Zh0ttGqvZieT4zw) [<img src="https://github.com/kritiksoman/tmp/blob/master/inst.png" width="50" height="50">](https://www.instagram.com/explore/tags/gimpml/) [<img src="https://github.com/kritiksoman/tmp/blob/master/arxiv.png" width="100" height="50">](https://arxiv.org/abs/2004.13060) [<img src="https://github.com/kritiksoman/tmp/blob/master/manual.png" width="100" height="50">](https://github.com/kritiksoman/GIMP-ML/wiki/User-Manual) [<img src="https://github.com/kritiksoman/tmp/blob/master/ref.png" width="100" height="50">](https://github.com/kritiksoman/GIMP-ML/wiki/References) [<img src="https://github.com/kritiksoman/tmp/blob/master/wiki.png" width="100" height="30">](https://en.wikipedia.org/wiki/GIMP#Extensions)<br>
|
||||
|
||||
:star: :star: :star: :star: are welcome. New tools will be added and existing will be improved with time.<br>
|
||||
|
||||
Updates: <br>
|
||||
[October 31] Use super-resolution as a filter for medium/large images. (Existing users should be able to update.)<br>
|
||||
[October 17] Added image enlightening.<br>
|
||||
[September 27] Added Force CPU use button and minor bug fixes. <br>
|
||||
[August 28] Added deep learning based dehazing and denoising. <br>
|
||||
[August 25] Simplified installation and updating method. <br>
|
||||
[August 2] Added deep matting and k-means. <br>
|
||||
[July 17] MonoDepth and Colorization models have been updated. <br>
|
||||
|
||||
# Screenshot of Menu
|
||||
|
||||
|
||||
# Installation Steps
|
||||
[1] Install [GIMP](https://www.gimp.org/downloads/) 2.10.<br>
|
||||
[2] Clone this repository: git clone https://github.com/kritiksoman/GIMP-ML.git <br>
|
||||
[3] Open terminal, go to GIMP-ML/gimp-plugins and run : <br>
|
||||
```bash installGimpML.sh```<br>
|
||||
[4] Open GIMP and go to Preferences -> Folders -> Plug-ins, add the folder gimp-plugins and restart GIMP. <br>
|
||||
[5] Go to Layer->GIMP-ML->update, click on ok with "update weights" set to yes and restart GIMP. (Weights ~ 1.5GB will be downloaded)<br>
|
||||
Manual install description if above is not working: [Link](https://github.com/kritiksoman/GIMP-ML/blob/master/INSTALLATION.md) <br>
|
||||
|
||||
# Update Steps
|
||||
[1] Go to Layer->GIMP-ML->update, click on ok with "update weights" set to NO and restart GIMP. <br>
|
||||
[2] Go to Layer->GIMP-ML->update, click on ok with "update weights" set to YES and restart GIMP. <br>
|
||||
|
||||
# Citation
|
||||
Please cite using the following bibtex entry:
|
||||
|
||||
```
|
||||
@article{soman2020GIMPML,
|
||||
title={GIMP-ML: Python Plugins for using Computer Vision Models in GIMP},
|
||||
author={Soman, Kritik},
|
||||
journal={arXiv preprint arXiv:2004.13060},
|
||||
year={2020}
|
||||
}
|
||||
```
|
||||
|
||||
# Tools
|
||||
| Name | License | Dataset |
|
||||
| ------------- |:-------------:| :-------------:|
|
||||
| facegen | [CC BY-NC-SA 4.0](https://github.com/switchablenorms/CelebAMask-HQ#dataset-agreement) | CelebAMask-HQ |
|
||||
| deblur | [BSD 3-clause](https://github.com/VITA-Group/DeblurGANv2/blob/master/LICENSE) | GoPro |
|
||||
| faceparse | [MIT](https://github.com/zllrunning/face-parsing.PyTorch/blob/master/LICENSE) | CelebAMask-HQ |
|
||||
| deepcolor | [MIT](https://github.com/junyanz/interactive-deep-colorization/blob/master/LICENSE) | ImageNet |
|
||||
| monodepth | [MIT](https://github.com/intel-isl/MiDaS/blob/master/LICENSE) | [Multiple](https://arxiv.org/pdf/1907.01341v3.pdf) |
|
||||
| super-resolution | [MIT](https://github.com/twtygqyy/pytorch-SRResNet/blob/master/LICENSE) | ImageNet |
|
||||
| deepmatting | [Non-commercial purposes](https://github.com/poppinace/indexnet_matting/blob/master/Adobe%20Deep%20Image%20Mattng%20Dataset%20License%20Agreement.pdf) | Adobe Deep Image Matting |
|
||||
| semantic-segmentation | MIT | COCO |
|
||||
| kmeans | [BSD](https://github.com/scipy/scipy/blob/master/LICENSE.txt) | - |
|
||||
| deep-dehazing | [MIT](https://github.com/MayankSingal/PyTorch-Image-Dehazing/blob/master/LICENSE) | [Custom](https://sites.google.com/site/boyilics/website-builder/project-page) |
|
||||
| deep-denoising | [GPL3](https://github.com/SaoYan/DnCNN-PyTorch/blob/master/LICENSE) | BSD68 |
|
||||
| enlighten | [BSD](https://github.com/VITA-Group/EnlightenGAN/blob/master/License) | [Custom](https://arxiv.org/pdf/1906.06972.pdf) |
|
||||
|
||||
Keywords: sample,setuptools,development
|
||||
Platform: UNKNOWN
|
||||
Classifier: Development Status :: 3 - Alpha
|
||||
Classifier: Intended Audience :: Developers
|
||||
Classifier: Topic :: Software Development :: Build Tools
|
||||
Classifier: Programming Language :: Python :: 3.8
|
||||
Requires-Python: >=2.7
|
||||
Description-Content-Type: text/markdown
|
||||
@ -0,0 +1,19 @@
|
||||
README.md
|
||||
setup.py
|
||||
gimpml/__init__.py
|
||||
gimpml.egg-info/PKG-INFO
|
||||
gimpml.egg-info/SOURCES.txt
|
||||
gimpml.egg-info/dependency_links.txt
|
||||
gimpml.egg-info/requires.txt
|
||||
gimpml.egg-info/top_level.txt
|
||||
gimpml/plugins/__init__.py
|
||||
gimpml/plugins/plugin_utils.py
|
||||
gimpml/plugins/monodepth/__init__.py
|
||||
gimpml/plugins/monodepth/monodepth.py
|
||||
gimpml/tools/__init__.py
|
||||
gimpml/tools/complete_install.py
|
||||
gimpml/tools/monodepth.py
|
||||
gimpml/tools/MiDaS/MiDaS_utils.py
|
||||
gimpml/tools/MiDaS/__init__.py
|
||||
gimpml/tools/MiDaS/mono_run.py
|
||||
gimpml/tools/MiDaS/monodepth_net.py
|
||||
@ -0,0 +1 @@
|
||||
|
||||
@ -0,0 +1,16 @@
|
||||
numpy
|
||||
scipy
|
||||
typing
|
||||
requests
|
||||
opencv-python<=4.3
|
||||
pretrainedmodels
|
||||
|
||||
[:python_version <= "2.7"]
|
||||
torchvision==0.5.0
|
||||
future
|
||||
torch==1.4.0
|
||||
enum
|
||||
|
||||
[:python_version > "2.7"]
|
||||
torchvision
|
||||
torch
|
||||
@ -0,0 +1 @@
|
||||
gimpml
|
||||
@ -0,0 +1,17 @@
|
||||
# from .kmeans import get_kmeans as kmeans
|
||||
from .tools.deblur import get_deblur as deblur
|
||||
# from .deepcolor import get_deepcolor as deepcolor
|
||||
from .tools.dehaze import get_dehaze as dehaze
|
||||
# from .deepdenoise import get_denoise as denoise
|
||||
# from .deepmatting import get_newalpha as matting
|
||||
from .tools.enlighten import get_enlighten as enlighten
|
||||
# from .facegen import get_newface as newface
|
||||
from .tools.faceparse import get_face as parseface
|
||||
# from .interpolateframes import get_inter as interpolateframe
|
||||
from .tools.monodepth import get_mono_depth as depth
|
||||
from .tools.complete_install import setup_python_weights
|
||||
# from .semseg import get_sem_seg as semseg
|
||||
from .tools.superresolution import get_super as super
|
||||
# from .inpainting import get_inpaint as inpaint
|
||||
# from .syncWeights import sync as sync
|
||||
|
||||
Binary file not shown.
@ -0,0 +1,44 @@
|
||||
import pickle
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
|
||||
plugin_loc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
||||
base_loc = os.path.expanduser("~") + '/GIMP-ML/'
|
||||
# base_loc = "D:/PycharmProjects/"
|
||||
sys.path.extend([plugin_loc + 'MiDaS'])
|
||||
# data_path = "D:/PycharmProjects/GIMP3-ML-pip/gimpml/"
|
||||
|
||||
from mono_run import run_depth
|
||||
from monodepth_net import MonoDepthNet
|
||||
import MiDaS_utils as MiDaS_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
|
||||
def get_mono_depth(input_image, cFlag = False):
|
||||
image = input_image / 255.0
|
||||
out = run_depth(image, base_loc + 'weights/MiDaS/model.pt', MonoDepthNet, MiDaS_utils, target_w=640, f=cFlag)
|
||||
out = np.repeat(out[:, :, np.newaxis], 3, axis=2)
|
||||
d1, d2 = input_image.shape[:2]
|
||||
out = cv2.resize(out, (d2, d1))
|
||||
# cv2.imwrite("/Users/kritiksoman/PycharmProjects/new/out.png", out)
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# # This will run when script is run as sub-process
|
||||
# dbfile = open(data_path + "data_input", 'rb')
|
||||
# data_input = pickle.load(dbfile)
|
||||
# dbfile.close()
|
||||
# # print(data)
|
||||
# data_output = {'args_input': {'processed': 1}, 'image_output': get_mono_depth(data_input['image'])}
|
||||
#
|
||||
# dbfile = open(data_path + "data_output", 'ab')
|
||||
# pickle.dump(data_output, dbfile) # source, destination
|
||||
# dbfile.close()
|
||||
|
||||
image = cv2.imread(os.path.join(base_loc, "cache.png"))[:, :, ::-1]
|
||||
output = get_mono_depth(image)
|
||||
cv2.imwrite(os.path.join(base_loc, 'cache.png'), output[:, :, ::-1])
|
||||
|
Before Width: | Height: | Size: 34 KiB After Width: | Height: | Size: 34 KiB |
@ -0,0 +1,105 @@
|
||||
#!/usr/bin/env python3
|
||||
#coding: utf-8
|
||||
"""
|
||||
.d8888b. 8888888 888b d888 8888888b. 888b d888 888
|
||||
d88P Y88b 888 8888b d8888 888 Y88b 8888b d8888 888
|
||||
888 888 888 88888b.d88888 888 888 88888b.d88888 888
|
||||
888 888 888Y88888P888 888 d88P 888Y88888P888 888
|
||||
888 88888 888 888 Y888P 888 8888888P" 888 Y888P 888 888
|
||||
888 888 888 888 Y8P 888 888 888 Y8P 888 888
|
||||
Y88b d88P 888 888 " 888 888 888 " 888 888
|
||||
"Y8888P88 8888888 888 888 888 888 888 88888888
|
||||
|
||||
|
||||
Opens the color palette as a new image file in GIMP.
|
||||
"""
|
||||
import gi
|
||||
|
||||
gi.require_version('Gimp', '3.0')
|
||||
from gi.repository import Gimp
|
||||
from gi.repository import GObject
|
||||
from gi.repository import GLib
|
||||
from gi.repository import Gio
|
||||
import time
|
||||
import sys
|
||||
import os
|
||||
import gettext
|
||||
|
||||
_ = gettext.gettext
|
||||
|
||||
|
||||
def N_(message): return message
|
||||
|
||||
|
||||
def colorpalette(procedure, run_mode, image, n_drawables, drawable, args, data):
|
||||
image_new = Gimp.Image.new(1200, 675, 0) # 0 for RGB
|
||||
display = Gimp.Display.new(image_new)
|
||||
|
||||
|
||||
result = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(
|
||||
os.path.join(os.path.dirname(os.path.realpath(__file__)), 'color_palette.png')))
|
||||
result_layer = result.get_active_layer()
|
||||
copy = Gimp.Layer.new_from_drawable(result_layer, image_new)
|
||||
copy.set_name("Color Palette")
|
||||
copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)# DIFFERENCE_LEGACY
|
||||
image_new.insert_layer(copy, None, -1)
|
||||
Gimp.displays_flush()
|
||||
|
||||
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
|
||||
|
||||
|
||||
class ColorPalette(Gimp.PlugIn):
|
||||
## Parameters ##
|
||||
__gproperties__ = {
|
||||
# "name": (str,
|
||||
# _("Layer name"),
|
||||
# _("Layer name"),
|
||||
# _("Clouds"),
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "color": (Gimp.RGB,
|
||||
# _("Fog color"),
|
||||
# _("Fog color"),
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "turbulence": (float,
|
||||
# _("Turbulence"),
|
||||
# _("Turbulence"),
|
||||
# 0.0, 10.0, 1.0,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "opacity": (float,
|
||||
# _("Opacity"),
|
||||
# _("Opacity"),
|
||||
# 0.0, 100.0, 100.0,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
}
|
||||
|
||||
## GimpPlugIn virtual methods ##
|
||||
def do_query_procedures(self):
|
||||
self.set_translation_domain("gimp30-python",
|
||||
Gio.file_new_for_path(Gimp.locale_directory()))
|
||||
|
||||
return ['colorpalette']
|
||||
|
||||
def do_create_procedure(self, name):
|
||||
procedure = Gimp.ImageProcedure.new(self, name,
|
||||
Gimp.PDBProcType.PLUGIN,
|
||||
colorpalette, None)
|
||||
procedure.set_image_types("RGB*, GRAY*");
|
||||
procedure.set_documentation(N_("Add a layer of fog"),
|
||||
"Adds a layer of fog to the image.",
|
||||
name)
|
||||
procedure.set_menu_label(N_("_Color Palette..."))
|
||||
procedure.set_attribution("Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2021")
|
||||
procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
|
||||
|
||||
# procedure.add_argument_from_property(self, "name")
|
||||
# TODO: add support for GBoxed values.
|
||||
# procedure.add_argument_from_property(self, "color")
|
||||
# procedure.add_argument_from_property(self, "turbulence")
|
||||
# procedure.add_argument_from_property(self, "opacity")
|
||||
return procedure
|
||||
|
||||
|
||||
Gimp.main(ColorPalette.__gtype__, sys.argv)
|
||||
@ -0,0 +1,227 @@
|
||||
#!/usr/bin/env python3
|
||||
#coding: utf-8
|
||||
"""
|
||||
.d8888b. 8888888 888b d888 8888888b. 888b d888 888
|
||||
d88P Y88b 888 8888b d8888 888 Y88b 8888b d8888 888
|
||||
888 888 888 88888b.d88888 888 888 88888b.d88888 888
|
||||
888 888 888Y88888P888 888 d88P 888Y88888P888 888
|
||||
888 88888 888 888 Y888P 888 8888888P" 888 Y888P 888 888
|
||||
888 888 888 888 Y8P 888 888 888 Y8P 888 888
|
||||
Y88b d88P 888 888 " 888 888 888 " 888 888
|
||||
"Y8888P88 8888888 888 888 888 888 888 88888888
|
||||
|
||||
|
||||
Deblur the current layer.
|
||||
"""
|
||||
import sys
|
||||
import gi
|
||||
gi.require_version('Gimp', '3.0')
|
||||
from gi.repository import Gimp
|
||||
gi.require_version('GimpUi', '3.0')
|
||||
from gi.repository import GimpUi
|
||||
from gi.repository import GObject
|
||||
from gi.repository import GLib
|
||||
from gi.repository import Gio
|
||||
gi.require_version('Gtk', '3.0')
|
||||
from gi.repository import Gtk
|
||||
|
||||
import gettext
|
||||
_ = gettext.gettext
|
||||
def N_(message): return message
|
||||
|
||||
import subprocess
|
||||
import pickle
|
||||
import os
|
||||
|
||||
def deblur(procedure, image, drawable, force_cpu, progress_bar):
|
||||
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools")
|
||||
with open(os.path.join(config_path, 'gimp_ml_config.pkl'), 'rb') as file:
|
||||
data_output = pickle.load(file)
|
||||
weight_path = data_output["weight_path"]
|
||||
python_path = data_output["python_path"]
|
||||
plugin_path = os.path.join(config_path, 'deblur.py')
|
||||
|
||||
Gimp.context_push()
|
||||
image.undo_group_start()
|
||||
|
||||
interlace, compression = 0, 2
|
||||
Gimp.get_pdb().run_procedure('file-png-save', [
|
||||
GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
|
||||
GObject.Value(Gimp.Image, image),
|
||||
GObject.Value(GObject.TYPE_INT, 1),
|
||||
GObject.Value(Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, drawable, 1)),
|
||||
GObject.Value(Gio.File, Gio.File.new_for_path(os.path.join(weight_path, '..', 'cache.png'))),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, interlace),
|
||||
GObject.Value(GObject.TYPE_INT, compression),
|
||||
# write all PNG chunks except oFFs(ets)
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, False),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
])
|
||||
|
||||
with open(os.path.join(weight_path, '..', 'gimp_ml_run.pkl'), 'wb') as file:
|
||||
pickle.dump({"force_cpu": bool(force_cpu)}, file)
|
||||
|
||||
subprocess.call([python_path, plugin_path])
|
||||
|
||||
result = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(os.path.join(weight_path, '..', 'cache.png')))
|
||||
result_layer = result.get_active_layer()
|
||||
copy = Gimp.Layer.new_from_drawable(result_layer, image)
|
||||
copy.set_name("Deblur")
|
||||
copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)#DIFFERENCE_LEGACY
|
||||
image.insert_layer(copy, None, -1)
|
||||
|
||||
image.undo_group_end()
|
||||
Gimp.context_pop()
|
||||
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
|
||||
|
||||
|
||||
|
||||
def run(procedure, run_mode, image, n_drawables, layer, args, data):
|
||||
# gio_file = args.index(0)
|
||||
# bucket_size = args.index(0)
|
||||
force_cpu = args.index(1)
|
||||
# output_format = args.index(2)
|
||||
|
||||
progress_bar = None
|
||||
config = None
|
||||
|
||||
if run_mode == Gimp.RunMode.INTERACTIVE:
|
||||
|
||||
config = procedure.create_config()
|
||||
|
||||
# Set properties from arguments. These properties will be changed by the UI.
|
||||
#config.set_property("file", gio_file)
|
||||
#config.set_property("bucket_size", bucket_size)
|
||||
config.set_property("force_cpu", force_cpu)
|
||||
#config.set_property("output_format", output_format)
|
||||
config.begin_run(image, run_mode, args)
|
||||
|
||||
GimpUi.init("deblur.py")
|
||||
use_header_bar = Gtk.Settings.get_default().get_property("gtk-dialogs-use-header")
|
||||
dialog = GimpUi.Dialog(use_header_bar=use_header_bar,
|
||||
title=_("Deblur..."))
|
||||
dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
|
||||
dialog.add_button("_OK", Gtk.ResponseType.OK)
|
||||
|
||||
vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL,
|
||||
homogeneous=False, spacing=10)
|
||||
dialog.get_content_area().add(vbox)
|
||||
vbox.show()
|
||||
|
||||
# Create grid to set all the properties inside.
|
||||
grid = Gtk.Grid()
|
||||
grid.set_column_homogeneous(False)
|
||||
grid.set_border_width(10)
|
||||
grid.set_column_spacing(10)
|
||||
grid.set_row_spacing(10)
|
||||
vbox.add(grid)
|
||||
grid.show()
|
||||
|
||||
# # Bucket size parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Bucket Size"))
|
||||
# grid.attach(label, 0, 1, 1, 1)
|
||||
# label.show()
|
||||
# spin = GimpUi.prop_spin_button_new(config, "bucket_size", step_increment=0.001, page_increment=0.1, digits=3)
|
||||
# grid.attach(spin, 1, 1, 1, 1)
|
||||
# spin.show()
|
||||
|
||||
# Force CPU parameter
|
||||
spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
|
||||
spin.set_tooltip_text(_("If checked, CPU is used for model inference."
|
||||
" Otherwise, GPU will be used if available."))
|
||||
grid.attach(spin, 1, 2, 1, 1)
|
||||
spin.show()
|
||||
|
||||
# # Output format parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Output Format"))
|
||||
# grid.attach(label, 0, 3, 1, 1)
|
||||
# label.show()
|
||||
# combo = GimpUi.prop_string_combo_box_new(config, "output_format", output_format_enum.get_tree_model(), 0, 1)
|
||||
# grid.attach(combo, 1, 3, 1, 1)
|
||||
# combo.show()
|
||||
|
||||
progress_bar = Gtk.ProgressBar()
|
||||
vbox.add(progress_bar)
|
||||
progress_bar.show()
|
||||
|
||||
dialog.show()
|
||||
if dialog.run() != Gtk.ResponseType.OK:
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.CANCEL,
|
||||
GLib.Error())
|
||||
|
||||
result = deblur(procedure, image, layer, force_cpu, progress_bar)
|
||||
|
||||
# If the execution was successful, save parameters so they will be restored next time we show dialog.
|
||||
if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
|
||||
config.end_run(Gimp.PDBStatusType.SUCCESS)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Deblur(Gimp.PlugIn):
|
||||
|
||||
## Parameters ##
|
||||
__gproperties__ = {
|
||||
# "filename": (str,
|
||||
# # TODO: I wanted this property to be a path (and not just str) , so I could use
|
||||
# # prop_file_chooser_button_new to open a file dialog. However, it fails without an error message.
|
||||
# # Gimp.ConfigPath,
|
||||
# _("Histogram _File"),
|
||||
# _("Histogram _File"),
|
||||
# "deblur.csv",
|
||||
# # Gimp.ConfigPathType.FILE,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "file": (Gio.File,
|
||||
# _("Histogram _File"),
|
||||
# "Histogram export file",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "bucket_size": (float,
|
||||
# _("_Bucket Size"),
|
||||
# "Bucket Size",
|
||||
# 0.001, 1.0, 0.01,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
"force_cpu": (bool,
|
||||
_("Force _CPU"),
|
||||
"Force CPU",
|
||||
False,
|
||||
GObject.ParamFlags.READWRITE),
|
||||
# "output_format": (str,
|
||||
# _("Output format"),
|
||||
# "Output format: 'pixel count', 'normalized', 'percent'",
|
||||
# "pixel count",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
}
|
||||
|
||||
## GimpPlugIn virtual methods ##
|
||||
def do_query_procedures(self):
|
||||
self.set_translation_domain("gimp30-python",
|
||||
Gio.file_new_for_path(Gimp.locale_directory()))
|
||||
return ['deblur']
|
||||
|
||||
def do_create_procedure(self, name):
|
||||
procedure = None
|
||||
if name == 'deblur':
|
||||
procedure = Gimp.ImageProcedure.new(self, name, Gimp.PDBProcType.PLUGIN, run, None)
|
||||
procedure.set_image_types("*")
|
||||
procedure.set_documentation (
|
||||
N_("Deblur the current layer."),
|
||||
globals()["__doc__"], # This includes the docstring, on the top of the file
|
||||
name)
|
||||
procedure.set_menu_label(N_("Deblur..."))
|
||||
procedure.set_attribution("Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2021")
|
||||
procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
|
||||
|
||||
# procedure.add_argument_from_property(self, "file")
|
||||
# procedure.add_argument_from_property(self, "bucket_size")
|
||||
procedure.add_argument_from_property(self, "force_cpu")
|
||||
# procedure.add_argument_from_property(self, "output_format")
|
||||
|
||||
return procedure
|
||||
|
||||
|
||||
Gimp.main(Deblur.__gtype__, sys.argv)
|
||||
@ -0,0 +1,227 @@
|
||||
#!/usr/bin/env python3
|
||||
#coding: utf-8
|
||||
"""
|
||||
.d8888b. 8888888 888b d888 8888888b. 888b d888 888
|
||||
d88P Y88b 888 8888b d8888 888 Y88b 8888b d8888 888
|
||||
888 888 888 88888b.d88888 888 888 88888b.d88888 888
|
||||
888 888 888Y88888P888 888 d88P 888Y88888P888 888
|
||||
888 88888 888 888 Y888P 888 8888888P" 888 Y888P 888 888
|
||||
888 888 888 888 Y8P 888 888 888 Y8P 888 888
|
||||
Y88b d88P 888 888 " 888 888 888 " 888 888
|
||||
"Y8888P88 8888888 888 888 888 888 888 88888888
|
||||
|
||||
|
||||
Dehazes the current layer.
|
||||
"""
|
||||
import sys
|
||||
import gi
|
||||
gi.require_version('Gimp', '3.0')
|
||||
from gi.repository import Gimp
|
||||
gi.require_version('GimpUi', '3.0')
|
||||
from gi.repository import GimpUi
|
||||
from gi.repository import GObject
|
||||
from gi.repository import GLib
|
||||
from gi.repository import Gio
|
||||
gi.require_version('Gtk', '3.0')
|
||||
from gi.repository import Gtk
|
||||
|
||||
import gettext
|
||||
_ = gettext.gettext
|
||||
def N_(message): return message
|
||||
|
||||
import subprocess
|
||||
import pickle
|
||||
import os
|
||||
|
||||
def dehaze(procedure, image, drawable, force_cpu, progress_bar):
|
||||
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools")
|
||||
with open(os.path.join(config_path, 'gimp_ml_config.pkl'), 'rb') as file:
|
||||
data_output = pickle.load(file)
|
||||
weight_path = data_output["weight_path"]
|
||||
python_path = data_output["python_path"]
|
||||
plugin_path = os.path.join(config_path, 'dehaze.py')
|
||||
|
||||
Gimp.context_push()
|
||||
image.undo_group_start()
|
||||
|
||||
interlace, compression = 0, 2
|
||||
Gimp.get_pdb().run_procedure('file-png-save', [
|
||||
GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
|
||||
GObject.Value(Gimp.Image, image),
|
||||
GObject.Value(GObject.TYPE_INT, 1),
|
||||
GObject.Value(Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, drawable, 1)),
|
||||
GObject.Value(Gio.File, Gio.File.new_for_path(os.path.join(weight_path, '..', 'cache.png'))),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, interlace),
|
||||
GObject.Value(GObject.TYPE_INT, compression),
|
||||
# write all PNG chunks except oFFs(ets)
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, False),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
])
|
||||
|
||||
with open(os.path.join(weight_path, '..', 'gimp_ml_run.pkl'), 'wb') as file:
|
||||
pickle.dump({"force_cpu": bool(force_cpu)}, file)
|
||||
|
||||
subprocess.call([python_path, plugin_path])
|
||||
|
||||
result = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(os.path.join(weight_path, '..', 'cache.png')))
|
||||
result_layer = result.get_active_layer()
|
||||
copy = Gimp.Layer.new_from_drawable(result_layer, image)
|
||||
copy.set_name("Dehaze")
|
||||
copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)#DIFFERENCE_LEGACY
|
||||
image.insert_layer(copy, None, -1)
|
||||
|
||||
image.undo_group_end()
|
||||
Gimp.context_pop()
|
||||
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
|
||||
|
||||
|
||||
|
||||
def run(procedure, run_mode, image, n_drawables, layer, args, data):
|
||||
# gio_file = args.index(0)
|
||||
# bucket_size = args.index(0)
|
||||
force_cpu = args.index(1)
|
||||
# output_format = args.index(2)
|
||||
|
||||
progress_bar = None
|
||||
config = None
|
||||
|
||||
if run_mode == Gimp.RunMode.INTERACTIVE:
|
||||
|
||||
config = procedure.create_config()
|
||||
|
||||
# Set properties from arguments. These properties will be changed by the UI.
|
||||
#config.set_property("file", gio_file)
|
||||
#config.set_property("bucket_size", bucket_size)
|
||||
config.set_property("force_cpu", force_cpu)
|
||||
#config.set_property("output_format", output_format)
|
||||
config.begin_run(image, run_mode, args)
|
||||
|
||||
GimpUi.init("dehaze.py")
|
||||
use_header_bar = Gtk.Settings.get_default().get_property("gtk-dialogs-use-header")
|
||||
dialog = GimpUi.Dialog(use_header_bar=use_header_bar,
|
||||
title=_("Dehaze..."))
|
||||
dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
|
||||
dialog.add_button("_OK", Gtk.ResponseType.OK)
|
||||
|
||||
vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL,
|
||||
homogeneous=False, spacing=10)
|
||||
dialog.get_content_area().add(vbox)
|
||||
vbox.show()
|
||||
|
||||
# Create grid to set all the properties inside.
|
||||
grid = Gtk.Grid()
|
||||
grid.set_column_homogeneous(False)
|
||||
grid.set_border_width(10)
|
||||
grid.set_column_spacing(10)
|
||||
grid.set_row_spacing(10)
|
||||
vbox.add(grid)
|
||||
grid.show()
|
||||
|
||||
# # Bucket size parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Bucket Size"))
|
||||
# grid.attach(label, 0, 1, 1, 1)
|
||||
# label.show()
|
||||
# spin = GimpUi.prop_spin_button_new(config, "bucket_size", step_increment=0.001, page_increment=0.1, digits=3)
|
||||
# grid.attach(spin, 1, 1, 1, 1)
|
||||
# spin.show()
|
||||
|
||||
# Force CPU parameter
|
||||
spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
|
||||
spin.set_tooltip_text(_("If checked, CPU is used for model inference."
|
||||
" Otherwise, GPU will be used if available."))
|
||||
grid.attach(spin, 1, 2, 1, 1)
|
||||
spin.show()
|
||||
|
||||
# # Output format parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Output Format"))
|
||||
# grid.attach(label, 0, 3, 1, 1)
|
||||
# label.show()
|
||||
# combo = GimpUi.prop_string_combo_box_new(config, "output_format", output_format_enum.get_tree_model(), 0, 1)
|
||||
# grid.attach(combo, 1, 3, 1, 1)
|
||||
# combo.show()
|
||||
|
||||
progress_bar = Gtk.ProgressBar()
|
||||
vbox.add(progress_bar)
|
||||
progress_bar.show()
|
||||
|
||||
dialog.show()
|
||||
if dialog.run() != Gtk.ResponseType.OK:
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.CANCEL,
|
||||
GLib.Error())
|
||||
|
||||
result = dehaze(procedure, image, layer, force_cpu, progress_bar)
|
||||
|
||||
# If the execution was successful, save parameters so they will be restored next time we show dialog.
|
||||
if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
|
||||
config.end_run(Gimp.PDBStatusType.SUCCESS)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Dehaze(Gimp.PlugIn):
|
||||
|
||||
## Parameters ##
|
||||
__gproperties__ = {
|
||||
# "filename": (str,
|
||||
# # TODO: I wanted this property to be a path (and not just str) , so I could use
|
||||
# # prop_file_chooser_button_new to open a file dialog. However, it fails without an error message.
|
||||
# # Gimp.ConfigPath,
|
||||
# _("Histogram _File"),
|
||||
# _("Histogram _File"),
|
||||
# "dehaze.csv",
|
||||
# # Gimp.ConfigPathType.FILE,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "file": (Gio.File,
|
||||
# _("Histogram _File"),
|
||||
# "Histogram export file",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "bucket_size": (float,
|
||||
# _("_Bucket Size"),
|
||||
# "Bucket Size",
|
||||
# 0.001, 1.0, 0.01,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
"force_cpu": (bool,
|
||||
_("Force _CPU"),
|
||||
"Force CPU",
|
||||
False,
|
||||
GObject.ParamFlags.READWRITE),
|
||||
# "output_format": (str,
|
||||
# _("Output format"),
|
||||
# "Output format: 'pixel count', 'normalized', 'percent'",
|
||||
# "pixel count",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
}
|
||||
|
||||
## GimpPlugIn virtual methods ##
|
||||
def do_query_procedures(self):
|
||||
self.set_translation_domain("gimp30-python",
|
||||
Gio.file_new_for_path(Gimp.locale_directory()))
|
||||
return ['dehaze']
|
||||
|
||||
def do_create_procedure(self, name):
|
||||
procedure = None
|
||||
if name == 'dehaze':
|
||||
procedure = Gimp.ImageProcedure.new(self, name, Gimp.PDBProcType.PLUGIN, run, None)
|
||||
procedure.set_image_types("*")
|
||||
procedure.set_documentation (
|
||||
N_("Dehazes the current layer."),
|
||||
globals()["__doc__"], # This includes the docstring, on the top of the file
|
||||
name)
|
||||
procedure.set_menu_label(N_("_Dehaze..."))
|
||||
procedure.set_attribution("Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2021")
|
||||
procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
|
||||
|
||||
# procedure.add_argument_from_property(self, "file")
|
||||
# procedure.add_argument_from_property(self, "bucket_size")
|
||||
procedure.add_argument_from_property(self, "force_cpu")
|
||||
# procedure.add_argument_from_property(self, "output_format")
|
||||
|
||||
return procedure
|
||||
|
||||
|
||||
Gimp.main(Dehaze.__gtype__, sys.argv)
|
||||
@ -0,0 +1,227 @@
|
||||
#!/usr/bin/env python3
|
||||
#coding: utf-8
|
||||
"""
|
||||
.d8888b. 8888888 888b d888 8888888b. 888b d888 888
|
||||
d88P Y88b 888 8888b d8888 888 Y88b 8888b d8888 888
|
||||
888 888 888 88888b.d88888 888 888 88888b.d88888 888
|
||||
888 888 888Y88888P888 888 d88P 888Y88888P888 888
|
||||
888 88888 888 888 Y888P 888 8888888P" 888 Y888P 888 888
|
||||
888 888 888 888 Y8P 888 888 888 Y8P 888 888
|
||||
Y88b d88P 888 888 " 888 888 888 " 888 888
|
||||
"Y8888P88 8888888 888 888 888 888 888 88888888
|
||||
|
||||
|
||||
denoises the current layer.
|
||||
"""
|
||||
import sys
|
||||
import gi
|
||||
gi.require_version('Gimp', '3.0')
|
||||
from gi.repository import Gimp
|
||||
gi.require_version('GimpUi', '3.0')
|
||||
from gi.repository import GimpUi
|
||||
from gi.repository import GObject
|
||||
from gi.repository import GLib
|
||||
from gi.repository import Gio
|
||||
gi.require_version('Gtk', '3.0')
|
||||
from gi.repository import Gtk
|
||||
|
||||
import gettext
|
||||
_ = gettext.gettext
|
||||
def N_(message): return message
|
||||
|
||||
import subprocess
|
||||
import pickle
|
||||
import os
|
||||
|
||||
def denoise(procedure, image, drawable, force_cpu, progress_bar):
|
||||
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "tools")
|
||||
with open(os.path.join(config_path, 'gimp_ml_config.pkl'), 'rb') as file:
|
||||
data_output = pickle.load(file)
|
||||
weight_path = data_output["weight_path"]
|
||||
python_path = data_output["python_path"]
|
||||
plugin_path = os.path.join(config_path, 'denoise.py')
|
||||
|
||||
Gimp.context_push()
|
||||
image.undo_group_start()
|
||||
|
||||
interlace, compression = 0, 2
|
||||
Gimp.get_pdb().run_procedure('file-png-save', [
|
||||
GObject.Value(Gimp.RunMode, Gimp.RunMode.NONINTERACTIVE),
|
||||
GObject.Value(Gimp.Image, image),
|
||||
GObject.Value(GObject.TYPE_INT, 1),
|
||||
GObject.Value(Gimp.ObjectArray, Gimp.ObjectArray.new(Gimp.Drawable, drawable, 1)),
|
||||
GObject.Value(Gio.File, Gio.File.new_for_path(os.path.join(weight_path, '..', 'cache.png'))),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, interlace),
|
||||
GObject.Value(GObject.TYPE_INT, compression),
|
||||
# write all PNG chunks except oFFs(ets)
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, False),
|
||||
GObject.Value(GObject.TYPE_BOOLEAN, True),
|
||||
])
|
||||
|
||||
with open(os.path.join(weight_path, '..', 'gimp_ml_run.pkl'), 'wb') as file:
|
||||
pickle.dump({"force_cpu": bool(force_cpu)}, file)
|
||||
|
||||
subprocess.call([python_path, plugin_path])
|
||||
|
||||
result = Gimp.file_load(Gimp.RunMode.NONINTERACTIVE, Gio.file_new_for_path(os.path.join(weight_path, '..', 'cache.png')))
|
||||
result_layer = result.get_active_layer()
|
||||
copy = Gimp.Layer.new_from_drawable(result_layer, image)
|
||||
copy.set_name("Denoise")
|
||||
copy.set_mode(Gimp.LayerMode.NORMAL_LEGACY)#DIFFERENCE_LEGACY
|
||||
image.insert_layer(copy, None, -1)
|
||||
|
||||
image.undo_group_end()
|
||||
Gimp.context_pop()
|
||||
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.SUCCESS, GLib.Error())
|
||||
|
||||
|
||||
|
||||
def run(procedure, run_mode, image, n_drawables, layer, args, data):
|
||||
# gio_file = args.index(0)
|
||||
# bucket_size = args.index(0)
|
||||
force_cpu = args.index(1)
|
||||
# output_format = args.index(2)
|
||||
|
||||
progress_bar = None
|
||||
config = None
|
||||
|
||||
if run_mode == Gimp.RunMode.INTERACTIVE:
|
||||
|
||||
config = procedure.create_config()
|
||||
|
||||
# Set properties from arguments. These properties will be changed by the UI.
|
||||
#config.set_property("file", gio_file)
|
||||
#config.set_property("bucket_size", bucket_size)
|
||||
config.set_property("force_cpu", force_cpu)
|
||||
#config.set_property("output_format", output_format)
|
||||
config.begin_run(image, run_mode, args)
|
||||
|
||||
GimpUi.init("denoise.py")
|
||||
use_header_bar = Gtk.Settings.get_default().get_property("gtk-dialogs-use-header")
|
||||
dialog = GimpUi.Dialog(use_header_bar=use_header_bar,
|
||||
title=_("Denoise..."))
|
||||
dialog.add_button("_Cancel", Gtk.ResponseType.CANCEL)
|
||||
dialog.add_button("_OK", Gtk.ResponseType.OK)
|
||||
|
||||
vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL,
|
||||
homogeneous=False, spacing=10)
|
||||
dialog.get_content_area().add(vbox)
|
||||
vbox.show()
|
||||
|
||||
# Create grid to set all the properties inside.
|
||||
grid = Gtk.Grid()
|
||||
grid.set_column_homogeneous(False)
|
||||
grid.set_border_width(10)
|
||||
grid.set_column_spacing(10)
|
||||
grid.set_row_spacing(10)
|
||||
vbox.add(grid)
|
||||
grid.show()
|
||||
|
||||
# # Bucket size parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Bucket Size"))
|
||||
# grid.attach(label, 0, 1, 1, 1)
|
||||
# label.show()
|
||||
# spin = GimpUi.prop_spin_button_new(config, "bucket_size", step_increment=0.001, page_increment=0.1, digits=3)
|
||||
# grid.attach(spin, 1, 1, 1, 1)
|
||||
# spin.show()
|
||||
|
||||
# Force CPU parameter
|
||||
spin = GimpUi.prop_check_button_new(config, "force_cpu", _("Force _CPU"))
|
||||
spin.set_tooltip_text(_("If checked, CPU is used for model inference."
|
||||
" Otherwise, GPU will be used if available."))
|
||||
grid.attach(spin, 1, 2, 1, 1)
|
||||
spin.show()
|
||||
|
||||
# # Output format parameter
|
||||
# label = Gtk.Label.new_with_mnemonic(_("_Output Format"))
|
||||
# grid.attach(label, 0, 3, 1, 1)
|
||||
# label.show()
|
||||
# combo = GimpUi.prop_string_combo_box_new(config, "output_format", output_format_enum.get_tree_model(), 0, 1)
|
||||
# grid.attach(combo, 1, 3, 1, 1)
|
||||
# combo.show()
|
||||
|
||||
progress_bar = Gtk.ProgressBar()
|
||||
vbox.add(progress_bar)
|
||||
progress_bar.show()
|
||||
|
||||
dialog.show()
|
||||
if dialog.run() != Gtk.ResponseType.OK:
|
||||
return procedure.new_return_values(Gimp.PDBStatusType.CANCEL,
|
||||
GLib.Error())
|
||||
|
||||
result = denoise(procedure, image, layer, force_cpu, progress_bar)
|
||||
|
||||
# If the execution was successful, save parameters so they will be restored next time we show dialog.
|
||||
if result.index(0) == Gimp.PDBStatusType.SUCCESS and config is not None:
|
||||
config.end_run(Gimp.PDBStatusType.SUCCESS)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Denoise(Gimp.PlugIn):
|
||||
|
||||
## Parameters ##
|
||||
__gproperties__ = {
|
||||
# "filename": (str,
|
||||
# # TODO: I wanted this property to be a path (and not just str) , so I could use
|
||||
# # prop_file_chooser_button_new to open a file dialog. However, it fails without an error message.
|
||||
# # Gimp.ConfigPath,
|
||||
# _("Histogram _File"),
|
||||
# _("Histogram _File"),
|
||||
# "denoise.csv",
|
||||
# # Gimp.ConfigPathType.FILE,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "file": (Gio.File,
|
||||
# _("Histogram _File"),
|
||||
# "Histogram export file",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
# "bucket_size": (float,
|
||||
# _("_Bucket Size"),
|
||||
# "Bucket Size",
|
||||
# 0.001, 1.0, 0.01,
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
"force_cpu": (bool,
|
||||
_("Force _CPU"),
|
||||
"Force CPU",
|
||||
False,
|
||||
GObject.ParamFlags.READWRITE),
|
||||
# "output_format": (str,
|
||||
# _("Output format"),
|
||||
# "Output format: 'pixel count', 'normalized', 'percent'",
|
||||
# "pixel count",
|
||||
# GObject.ParamFlags.READWRITE),
|
||||
}
|
||||
|
||||
## GimpPlugIn virtual methods ##
|
||||
def do_query_procedures(self):
|
||||
self.set_translation_domain("gimp30-python",
|
||||
Gio.file_new_for_path(Gimp.locale_directory()))
|
||||
return ['denoise']
|
||||
|
||||
def do_create_procedure(self, name):
|
||||
procedure = None
|
||||
if name == 'denoise':
|
||||
procedure = Gimp.ImageProcedure.new(self, name, Gimp.PDBProcType.PLUGIN, run, None)
|
||||
procedure.set_image_types("*")
|
||||
procedure.set_documentation (
|
||||
N_("Denoises the current layer."),
|
||||
globals()["__doc__"], # This includes the docstring, on the top of the file
|
||||
name)
|
||||
procedure.set_menu_label(N_("Denoise..."))
|
||||
procedure.set_attribution("Kritik Soman",
|
||||
"GIMP-ML",
|
||||
"2021")
|
||||
procedure.add_menu_path("<Image>/Layer/GIMP-ML/")
|
||||
|
||||
# procedure.add_argument_from_property(self, "file")
|
||||
# procedure.add_argument_from_property(self, "bucket_size")
|
||||
procedure.add_argument_from_property(self, "force_cpu")
|
||||
# procedure.add_argument_from_property(self, "output_format")
|
||||
|
||||
return procedure
|
||||
|
||||
|
||||
Gimp.main(Denoise.__gtype__, sys.argv)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue