Forked from iperov/DeepFaceLab

This commit is contained in:
Alex
2020-08-04 01:04:26 +03:00
parent 3e7ee22ae3
commit 71ceb60719
177 changed files with 1 additions and 20433 deletions
-19
View File
@@ -1,19 +0,0 @@
THIS IS NOT TECH SUPPORT FOR NEWBIE FAKERS
POST ONLY ISSUES RELATED TO BUGS OR CODE
## Expected behavior
*Describe, in some detail, what you are trying to do and what the output is that you expect from the program.*
## Actual behavior
*Describe, in some detail, what the program does instead. Be sure to include any error message or screenshots.*
## Steps to reproduce
*Describe, in some detail, the steps you tried that resulted in the behavior described above.*
## Other relevant information
- **Command lined used (if not specified in steps to reproduce)**: main.py ...
- **Operating system and version:** Windows, macOS, Linux
- **Python version:** 3.5, 3.6.4, ... (if you are not using prebuilt windows binary)
-8
View File
@@ -1,8 +0,0 @@
*
!*.py
!*.md
!*.txt
!*.jpg
!requirements*
!Dockerfile*
!*.sh
-26
View File
@@ -1,26 +0,0 @@
{
// Use IntelliSense to learn about possible attributes.
// Hover to view descriptions of existing attributes.
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0",
"configurations": [
{
"name": "DFL train TEST",
"subProcess": true,
"justMyCode": true,
"type": "python",
"request": "launch",
"program": "${env:DFL_ROOT}\\main.py",
"pythonPath": "${env:PYTHONEXECUTABLE}",
"cwd": "${env:WORKSPACE}",
"console": "integratedTerminal",
"args": ["train",
"--training-data-src-dir", "${env:WORKSPACE}\\data_src\\aligned",
"--training-data-dst-dir", "${env:WORKSPACE}\\data_dst\\aligned",
"--model-dir", "${env:WORKSPACE}\\model",
"--model", "TEST"
]
}
]
}
View File
View File
-5
View File
@@ -1,5 +0,0 @@
Please don't ruin the code and this good (as I think) architecture.
Please follow the same logic and brevity/pithiness.
Don't abstract the code into huge classes if you only win some lines of code in one place, because this can prevent programmers from understanding it quickly.
-12
View File
@@ -1,12 +0,0 @@
from pathlib import Path
from .DFLJPG import DFLJPG
class DFLIMG():
@staticmethod
def load(filepath, loader_func=None):
if filepath.suffix == '.jpg':
return DFLJPG.load ( str(filepath), loader_func=loader_func )
else:
return None
-317
View File
@@ -1,317 +0,0 @@
import pickle
import struct
import traceback
import cv2
import numpy as np
from core import imagelib
from core.cv2ex import *
from core.imagelib import SegIEPolys
from core.interact import interact as io
from core.structex import *
from facelib import FaceType
class DFLJPG(object):
def __init__(self, filename):
self.filename = filename
self.data = b""
self.length = 0
self.chunks = []
self.dfl_dict = None
self.shape = None
self.img = None
@staticmethod
def load_raw(filename, loader_func=None):
try:
if loader_func is not None:
data = loader_func(filename)
else:
with open(filename, "rb") as f:
data = f.read()
except:
raise FileNotFoundError(filename)
try:
inst = DFLJPG(filename)
inst.data = data
inst.length = len(data)
inst_length = inst.length
chunks = []
data_counter = 0
while data_counter < inst_length:
chunk_m_l, chunk_m_h = struct.unpack ("BB", data[data_counter:data_counter+2])
data_counter += 2
if chunk_m_l != 0xFF:
raise ValueError(f"No Valid JPG info in {filename}")
chunk_name = None
chunk_size = None
chunk_data = None
chunk_ex_data = None
is_unk_chunk = False
if chunk_m_h & 0xF0 == 0xD0:
n = chunk_m_h & 0x0F
if n >= 0 and n <= 7:
chunk_name = "RST%d" % (n)
chunk_size = 0
elif n == 0x8:
chunk_name = "SOI"
chunk_size = 0
if len(chunks) != 0:
raise Exception("")
elif n == 0x9:
chunk_name = "EOI"
chunk_size = 0
elif n == 0xA:
chunk_name = "SOS"
elif n == 0xB:
chunk_name = "DQT"
elif n == 0xD:
chunk_name = "DRI"
chunk_size = 2
else:
is_unk_chunk = True
elif chunk_m_h & 0xF0 == 0xC0:
n = chunk_m_h & 0x0F
if n == 0:
chunk_name = "SOF0"
elif n == 2:
chunk_name = "SOF2"
elif n == 4:
chunk_name = "DHT"
else:
is_unk_chunk = True
elif chunk_m_h & 0xF0 == 0xE0:
n = chunk_m_h & 0x0F
chunk_name = "APP%d" % (n)
else:
is_unk_chunk = True
#if is_unk_chunk:
# #raise ValueError(f"Unknown chunk {chunk_m_h} in {filename}")
# io.log_info(f"Unknown chunk {chunk_m_h} in {filename}")
if chunk_size == None: #variable size
chunk_size, = struct.unpack (">H", data[data_counter:data_counter+2])
chunk_size -= 2
data_counter += 2
if chunk_size > 0:
chunk_data = data[data_counter:data_counter+chunk_size]
data_counter += chunk_size
if chunk_name == "SOS":
c = data_counter
while c < inst_length and (data[c] != 0xFF or data[c+1] != 0xD9):
c += 1
chunk_ex_data = data[data_counter:c]
data_counter = c
chunks.append ({'name' : chunk_name,
'm_h' : chunk_m_h,
'data' : chunk_data,
'ex_data' : chunk_ex_data,
})
inst.chunks = chunks
return inst
except Exception as e:
raise Exception (f"Corrupted JPG file {filename} {e}")
@staticmethod
def load(filename, loader_func=None):
try:
inst = DFLJPG.load_raw (filename, loader_func=loader_func)
inst.dfl_dict = {}
for chunk in inst.chunks:
if chunk['name'] == 'APP0':
d, c = chunk['data'], 0
c, id, _ = struct_unpack (d, c, "=4sB")
if id == b"JFIF":
c, ver_major, ver_minor, units, Xdensity, Ydensity, Xthumbnail, Ythumbnail = struct_unpack (d, c, "=BBBHHBB")
else:
raise Exception("Unknown jpeg ID: %s" % (id) )
elif chunk['name'] == 'SOF0' or chunk['name'] == 'SOF2':
d, c = chunk['data'], 0
c, precision, height, width = struct_unpack (d, c, ">BHH")
inst.shape = (height, width, 3)
elif chunk['name'] == 'APP15':
if type(chunk['data']) == bytes:
inst.dfl_dict = pickle.loads(chunk['data'])
return inst
except Exception as e:
io.log_err (f'Exception occured while DFLJPG.load : {traceback.format_exc()}')
return None
def has_data(self):
return len(self.dfl_dict.keys()) != 0
def save(self):
try:
with open(self.filename, "wb") as f:
f.write ( self.dump() )
except:
raise Exception( f'cannot save {self.filename}' )
def dump(self):
data = b""
dict_data = self.dfl_dict
# Remove None keys
for key in list(dict_data.keys()):
if dict_data[key] is None:
dict_data.pop(key)
for chunk in self.chunks:
if chunk['name'] == 'APP15':
self.chunks.remove(chunk)
break
last_app_chunk = 0
for i, chunk in enumerate (self.chunks):
if chunk['m_h'] & 0xF0 == 0xE0:
last_app_chunk = i
dflchunk = {'name' : 'APP15',
'm_h' : 0xEF,
'data' : pickle.dumps(dict_data),
'ex_data' : None,
}
self.chunks.insert (last_app_chunk+1, dflchunk)
for chunk in self.chunks:
data += struct.pack ("BB", 0xFF, chunk['m_h'] )
chunk_data = chunk['data']
if chunk_data is not None:
data += struct.pack (">H", len(chunk_data)+2 )
data += chunk_data
chunk_ex_data = chunk['ex_data']
if chunk_ex_data is not None:
data += chunk_ex_data
return data
def get_img(self):
if self.img is None:
self.img = cv2_imread(self.filename)
return self.img
def get_shape(self):
if self.shape is None:
img = self.get_img()
if img is not None:
self.shape = img.shape
return self.shape
def get_height(self):
for chunk in self.chunks:
if type(chunk) == IHDR:
return chunk.height
return 0
def get_dict(self):
return self.dfl_dict
def set_dict (self, dict_data=None):
self.dfl_dict = dict_data
def get_face_type(self): return self.dfl_dict.get('face_type', FaceType.toString (FaceType.FULL) )
def set_face_type(self, face_type): self.dfl_dict['face_type'] = face_type
def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
def set_landmarks(self, landmarks): self.dfl_dict['landmarks'] = landmarks
def get_eyebrows_expand_mod(self): return self.dfl_dict.get ('eyebrows_expand_mod', 1.0)
def set_eyebrows_expand_mod(self, eyebrows_expand_mod): self.dfl_dict['eyebrows_expand_mod'] = eyebrows_expand_mod
def get_source_filename(self): return self.dfl_dict.get ('source_filename', None)
def set_source_filename(self, source_filename): self.dfl_dict['source_filename'] = source_filename
def get_source_rect(self): return self.dfl_dict.get ('source_rect', None)
def set_source_rect(self, source_rect): self.dfl_dict['source_rect'] = source_rect
def get_source_landmarks(self): return np.array ( self.dfl_dict.get('source_landmarks', None) )
def set_source_landmarks(self, source_landmarks): self.dfl_dict['source_landmarks'] = source_landmarks
def get_image_to_face_mat(self):
mat = self.dfl_dict.get ('image_to_face_mat', None)
if mat is not None:
return np.array (mat)
return None
def set_image_to_face_mat(self, image_to_face_mat): self.dfl_dict['image_to_face_mat'] = image_to_face_mat
def has_seg_ie_polys(self):
return self.dfl_dict.get('seg_ie_polys',None) is not None
def get_seg_ie_polys(self):
d = self.dfl_dict.get('seg_ie_polys',None)
if d is not None:
d = SegIEPolys.load(d)
else:
d = SegIEPolys()
return d
def set_seg_ie_polys(self, seg_ie_polys):
if seg_ie_polys is not None:
if not isinstance(seg_ie_polys, SegIEPolys):
raise ValueError('seg_ie_polys should be instance of SegIEPolys')
if seg_ie_polys.has_polys():
seg_ie_polys = seg_ie_polys.dump()
else:
seg_ie_polys = None
self.dfl_dict['seg_ie_polys'] = seg_ie_polys
def has_xseg_mask(self):
return self.dfl_dict.get('xseg_mask',None) is not None
def get_xseg_mask(self):
mask_buf = self.dfl_dict.get('xseg_mask',None)
if mask_buf is None:
return None
img = cv2.imdecode(mask_buf, cv2.IMREAD_UNCHANGED)
if len(img.shape) == 2:
img = img[...,None]
return img.astype(np.float32) / 255.0
def set_xseg_mask(self, mask_a):
if mask_a is None:
self.dfl_dict['xseg_mask'] = None
return
mask_a = imagelib.normalize_channels(mask_a, 1)
img_data = np.clip( mask_a*255, 0, 255 ).astype(np.uint8)
data_max_len = 4096
ret, buf = cv2.imencode('.png', img_data)
if not ret or len(buf) > data_max_len:
for jpeg_quality in range(100,-1,-1):
ret, buf = cv2.imencode( '.jpg', img_data, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality] )
if ret and len(buf) <= data_max_len:
break
if not ret:
raise Exception("set_xseg_mask: unable to generate image data for set_xseg_mask")
self.dfl_dict['xseg_mask'] = buf
-2
View File
@@ -1,2 +0,0 @@
from .DFLIMG import DFLIMG
from .DFLJPG import DFLJPG
-674
View File
@@ -1,674 +0,0 @@
GNU GENERAL PUBLIC LICENSE
Version 3, 29 June 2007
Copyright (C) 2007 Free Software Foundation, Inc. <http://fsf.org/>
Everyone is permitted to copy and distribute verbatim copies
of this license document, but changing it is not allowed.
Preamble
The GNU General Public License is a free, copyleft license for
software and other kinds of works.
The licenses for most software and other practical works are designed
to take away your freedom to share and change the works. By contrast,
the GNU General Public License is intended to guarantee your freedom to
share and change all versions of a program--to make sure it remains free
software for all its users. We, the Free Software Foundation, use the
GNU General Public License for most of our software; it applies also to
any other work released this way by its authors. You can apply it to
your programs, too.
When we speak of free software, we are referring to freedom, not
price. Our General Public Licenses are designed to make sure that you
have the freedom to distribute copies of free software (and charge for
them if you wish), that you receive source code or can get it if you
want it, that you can change the software or use pieces of it in new
free programs, and that you know you can do these things.
To protect your rights, we need to prevent others from denying you
these rights or asking you to surrender the rights. Therefore, you have
certain responsibilities if you distribute copies of the software, or if
you modify it: responsibilities to respect the freedom of others.
For example, if you distribute copies of such a program, whether
gratis or for a fee, you must pass on to the recipients the same
freedoms that you received. You must make sure that they, too, receive
or can get the source code. And you must show them these terms so they
know their rights.
Developers that use the GNU GPL protect your rights with two steps:
(1) assert copyright on the software, and (2) offer you this License
giving you legal permission to copy, distribute and/or modify it.
For the developers' and authors' protection, the GPL clearly explains
that there is no warranty for this free software. For both users' and
authors' sake, the GPL requires that modified versions be marked as
changed, so that their problems will not be attributed erroneously to
authors of previous versions.
Some devices are designed to deny users access to install or run
modified versions of the software inside them, although the manufacturer
can do so. This is fundamentally incompatible with the aim of
protecting users' freedom to change the software. The systematic
pattern of such abuse occurs in the area of products for individuals to
use, which is precisely where it is most unacceptable. Therefore, we
have designed this version of the GPL to prohibit the practice for those
products. If such problems arise substantially in other domains, we
stand ready to extend this provision to those domains in future versions
of the GPL, as needed to protect the freedom of users.
Finally, every program is threatened constantly by software patents.
States should not allow patents to restrict development and use of
software on general-purpose computers, but in those that do, we wish to
avoid the special danger that patents applied to a free program could
make it effectively proprietary. To prevent this, the GPL assures that
patents cannot be used to render the program non-free.
The precise terms and conditions for copying, distribution and
modification follow.
TERMS AND CONDITIONS
0. Definitions.
"This License" refers to version 3 of the GNU General Public License.
"Copyright" also means copyright-like laws that apply to other kinds of
works, such as semiconductor masks.
"The Program" refers to any copyrightable work licensed under this
License. Each licensee is addressed as "you". "Licensees" and
"recipients" may be individuals or organizations.
To "modify" a work means to copy from or adapt all or part of the work
in a fashion requiring copyright permission, other than the making of an
exact copy. The resulting work is called a "modified version" of the
earlier work or a work "based on" the earlier work.
A "covered work" means either the unmodified Program or a work based
on the Program.
To "propagate" a work means to do anything with it that, without
permission, would make you directly or secondarily liable for
infringement under applicable copyright law, except executing it on a
computer or modifying a private copy. Propagation includes copying,
distribution (with or without modification), making available to the
public, and in some countries other activities as well.
To "convey" a work means any kind of propagation that enables other
parties to make or receive copies. Mere interaction with a user through
a computer network, with no transfer of a copy, is not conveying.
An interactive user interface displays "Appropriate Legal Notices"
to the extent that it includes a convenient and prominently visible
feature that (1) displays an appropriate copyright notice, and (2)
tells the user that there is no warranty for the work (except to the
extent that warranties are provided), that licensees may convey the
work under this License, and how to view a copy of this License. If
the interface presents a list of user commands or options, such as a
menu, a prominent item in the list meets this criterion.
1. Source Code.
The "source code" for a work means the preferred form of the work
for making modifications to it. "Object code" means any non-source
form of a work.
A "Standard Interface" means an interface that either is an official
standard defined by a recognized standards body, or, in the case of
interfaces specified for a particular programming language, one that
is widely used among developers working in that language.
The "System Libraries" of an executable work include anything, other
than the work as a whole, that (a) is included in the normal form of
packaging a Major Component, but which is not part of that Major
Component, and (b) serves only to enable use of the work with that
Major Component, or to implement a Standard Interface for which an
implementation is available to the public in source code form. A
"Major Component", in this context, means a major essential component
(kernel, window system, and so on) of the specific operating system
(if any) on which the executable work runs, or a compiler used to
produce the work, or an object code interpreter used to run it.
The "Corresponding Source" for a work in object code form means all
the source code needed to generate, install, and (for an executable
work) run the object code and to modify the work, including scripts to
control those activities. However, it does not include the work's
System Libraries, or general-purpose tools or generally available free
programs which are used unmodified in performing those activities but
which are not part of the work. For example, Corresponding Source
includes interface definition files associated with source files for
the work, and the source code for shared libraries and dynamically
linked subprograms that the work is specifically designed to require,
such as by intimate data communication or control flow between those
subprograms and other parts of the work.
The Corresponding Source need not include anything that users
can regenerate automatically from other parts of the Corresponding
Source.
The Corresponding Source for a work in source code form is that
same work.
2. Basic Permissions.
All rights granted under this License are granted for the term of
copyright on the Program, and are irrevocable provided the stated
conditions are met. This License explicitly affirms your unlimited
permission to run the unmodified Program. The output from running a
covered work is covered by this License only if the output, given its
content, constitutes a covered work. This License acknowledges your
rights of fair use or other equivalent, as provided by copyright law.
You may make, run and propagate covered works that you do not
convey, without conditions so long as your license otherwise remains
in force. You may convey covered works to others for the sole purpose
of having them make modifications exclusively for you, or provide you
with facilities for running those works, provided that you comply with
the terms of this License in conveying all material for which you do
not control copyright. Those thus making or running the covered works
for you must do so exclusively on your behalf, under your direction
and control, on terms that prohibit them from making any copies of
your copyrighted material outside their relationship with you.
Conveying under any other circumstances is permitted solely under
the conditions stated below. Sublicensing is not allowed; section 10
makes it unnecessary.
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
No covered work shall be deemed part of an effective technological
measure under any applicable law fulfilling obligations under article
11 of the WIPO copyright treaty adopted on 20 December 1996, or
similar laws prohibiting or restricting circumvention of such
measures.
When you convey a covered work, you waive any legal power to forbid
circumvention of technological measures to the extent such circumvention
is effected by exercising rights under this License with respect to
the covered work, and you disclaim any intention to limit operation or
modification of the work as a means of enforcing, against the work's
users, your or third parties' legal rights to forbid circumvention of
technological measures.
4. Conveying Verbatim Copies.
You may convey verbatim copies of the Program's source code as you
receive it, in any medium, provided that you conspicuously and
appropriately publish on each copy an appropriate copyright notice;
keep intact all notices stating that this License and any
non-permissive terms added in accord with section 7 apply to the code;
keep intact all notices of the absence of any warranty; and give all
recipients a copy of this License along with the Program.
You may charge any price or no price for each copy that you convey,
and you may offer support or warranty protection for a fee.
5. Conveying Modified Source Versions.
You may convey a work based on the Program, or the modifications to
produce it from the Program, in the form of source code under the
terms of section 4, provided that you also meet all of these conditions:
a) The work must carry prominent notices stating that you modified
it, and giving a relevant date.
b) The work must carry prominent notices stating that it is
released under this License and any conditions added under section
7. This requirement modifies the requirement in section 4 to
"keep intact all notices".
c) You must license the entire work, as a whole, under this
License to anyone who comes into possession of a copy. This
License will therefore apply, along with any applicable section 7
additional terms, to the whole of the work, and all its parts,
regardless of how they are packaged. This License gives no
permission to license the work in any other way, but it does not
invalidate such permission if you have separately received it.
d) If the work has interactive user interfaces, each must display
Appropriate Legal Notices; however, if the Program has interactive
interfaces that do not display Appropriate Legal Notices, your
work need not make them do so.
A compilation of a covered work with other separate and independent
works, which are not by their nature extensions of the covered work,
and which are not combined with it such as to form a larger program,
in or on a volume of a storage or distribution medium, is called an
"aggregate" if the compilation and its resulting copyright are not
used to limit the access or legal rights of the compilation's users
beyond what the individual works permit. Inclusion of a covered work
in an aggregate does not cause this License to apply to the other
parts of the aggregate.
6. Conveying Non-Source Forms.
You may convey a covered work in object code form under the terms
of sections 4 and 5, provided that you also convey the
machine-readable Corresponding Source under the terms of this License,
in one of these ways:
a) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by the
Corresponding Source fixed on a durable physical medium
customarily used for software interchange.
b) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by a
written offer, valid for at least three years and valid for as
long as you offer spare parts or customer support for that product
model, to give anyone who possesses the object code either (1) a
copy of the Corresponding Source for all the software in the
product that is covered by this License, on a durable physical
medium customarily used for software interchange, for a price no
more than your reasonable cost of physically performing this
conveying of source, or (2) access to copy the
Corresponding Source from a network server at no charge.
c) Convey individual copies of the object code with a copy of the
written offer to provide the Corresponding Source. This
alternative is allowed only occasionally and noncommercially, and
only if you received the object code with such an offer, in accord
with subsection 6b.
d) Convey the object code by offering access from a designated
place (gratis or for a charge), and offer equivalent access to the
Corresponding Source in the same way through the same place at no
further charge. You need not require recipients to copy the
Corresponding Source along with the object code. If the place to
copy the object code is a network server, the Corresponding Source
may be on a different server (operated by you or a third party)
that supports equivalent copying facilities, provided you maintain
clear directions next to the object code saying where to find the
Corresponding Source. Regardless of what server hosts the
Corresponding Source, you remain obligated to ensure that it is
available for as long as needed to satisfy these requirements.
e) Convey the object code using peer-to-peer transmission, provided
you inform other peers where the object code and Corresponding
Source of the work are being offered to the general public at no
charge under subsection 6d.
A separable portion of the object code, whose source code is excluded
from the Corresponding Source as a System Library, need not be
included in conveying the object code work.
A "User Product" is either (1) a "consumer product", which means any
tangible personal property which is normally used for personal, family,
or household purposes, or (2) anything designed or sold for incorporation
into a dwelling. In determining whether a product is a consumer product,
doubtful cases shall be resolved in favor of coverage. For a particular
product received by a particular user, "normally used" refers to a
typical or common use of that class of product, regardless of the status
of the particular user or of the way in which the particular user
actually uses, or expects or is expected to use, the product. A product
is a consumer product regardless of whether the product has substantial
commercial, industrial or non-consumer uses, unless such uses represent
the only significant mode of use of the product.
"Installation Information" for a User Product means any methods,
procedures, authorization keys, or other information required to install
and execute modified versions of a covered work in that User Product from
a modified version of its Corresponding Source. The information must
suffice to ensure that the continued functioning of the modified object
code is in no case prevented or interfered with solely because
modification has been made.
If you convey an object code work under this section in, or with, or
specifically for use in, a User Product, and the conveying occurs as
part of a transaction in which the right of possession and use of the
User Product is transferred to the recipient in perpetuity or for a
fixed term (regardless of how the transaction is characterized), the
Corresponding Source conveyed under this section must be accompanied
by the Installation Information. But this requirement does not apply
if neither you nor any third party retains the ability to install
modified object code on the User Product (for example, the work has
been installed in ROM).
The requirement to provide Installation Information does not include a
requirement to continue to provide support service, warranty, or updates
for a work that has been modified or installed by the recipient, or for
the User Product in which it has been modified or installed. Access to a
network may be denied when the modification itself materially and
adversely affects the operation of the network or violates the rules and
protocols for communication across the network.
Corresponding Source conveyed, and Installation Information provided,
in accord with this section must be in a format that is publicly
documented (and with an implementation available to the public in
source code form), and must require no special password or key for
unpacking, reading or copying.
7. Additional Terms.
"Additional permissions" are terms that supplement the terms of this
License by making exceptions from one or more of its conditions.
Additional permissions that are applicable to the entire Program shall
be treated as though they were included in this License, to the extent
that they are valid under applicable law. If additional permissions
apply only to part of the Program, that part may be used separately
under those permissions, but the entire Program remains governed by
this License without regard to the additional permissions.
When you convey a copy of a covered work, you may at your option
remove any additional permissions from that copy, or from any part of
it. (Additional permissions may be written to require their own
removal in certain cases when you modify the work.) You may place
additional permissions on material, added by you to a covered work,
for which you have or can give appropriate copyright permission.
Notwithstanding any other provision of this License, for material you
add to a covered work, you may (if authorized by the copyright holders of
that material) supplement the terms of this License with terms:
a) Disclaiming warranty or limiting liability differently from the
terms of sections 15 and 16 of this License; or
b) Requiring preservation of specified reasonable legal notices or
author attributions in that material or in the Appropriate Legal
Notices displayed by works containing it; or
c) Prohibiting misrepresentation of the origin of that material, or
requiring that modified versions of such material be marked in
reasonable ways as different from the original version; or
d) Limiting the use for publicity purposes of names of licensors or
authors of the material; or
e) Declining to grant rights under trademark law for use of some
trade names, trademarks, or service marks; or
f) Requiring indemnification of licensors and authors of that
material by anyone who conveys the material (or modified versions of
it) with contractual assumptions of liability to the recipient, for
any liability that these contractual assumptions directly impose on
those licensors and authors.
All other non-permissive additional terms are considered "further
restrictions" within the meaning of section 10. If the Program as you
received it, or any part of it, contains a notice stating that it is
governed by this License along with a term that is a further
restriction, you may remove that term. If a license document contains
a further restriction but permits relicensing or conveying under this
License, you may add to a covered work material governed by the terms
of that license document, provided that the further restriction does
not survive such relicensing or conveying.
If you add terms to a covered work in accord with this section, you
must place, in the relevant source files, a statement of the
additional terms that apply to those files, or a notice indicating
where to find the applicable terms.
Additional terms, permissive or non-permissive, may be stated in the
form of a separately written license, or stated as exceptions;
the above requirements apply either way.
8. Termination.
You may not propagate or modify a covered work except as expressly
provided under this License. Any attempt otherwise to propagate or
modify it is void, and will automatically terminate your rights under
this License (including any patent licenses granted under the third
paragraph of section 11).
However, if you cease all violation of this License, then your
license from a particular copyright holder is reinstated (a)
provisionally, unless and until the copyright holder explicitly and
finally terminates your license, and (b) permanently, if the copyright
holder fails to notify you of the violation by some reasonable means
prior to 60 days after the cessation.
Moreover, your license from a particular copyright holder is
reinstated permanently if the copyright holder notifies you of the
violation by some reasonable means, this is the first time you have
received notice of violation of this License (for any work) from that
copyright holder, and you cure the violation prior to 30 days after
your receipt of the notice.
Termination of your rights under this section does not terminate the
licenses of parties who have received copies or rights from you under
this License. If your rights have been terminated and not permanently
reinstated, you do not qualify to receive new licenses for the same
material under section 10.
9. Acceptance Not Required for Having Copies.
You are not required to accept this License in order to receive or
run a copy of the Program. Ancillary propagation of a covered work
occurring solely as a consequence of using peer-to-peer transmission
to receive a copy likewise does not require acceptance. However,
nothing other than this License grants you permission to propagate or
modify any covered work. These actions infringe copyright if you do
not accept this License. Therefore, by modifying or propagating a
covered work, you indicate your acceptance of this License to do so.
10. Automatic Licensing of Downstream Recipients.
Each time you convey a covered work, the recipient automatically
receives a license from the original licensors, to run, modify and
propagate that work, subject to this License. You are not responsible
for enforcing compliance by third parties with this License.
An "entity transaction" is a transaction transferring control of an
organization, or substantially all assets of one, or subdividing an
organization, or merging organizations. If propagation of a covered
work results from an entity transaction, each party to that
transaction who receives a copy of the work also receives whatever
licenses to the work the party's predecessor in interest had or could
give under the previous paragraph, plus a right to possession of the
Corresponding Source of the work from the predecessor in interest, if
the predecessor has it or can get it with reasonable efforts.
You may not impose any further restrictions on the exercise of the
rights granted or affirmed under this License. For example, you may
not impose a license fee, royalty, or other charge for exercise of
rights granted under this License, and you may not initiate litigation
(including a cross-claim or counterclaim in a lawsuit) alleging that
any patent claim is infringed by making, using, selling, offering for
sale, or importing the Program or any portion of it.
11. Patents.
A "contributor" is a copyright holder who authorizes use under this
License of the Program or a work on which the Program is based. The
work thus licensed is called the contributor's "contributor version".
A contributor's "essential patent claims" are all patent claims
owned or controlled by the contributor, whether already acquired or
hereafter acquired, that would be infringed by some manner, permitted
by this License, of making, using, or selling its contributor version,
but do not include claims that would be infringed only as a
consequence of further modification of the contributor version. For
purposes of this definition, "control" includes the right to grant
patent sublicenses in a manner consistent with the requirements of
this License.
Each contributor grants you a non-exclusive, worldwide, royalty-free
patent license under the contributor's essential patent claims, to
make, use, sell, offer for sale, import and otherwise run, modify and
propagate the contents of its contributor version.
In the following three paragraphs, a "patent license" is any express
agreement or commitment, however denominated, not to enforce a patent
(such as an express permission to practice a patent or covenant not to
sue for patent infringement). To "grant" such a patent license to a
party means to make such an agreement or commitment not to enforce a
patent against the party.
If you convey a covered work, knowingly relying on a patent license,
and the Corresponding Source of the work is not available for anyone
to copy, free of charge and under the terms of this License, through a
publicly available network server or other readily accessible means,
then you must either (1) cause the Corresponding Source to be so
available, or (2) arrange to deprive yourself of the benefit of the
patent license for this particular work, or (3) arrange, in a manner
consistent with the requirements of this License, to extend the patent
license to downstream recipients. "Knowingly relying" means you have
actual knowledge that, but for the patent license, your conveying the
covered work in a country, or your recipient's use of the covered work
in a country, would infringe one or more identifiable patents in that
country that you have reason to believe are valid.
If, pursuant to or in connection with a single transaction or
arrangement, you convey, or propagate by procuring conveyance of, a
covered work, and grant a patent license to some of the parties
receiving the covered work authorizing them to use, propagate, modify
or convey a specific copy of the covered work, then the patent license
you grant is automatically extended to all recipients of the covered
work and works based on it.
A patent license is "discriminatory" if it does not include within
the scope of its coverage, prohibits the exercise of, or is
conditioned on the non-exercise of one or more of the rights that are
specifically granted under this License. You may not convey a covered
work if you are a party to an arrangement with a third party that is
in the business of distributing software, under which you make payment
to the third party based on the extent of your activity of conveying
the work, and under which the third party grants, to any of the
parties who would receive the covered work from you, a discriminatory
patent license (a) in connection with copies of the covered work
conveyed by you (or copies made from those copies), or (b) primarily
for and in connection with specific products or compilations that
contain the covered work, unless you entered into that arrangement,
or that patent license was granted, prior to 28 March 2007.
Nothing in this License shall be construed as excluding or limiting
any implied license or other defenses to infringement that may
otherwise be available to you under applicable patent law.
12. No Surrender of Others' Freedom.
If conditions are imposed on you (whether by court order, agreement or
otherwise) that contradict the conditions of this License, they do not
excuse you from the conditions of this License. If you cannot convey a
covered work so as to satisfy simultaneously your obligations under this
License and any other pertinent obligations, then as a consequence you may
not convey it at all. For example, if you agree to terms that obligate you
to collect a royalty for further conveying from those to whom you convey
the Program, the only way you could satisfy both those terms and this
License would be to refrain entirely from conveying the Program.
13. Use with the GNU Affero General Public License.
Notwithstanding any other provision of this License, you have
permission to link or combine any covered work with a work licensed
under version 3 of the GNU Affero General Public License into a single
combined work, and to convey the resulting work. The terms of this
License will continue to apply to the part which is the covered work,
but the special requirements of the GNU Affero General Public License,
section 13, concerning interaction through a network will apply to the
combination as such.
14. Revised Versions of this License.
The Free Software Foundation may publish revised and/or new versions of
the GNU General Public License from time to time. Such new versions will
be similar in spirit to the present version, but may differ in detail to
address new problems or concerns.
Each version is given a distinguishing version number. If the
Program specifies that a certain numbered version of the GNU General
Public License "or any later version" applies to it, you have the
option of following the terms and conditions either of that numbered
version or of any later version published by the Free Software
Foundation. If the Program does not specify a version number of the
GNU General Public License, you may choose any version ever published
by the Free Software Foundation.
If the Program specifies that a proxy can decide which future
versions of the GNU General Public License can be used, that proxy's
public statement of acceptance of a version permanently authorizes you
to choose that version for the Program.
Later license versions may give you additional or different
permissions. However, no additional obligations are imposed on any
author or copyright holder as a result of your choosing to follow a
later version.
15. Disclaimer of Warranty.
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
16. Limitation of Liability.
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
SUCH DAMAGES.
17. Interpretation of Sections 15 and 16.
If the disclaimer of warranty and limitation of liability provided
above cannot be given local legal effect according to their terms,
reviewing courts shall apply local law that most closely approximates
an absolute waiver of all civil liability in connection with the
Program, unless a warranty or assumption of liability accompanies a
copy of the Program in return for a fee.
END OF TERMS AND CONDITIONS
How to Apply These Terms to Your New Programs
If you develop a new program, and you want it to be of the greatest
possible use to the public, the best way to achieve this is to make it
free software which everyone can redistribute and change under these terms.
To do so, attach the following notices to the program. It is safest
to attach them to the start of each source file to most effectively
state the exclusion of warranty; and each file should have at least
the "copyright" line and a pointer to where the full notice is found.
<one line to give the program's name and a brief idea of what it does.>
Copyright (C) <year> <name of author>
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 <http://www.gnu.org/licenses/>.
Also add information on how to contact you by electronic and paper mail.
If the program does terminal interaction, make it output a short
notice like this when it starts in an interactive mode:
<program> Copyright (C) <year> <name of author>
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
This is free software, and you are welcome to redistribute it
under certain conditions; type `show c' for details.
The hypothetical commands `show w' and `show c' should show the appropriate
parts of the General Public License. Of course, your program's commands
might be different; for a GUI interface, you would use an "about box".
You should also get your employer (if you work as a programmer) or school,
if any, to sign a "copyright disclaimer" for the program, if necessary.
For more information on this, and how to apply and follow the GNU GPL, see
<http://www.gnu.org/licenses/>.
The GNU General Public License does not permit incorporating your program
into proprietary programs. If your program is a subroutine library, you
may consider it more useful to permit linking proprietary applications with
the library. If this is what you want to do, use the GNU Lesser General
Public License instead of this License. But first, please read
<http://www.gnu.org/philosophy/why-not-lgpl.html>.
+1 -1
View File
@@ -1,6 +1,6 @@
<table align="center" border="0"><tr><td align="center" width="9999">
# DeepFaceLab
# DeepFaceLab Linux
<a href="https://arxiv.org/abs/2005.05535">
-10
View File
@@ -1,10 +0,0 @@
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
class QCursorDB():
@staticmethod
def initialize(cursor_path):
QCursorDB.cross_red = QCursor ( QPixmap ( str(cursor_path / 'cross_red.png') ) )
QCursorDB.cross_green = QCursor ( QPixmap ( str(cursor_path / 'cross_green.png') ) )
QCursorDB.cross_blue = QCursor ( QPixmap ( str(cursor_path / 'cross_blue.png') ) )
-25
View File
@@ -1,25 +0,0 @@
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
class QIconDB():
@staticmethod
def initialize(icon_path):
QIconDB.app_icon = QIcon ( str(icon_path / 'app_icon.png') )
QIconDB.delete_poly = QIcon ( str(icon_path / 'delete_poly.png') )
QIconDB.undo_pt = QIcon ( str(icon_path / 'undo_pt.png') )
QIconDB.redo_pt = QIcon ( str(icon_path / 'redo_pt.png') )
QIconDB.poly_color_red = QIcon ( str(icon_path / 'poly_color_red.png') )
QIconDB.poly_color_green = QIcon ( str(icon_path / 'poly_color_green.png') )
QIconDB.poly_color_blue = QIcon ( str(icon_path / 'poly_color_blue.png') )
QIconDB.poly_type_include = QIcon ( str(icon_path / 'poly_type_include.png') )
QIconDB.poly_type_exclude = QIcon ( str(icon_path / 'poly_type_exclude.png') )
QIconDB.left = QIcon ( str(icon_path / 'left.png') )
QIconDB.right = QIcon ( str(icon_path / 'right.png') )
QIconDB.pt_edit_mode = QIcon ( str(icon_path / 'pt_edit_mode.png') )
QIconDB.view_lock_center = QIcon ( str(icon_path / 'view_lock_center.png') )
QIconDB.view_baked = QIcon ( str(icon_path / 'view_baked.png') )
QIconDB.view_xseg = QIcon ( str(icon_path / 'view_xseg.png') )
QIconDB.view_xseg_overlay = QIcon ( str(icon_path / 'view_xseg_overlay.png') )
-8
View File
@@ -1,8 +0,0 @@
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
class QImageDB():
@staticmethod
def initialize(image_path):
QImageDB.intro = QImage ( str(image_path / 'intro.png') )
-97
View File
@@ -1,97 +0,0 @@
from localization import system_language
class QStringDB():
@staticmethod
def initialize():
lang = system_language
if lang not in ['en','ru','zh']:
lang = 'en'
QStringDB.btn_poly_color_red_tip = { 'en' : 'Poly color scheme red',
'ru' : 'Красная цветовая схема полигонов',
'zh' : '选区配色方案红色',
}[lang]
QStringDB.btn_poly_color_green_tip = { 'en' : 'Poly color scheme green',
'ru' : 'Зелёная цветовая схема полигонов',
'zh' : '选区配色方案绿色',
}[lang]
QStringDB.btn_poly_color_blue_tip = { 'en' : 'Poly color scheme blue',
'ru' : 'Синяя цветовая схема полигонов',
'zh' : '选区配色方案蓝色',
}[lang]
QStringDB.btn_view_baked_mask_tip = { 'en' : 'View baked mask',
'ru' : 'Посмотреть запечёную маску',
'zh' : '查看遮罩通道',
}[lang]
QStringDB.btn_view_xseg_mask_tip = { 'en' : 'View trained XSeg mask',
'ru' : 'Посмотреть тренированную XSeg маску',
'zh' : '查看导入后的XSeg遮罩',
}[lang]
QStringDB.btn_view_xseg_overlay_mask_tip = { 'en' : 'View trained XSeg mask overlay face',
'ru' : 'Посмотреть тренированную XSeg маску поверх лица',
'zh' : '查看导入后的XSeg遮罩于脸上方',
}[lang]
QStringDB.btn_poly_type_include_tip = { 'en' : 'Poly include mode',
'ru' : 'Режим полигонов - включение',
'zh' : '包含选区模式',
}[lang]
QStringDB.btn_poly_type_exclude_tip = { 'en' : 'Poly exclude mode',
'ru' : 'Режим полигонов - исключение',
'zh' : '排除选区模式',
}[lang]
QStringDB.btn_undo_pt_tip = { 'en' : 'Undo point',
'ru' : 'Отменить точку',
'zh' : '撤消点',
}[lang]
QStringDB.btn_redo_pt_tip = { 'en' : 'Redo point',
'ru' : 'Повторить точку',
'zh' : '重做点',
}[lang]
QStringDB.btn_delete_poly_tip = { 'en' : 'Delete poly',
'ru' : 'Удалить полигон',
'zh' : '删除选区',
}[lang]
QStringDB.btn_pt_edit_mode_tip = { 'en' : 'Add/delete point mode ( HOLD CTRL )',
'ru' : 'Режим добавления/удаления точек ( удерживайте CTRL )',
'zh' : '点加/删除模式 ( 按住CTRL )',
}[lang]
QStringDB.btn_view_lock_center_tip = { 'en' : 'Lock cursor at the center ( HOLD SHIFT )',
'ru' : 'Заблокировать курсор в центре ( удерживайте SHIFT )',
'zh' : '将光标锁定在中心 ( 按住SHIFT )',
}[lang]
QStringDB.btn_prev_image_tip = { 'en' : 'Save and Prev image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
'ru' : 'Сохранить и предыдущее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
'zh' : '保存并转到上一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
}[lang]
QStringDB.btn_next_image_tip = { 'en' : 'Save and Next image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
'ru' : 'Сохранить и следующее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
'zh' : '保存并转到下一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
}[lang]
QStringDB.loading_tip = {'en' : 'Loading',
'ru' : 'Загрузка',
'zh' : '正在载入',
}[lang]
QStringDB.labeled_tip = {'en' : 'labeled',
'ru' : 'размечено',
'zh' : '标记的',
}[lang]
File diff suppressed because it is too large Load Diff
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.6 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.6 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.6 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 5.5 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 4.7 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.6 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.7 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 8.4 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 9.0 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 8.9 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 6.3 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 5.6 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 4.2 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 5.4 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.7 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 5.4 KiB

Binary file not shown.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 2.6 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 8.1 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 4.0 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 9.5 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 12 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 30 KiB

Binary file not shown.
-9
View File
@@ -1,9 +0,0 @@
theme: jekyll-theme-cayman
plugins:
- jekyll-relative-links
relative_links:
enabled: true
collections: true
include:
- README.md
-40
View File
@@ -1,40 +0,0 @@
import cv2
import numpy as np
from pathlib import Path
from core.interact import interact as io
from core import imagelib
import traceback
def cv2_imread(filename, flags=cv2.IMREAD_UNCHANGED, loader_func=None, verbose=True):
"""
allows to open non-english characters path
"""
try:
if loader_func is not None:
bytes = bytearray(loader_func(filename))
else:
with open(filename, "rb") as stream:
bytes = bytearray(stream.read())
numpyarray = np.asarray(bytes, dtype=np.uint8)
return cv2.imdecode(numpyarray, flags)
except:
if verbose:
io.log_err(f"Exception occured in cv2_imread : {traceback.format_exc()}")
return None
def cv2_imwrite(filename, img, *args):
ret, buf = cv2.imencode( Path(filename).suffix, img, *args)
if ret == True:
try:
with open(filename, "wb") as stream:
stream.write( buf )
except:
pass
def cv2_resize(x, *args, **kwargs):
h,w,c = x.shape
x = cv2.resize(x, *args, **kwargs)
x = imagelib.normalize_channels(x, c)
return x
-152
View File
@@ -1,152 +0,0 @@
import numpy as np
import cv2
from enum import IntEnum
class SegIEPolyType(IntEnum):
EXCLUDE = 0
INCLUDE = 1
class SegIEPoly():
def __init__(self, type=None, pts=None, **kwargs):
self.type = type
if pts is None:
pts = np.empty( (0,2), dtype=np.float32 )
else:
pts = np.float32(pts)
self.pts = pts
self.n_max = self.n = len(pts)
def dump(self):
return {'type': int(self.type),
'pts' : self.get_pts(),
}
def identical(self, b):
if self.n != b.n:
return False
return (self.pts[0:self.n] == b.pts[0:b.n]).all()
def get_type(self):
return self.type
def add_pt(self, x, y):
self.pts = np.append(self.pts[0:self.n], [ ( float(x), float(y) ) ], axis=0).astype(np.float32)
self.n_max = self.n = self.n + 1
def undo(self):
self.n = max(0, self.n-1)
return self.n
def redo(self):
self.n = min(len(self.pts), self.n+1)
return self.n
def redo_clip(self):
self.pts = self.pts[0:self.n]
self.n_max = self.n
def insert_pt(self, n, pt):
if n < 0 or n > self.n:
raise ValueError("insert_pt out of range")
self.pts = np.concatenate( (self.pts[0:n], pt[None,...].astype(np.float32), self.pts[n:]), axis=0)
self.n_max = self.n = self.n+1
def remove_pt(self, n):
if n < 0 or n >= self.n:
raise ValueError("remove_pt out of range")
self.pts = np.concatenate( (self.pts[0:n], self.pts[n+1:]), axis=0)
self.n_max = self.n = self.n-1
def get_last_point(self):
return self.pts[self.n-1].copy()
def get_pts(self):
return self.pts[0:self.n].copy()
def get_pts_count(self):
return self.n
def set_point(self, id, pt):
self.pts[id] = pt
def set_points(self, pts):
self.pts = np.array(pts)
self.n_max = self.n = len(pts)
class SegIEPolys():
def __init__(self):
self.polys = []
def identical(self, b):
polys_len = len(self.polys)
o_polys_len = len(b.polys)
if polys_len != o_polys_len:
return False
return all ([ a_poly.identical(b_poly) for a_poly, b_poly in zip(self.polys, b.polys) ])
def add_poly(self, ie_poly_type):
poly = SegIEPoly(ie_poly_type)
self.polys.append (poly)
return poly
def remove_poly(self, poly):
if poly in self.polys:
self.polys.remove(poly)
def has_polys(self):
return len(self.polys) != 0
def get_poly(self, id):
return self.polys[id]
def get_polys(self):
return self.polys
def get_pts_count(self):
return sum([poly.get_pts_count() for poly in self.polys])
def sort(self):
poly_by_type = { SegIEPolyType.EXCLUDE : [], SegIEPolyType.INCLUDE : [] }
for poly in self.polys:
poly_by_type[poly.type].append(poly)
self.polys = poly_by_type[SegIEPolyType.INCLUDE] + poly_by_type[SegIEPolyType.EXCLUDE]
def __iter__(self):
for poly in self.polys:
yield poly
def overlay_mask(self, mask):
h,w,c = mask.shape
white = (1,)*c
black = (0,)*c
for poly in self.polys:
pts = poly.get_pts().astype(np.int32)
if len(pts) != 0:
cv2.fillPoly(mask, [pts], white if poly.type == SegIEPolyType.INCLUDE else black )
def dump(self):
return {'polys' : [ poly.dump() for poly in self.polys ] }
@staticmethod
def load(data=None):
ie_polys = SegIEPolys()
if data is not None:
if isinstance(data, list):
# Backward comp
ie_polys.polys = [ SegIEPoly(type=type, pts=pts) for (type, pts) in data ]
elif isinstance(data, dict):
ie_polys.polys = [ SegIEPoly(**poly_cfg) for poly_cfg in data['polys'] ]
ie_polys.sort()
return ie_polys
-27
View File
@@ -1,27 +0,0 @@
from .estimate_sharpness import estimate_sharpness
from .equalize_and_stack_square import equalize_and_stack_square
from .text import get_text_image, get_draw_text_lines
from .draw import draw_polygon, draw_rect
from .morph import morph_by_points
from .warp import gen_warp_params, warp_by_params
from .reduce_colors import reduce_colors
from .color_transfer import color_transfer, color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer
from .common import normalize_channels, cut_odd_image, overlay_alpha_image
from .SegIEPolys import *
from .blursharpen import LinearMotionBlur, blursharpen
from .filters import apply_random_rgb_levels, \
apply_random_hsv_shift, \
apply_random_motion_blur, \
apply_random_gaussian_blur, \
apply_random_bilinear_resize
-38
View File
@@ -1,38 +0,0 @@
import cv2
import numpy as np
def LinearMotionBlur(image, size, angle):
k = np.zeros((size, size), dtype=np.float32)
k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )
k = k * ( 1.0 / np.sum(k) )
return cv2.filter2D(image, -1, k)
def blursharpen (img, sharpen_mode=0, kernel_size=3, amount=100):
if kernel_size % 2 == 0:
kernel_size += 1
if amount > 0:
if sharpen_mode == 1: #box
kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
kernel[ kernel_size//2, kernel_size//2] = 1.0
box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2)
kernel = kernel + (kernel - box_filter) * amount
return cv2.filter2D(img, -1, kernel)
elif sharpen_mode == 2: #gaussian
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
return img
elif amount < 0:
n = -amount
while n > 0:
img_blur = cv2.medianBlur(img, 5)
if int(n / 10) != 0:
img = img_blur
else:
pass_power = (n % 10) / 10.0
img = img*(1.0-pass_power)+img_blur*pass_power
n = max(n-10,0)
return img
return img
-368
View File
@@ -1,368 +0,0 @@
import cv2
import numpy as np
from numpy import linalg as npla
import scipy as sp
def color_transfer_sot(src,trg, steps=10, batch_size=5, reg_sigmaXY=16.0, reg_sigmaV=5.0):
"""
Color Transform via Sliced Optimal Transfer
ported by @iperov from https://github.com/dcoeurjo/OTColorTransfer
src - any float range any channel image
dst - any float range any channel image, same shape as src
steps - number of solver steps
batch_size - solver batch size
reg_sigmaXY - apply regularization and sigmaXY of filter, otherwise set to 0.0
reg_sigmaV - sigmaV of filter
return value - clip it manually
"""
if not np.issubdtype(src.dtype, np.floating):
raise ValueError("src value must be float")
if not np.issubdtype(trg.dtype, np.floating):
raise ValueError("trg value must be float")
if len(src.shape) != 3:
raise ValueError("src shape must have rank 3 (h,w,c)")
if src.shape != trg.shape:
raise ValueError("src and trg shapes must be equal")
src_dtype = src.dtype
h,w,c = src.shape
new_src = src.copy()
advect = np.empty ( (h*w,c), dtype=src_dtype )
for step in range (steps):
advect.fill(0)
for batch in range (batch_size):
dir = np.random.normal(size=c).astype(src_dtype)
dir /= npla.norm(dir)
projsource = np.sum( new_src*dir, axis=-1).reshape ((h*w))
projtarget = np.sum( trg*dir, axis=-1).reshape ((h*w))
idSource = np.argsort (projsource)
idTarget = np.argsort (projtarget)
a = projtarget[idTarget]-projsource[idSource]
for i_c in range(c):
advect[idSource,i_c] += a * dir[i_c]
new_src += advect.reshape( (h,w,c) ) / batch_size
if reg_sigmaXY != 0.0:
src_diff = new_src-src
src_diff_filt = cv2.bilateralFilter (src_diff, 0, reg_sigmaV, reg_sigmaXY )
if len(src_diff_filt.shape) == 2:
src_diff_filt = src_diff_filt[...,None]
new_src = src + src_diff_filt
return new_src
def color_transfer_mkl(x0, x1):
eps = np.finfo(float).eps
h,w,c = x0.shape
h1,w1,c1 = x1.shape
x0 = x0.reshape ( (h*w,c) )
x1 = x1.reshape ( (h1*w1,c1) )
a = np.cov(x0.T)
b = np.cov(x1.T)
Da2, Ua = np.linalg.eig(a)
Da = np.diag(np.sqrt(Da2.clip(eps, None)))
C = np.dot(np.dot(np.dot(np.dot(Da, Ua.T), b), Ua), Da)
Dc2, Uc = np.linalg.eig(C)
Dc = np.diag(np.sqrt(Dc2.clip(eps, None)))
Da_inv = np.diag(1./(np.diag(Da)))
t = np.dot(np.dot(np.dot(np.dot(np.dot(np.dot(Ua, Da_inv), Uc), Dc), Uc.T), Da_inv), Ua.T)
mx0 = np.mean(x0, axis=0)
mx1 = np.mean(x1, axis=0)
result = np.dot(x0-mx0, t) + mx1
return np.clip ( result.reshape ( (h,w,c) ).astype(x0.dtype), 0, 1)
def color_transfer_idt(i0, i1, bins=256, n_rot=20):
import scipy.stats
relaxation = 1 / n_rot
h,w,c = i0.shape
h1,w1,c1 = i1.shape
i0 = i0.reshape ( (h*w,c) )
i1 = i1.reshape ( (h1*w1,c1) )
n_dims = c
d0 = i0.T
d1 = i1.T
for i in range(n_rot):
r = sp.stats.special_ortho_group.rvs(n_dims).astype(np.float32)
d0r = np.dot(r, d0)
d1r = np.dot(r, d1)
d_r = np.empty_like(d0)
for j in range(n_dims):
lo = min(d0r[j].min(), d1r[j].min())
hi = max(d0r[j].max(), d1r[j].max())
p0r, edges = np.histogram(d0r[j], bins=bins, range=[lo, hi])
p1r, _ = np.histogram(d1r[j], bins=bins, range=[lo, hi])
cp0r = p0r.cumsum().astype(np.float32)
cp0r /= cp0r[-1]
cp1r = p1r.cumsum().astype(np.float32)
cp1r /= cp1r[-1]
f = np.interp(cp0r, cp1r, edges[1:])
d_r[j] = np.interp(d0r[j], edges[1:], f, left=0, right=bins)
d0 = relaxation * np.linalg.solve(r, (d_r - d0r)) + d0
return np.clip ( d0.T.reshape ( (h,w,c) ).astype(i0.dtype) , 0, 1)
def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, source_mask=None, target_mask=None):
"""
Transfers the color distribution from the source to the target
image using the mean and standard deviations of the L*a*b*
color space.
This implementation is (loosely) based on to the "Color Transfer
between Images" paper by Reinhard et al., 2001.
Parameters:
-------
source: NumPy array
OpenCV image in BGR color space (the source image)
target: NumPy array
OpenCV image in BGR color space (the target image)
clip: Should components of L*a*b* image be scaled by np.clip before
converting back to BGR color space?
If False then components will be min-max scaled appropriately.
Clipping will keep target image brightness truer to the input.
Scaling will adjust image brightness to avoid washed out portions
in the resulting color transfer that can be caused by clipping.
preserve_paper: Should color transfer strictly follow methodology
layed out in original paper? The method does not always produce
aesthetically pleasing results.
If False then L*a*b* components will scaled using the reciprocal of
the scaling factor proposed in the paper. This method seems to produce
more consistently aesthetically pleasing results
Returns:
-------
transfer: NumPy array
OpenCV image (w, h, 3) NumPy array (uint8)
"""
# convert the images from the RGB to L*ab* color space, being
# sure to utilizing the floating point data type (note: OpenCV
# expects floats to be 32-bit, so use that instead of 64-bit)
source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32)
target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32)
# compute color statistics for the source and target images
src_input = source if source_mask is None else source*source_mask
tgt_input = target if target_mask is None else target*target_mask
(lMeanSrc, lStdSrc, aMeanSrc, aStdSrc, bMeanSrc, bStdSrc) = lab_image_stats(src_input)
(lMeanTar, lStdTar, aMeanTar, aStdTar, bMeanTar, bStdTar) = lab_image_stats(tgt_input)
# subtract the means from the target image
(l, a, b) = cv2.split(target)
l -= lMeanTar
a -= aMeanTar
b -= bMeanTar
if preserve_paper:
# scale by the standard deviations using paper proposed factor
l = (lStdTar / lStdSrc) * l
a = (aStdTar / aStdSrc) * a
b = (bStdTar / bStdSrc) * b
else:
# scale by the standard deviations using reciprocal of paper proposed factor
l = (lStdSrc / lStdTar) * l
a = (aStdSrc / aStdTar) * a
b = (bStdSrc / bStdTar) * b
# add in the source mean
l += lMeanSrc
a += aMeanSrc
b += bMeanSrc
# clip/scale the pixel intensities to [0, 255] if they fall
# outside this range
l = _scale_array(l, clip=clip)
a = _scale_array(a, clip=clip)
b = _scale_array(b, clip=clip)
# merge the channels together and convert back to the RGB color
# space, being sure to utilize the 8-bit unsigned integer data
# type
transfer = cv2.merge([l, a, b])
transfer = cv2.cvtColor(transfer.astype(np.uint8), cv2.COLOR_LAB2BGR)
# return the color transferred image
return transfer
def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5):
'''
Matches the colour distribution of the target image to that of the source image
using a linear transform.
Images are expected to be of form (w,h,c) and float in [0,1].
Modes are chol, pca or sym for different choices of basis.
'''
mu_t = target_img.mean(0).mean(0)
t = target_img - mu_t
t = t.transpose(2,0,1).reshape( t.shape[-1],-1)
Ct = t.dot(t.T) / t.shape[1] + eps * np.eye(t.shape[0])
mu_s = source_img.mean(0).mean(0)
s = source_img - mu_s
s = s.transpose(2,0,1).reshape( s.shape[-1],-1)
Cs = s.dot(s.T) / s.shape[1] + eps * np.eye(s.shape[0])
if mode == 'chol':
chol_t = np.linalg.cholesky(Ct)
chol_s = np.linalg.cholesky(Cs)
ts = chol_s.dot(np.linalg.inv(chol_t)).dot(t)
if mode == 'pca':
eva_t, eve_t = np.linalg.eigh(Ct)
Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
eva_s, eve_s = np.linalg.eigh(Cs)
Qs = eve_s.dot(np.sqrt(np.diag(eva_s))).dot(eve_s.T)
ts = Qs.dot(np.linalg.inv(Qt)).dot(t)
if mode == 'sym':
eva_t, eve_t = np.linalg.eigh(Ct)
Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
Qt_Cs_Qt = Qt.dot(Cs).dot(Qt)
eva_QtCsQt, eve_QtCsQt = np.linalg.eigh(Qt_Cs_Qt)
QtCsQt = eve_QtCsQt.dot(np.sqrt(np.diag(eva_QtCsQt))).dot(eve_QtCsQt.T)
ts = np.linalg.inv(Qt).dot(QtCsQt).dot(np.linalg.inv(Qt)).dot(t)
matched_img = ts.reshape(*target_img.transpose(2,0,1).shape).transpose(1,2,0)
matched_img += mu_s
matched_img[matched_img>1] = 1
matched_img[matched_img<0] = 0
return np.clip(matched_img.astype(source_img.dtype), 0, 1)
def lab_image_stats(image):
# compute the mean and standard deviation of each channel
(l, a, b) = cv2.split(image)
(lMean, lStd) = (l.mean(), l.std())
(aMean, aStd) = (a.mean(), a.std())
(bMean, bStd) = (b.mean(), b.std())
# return the color statistics
return (lMean, lStd, aMean, aStd, bMean, bStd)
def _scale_array(arr, clip=True):
if clip:
return np.clip(arr, 0, 255)
mn = arr.min()
mx = arr.max()
scale_range = (max([mn, 0]), min([mx, 255]))
if mn < scale_range[0] or mx > scale_range[1]:
return (scale_range[1] - scale_range[0]) * (arr - mn) / (mx - mn) + scale_range[0]
return arr
def channel_hist_match(source, template, hist_match_threshold=255, mask=None):
# Code borrowed from:
# https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x
masked_source = source
masked_template = template
if mask is not None:
masked_source = source * mask
masked_template = template * mask
oldshape = source.shape
source = source.ravel()
template = template.ravel()
masked_source = masked_source.ravel()
masked_template = masked_template.ravel()
s_values, bin_idx, s_counts = np.unique(source, return_inverse=True,
return_counts=True)
t_values, t_counts = np.unique(template, return_counts=True)
s_quantiles = np.cumsum(s_counts).astype(np.float64)
s_quantiles = hist_match_threshold * s_quantiles / s_quantiles[-1]
t_quantiles = np.cumsum(t_counts).astype(np.float64)
t_quantiles = 255 * t_quantiles / t_quantiles[-1]
interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)
return interp_t_values[bin_idx].reshape(oldshape)
def color_hist_match(src_im, tar_im, hist_match_threshold=255):
h,w,c = src_im.shape
matched_R = channel_hist_match(src_im[:,:,0], tar_im[:,:,0], hist_match_threshold, None)
matched_G = channel_hist_match(src_im[:,:,1], tar_im[:,:,1], hist_match_threshold, None)
matched_B = channel_hist_match(src_im[:,:,2], tar_im[:,:,2], hist_match_threshold, None)
to_stack = (matched_R, matched_G, matched_B)
for i in range(3, c):
to_stack += ( src_im[:,:,i],)
matched = np.stack(to_stack, axis=-1).astype(src_im.dtype)
return matched
def color_transfer_mix(img_src,img_trg):
img_src = np.clip(img_src*255.0, 0, 255).astype(np.uint8)
img_trg = np.clip(img_trg*255.0, 0, 255).astype(np.uint8)
img_src_lab = cv2.cvtColor(img_src, cv2.COLOR_BGR2LAB)
img_trg_lab = cv2.cvtColor(img_trg, cv2.COLOR_BGR2LAB)
rct_light = np.clip ( linear_color_transfer(img_src_lab[...,0:1].astype(np.float32)/255.0,
img_trg_lab[...,0:1].astype(np.float32)/255.0 )[...,0]*255.0,
0, 255).astype(np.uint8)
img_src_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8)
img_src_lab = cv2.cvtColor(img_src_lab, cv2.COLOR_LAB2BGR)
img_trg_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8)
img_trg_lab = cv2.cvtColor(img_trg_lab, cv2.COLOR_LAB2BGR)
img_rct = color_transfer_sot( img_src_lab.astype(np.float32), img_trg_lab.astype(np.float32) )
img_rct = np.clip(img_rct, 0, 255).astype(np.uint8)
img_rct = cv2.cvtColor(img_rct, cv2.COLOR_BGR2LAB)
img_rct[...,0] = rct_light
img_rct = cv2.cvtColor(img_rct, cv2.COLOR_LAB2BGR)
return (img_rct / 255.0).astype(np.float32)
def color_transfer(ct_mode, img_src, img_trg):
"""
color transfer for [0,1] float32 inputs
"""
if ct_mode == 'lct':
out = linear_color_transfer (img_src, img_trg)
elif ct_mode == 'rct':
out = reinhard_color_transfer ( np.clip( img_src*255, 0, 255 ).astype(np.uint8),
np.clip( img_trg*255, 0, 255 ).astype(np.uint8) )
out = np.clip( out.astype(np.float32) / 255.0, 0.0, 1.0)
elif ct_mode == 'mkl':
out = color_transfer_mkl (img_src, img_trg)
elif ct_mode == 'idt':
out = color_transfer_idt (img_src, img_trg)
elif ct_mode == 'sot':
out = color_transfer_sot (img_src, img_trg)
out = np.clip( out, 0.0, 1.0)
else:
raise ValueError(f"unknown ct_mode {ct_mode}")
return out
-47
View File
@@ -1,47 +0,0 @@
import numpy as np
def normalize_channels(img, target_channels):
img_shape_len = len(img.shape)
if img_shape_len == 2:
h, w = img.shape
c = 0
elif img_shape_len == 3:
h, w, c = img.shape
else:
raise ValueError("normalize: incorrect image dimensions.")
if c == 0 and target_channels > 0:
img = img[...,np.newaxis]
c = 1
if c == 1 and target_channels > 1:
img = np.repeat (img, target_channels, -1)
c = target_channels
if c > target_channels:
img = img[...,0:target_channels]
c = target_channels
return img
def cut_odd_image(img):
h, w, c = img.shape
wm, hm = w % 2, h % 2
if wm + hm != 0:
img = img[0:h-hm,0:w-wm,:]
return img
def overlay_alpha_image(img_target, img_source, xy_offset=(0,0) ):
(h,w,c) = img_source.shape
if c != 4:
raise ValueError("overlay_alpha_image, img_source must have 4 channels")
x1, x2 = xy_offset[0], xy_offset[0] + w
y1, y2 = xy_offset[1], xy_offset[1] + h
alpha_s = img_source[:, :, 3] / 255.0
alpha_l = 1.0 - alpha_s
for c in range(0, 3):
img_target[y1:y2, x1:x2, c] = (alpha_s * img_source[:, :, c] +
alpha_l * img_target[y1:y2, x1:x2, c])
-13
View File
@@ -1,13 +0,0 @@
import numpy as np
import cv2
def draw_polygon (image, points, color, thickness = 1):
points_len = len(points)
for i in range (0, points_len):
p0 = tuple( points[i] )
p1 = tuple( points[ (i+1) % points_len] )
cv2.line (image, p0, p1, color, thickness=thickness)
def draw_rect(image, rect, color, thickness=1):
l,t,r,b = rect
draw_polygon (image, [ (l,t), (r,t), (r,b), (l,b ) ], color, thickness)
@@ -1,45 +0,0 @@
import numpy as np
import cv2
def equalize_and_stack_square (images, axis=1):
max_c = max ([ 1 if len(image.shape) == 2 else image.shape[2] for image in images ] )
target_wh = 99999
for i,image in enumerate(images):
if len(image.shape) == 2:
h,w = image.shape
c = 1
else:
h,w,c = image.shape
if h < target_wh:
target_wh = h
if w < target_wh:
target_wh = w
for i,image in enumerate(images):
if len(image.shape) == 2:
h,w = image.shape
c = 1
else:
h,w,c = image.shape
if c < max_c:
if c == 1:
if len(image.shape) == 2:
image = np.expand_dims ( image, -1 )
image = np.concatenate ( (image,)*max_c, -1 )
elif c == 2: #GA
image = np.expand_dims ( image[...,0], -1 )
image = np.concatenate ( (image,)*max_c, -1 )
else:
image = np.concatenate ( (image, np.ones((h,w,max_c - c))), -1 )
if h != target_wh or w != target_wh:
image = cv2.resize ( image, (target_wh, target_wh) )
h,w,c = image.shape
images[i] = image
return np.concatenate ( images, axis = 1 )
-278
View File
@@ -1,278 +0,0 @@
"""
Copyright (c) 2009-2010 Arizona Board of Regents. All Rights Reserved.
Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu)
Image, Video, and Usabilty (IVU) Lab, http://ivulab.asu.edu , Arizona State University
This copyright statement may not be removed from any file containing it or from modifications to these files.
This copyright notice must also be included in any file or product that is derived from the source files.
Redistribution and use of this code in source and binary forms, with or without modification, are permitted provided that the
following conditions are met:
- Redistribution's of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
- Redistribution's in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the distribution.
- The Image, Video, and Usability Laboratory (IVU Lab, http://ivulab.asu.edu) is acknowledged in any publication that
reports research results using this code, copies of this code, or modifications of this code.
The code and our papers are to be cited in the bibliography as:
N. D. Narvekar and L. J. Karam, "CPBD Sharpness Metric Software", http://ivulab.asu.edu/Quality/CPBD
N. D. Narvekar and L. J. Karam, "A No-Reference Image Blur Metric Based on the Cumulative
Probability of Blur Detection (CPBD)," accepted and to appear in the IEEE Transactions on Image Processing, 2011.
N. D. Narvekar and L. J. Karam, "An Improved No-Reference Sharpness Metric Based on the Probability of Blur Detection," International Workshop on Video Processing and Quality Metrics for Consumer Electronics (VPQM), January 2010, http://www.vpqm.org (pdf)
N. D. Narvekar and L. J. Karam, "A No Reference Perceptual Quality Metric based on Cumulative Probability of Blur Detection," First International Workshop on the Quality of Multimedia Experience (QoMEX), pp. 87-91, July 2009.
DISCLAIMER:
This software is provided by the copyright holders and contributors "as is" and any express or implied warranties, including, but not limited to, the implied warranties of merchantability and fitness for a particular purpose are disclaimed. In no event shall the Arizona Board of Regents, Arizona State University, IVU Lab members, authors or contributors be liable for any direct, indirect, incidental, special, exemplary, or consequential damages (including, but not limited to, procurement of substitute
goods or services; loss of use, data, or profits; or business interruption) however caused and on any theory of liability, whether in contract, strict liability, or tort (including negligence or otherwise) arising in any way out of the use of this software, even if advised of the possibility of such damage.
"""
import numpy as np
import cv2
from math import atan2, pi
def sobel(image):
# type: (numpy.ndarray) -> numpy.ndarray
"""
Find edges using the Sobel approximation to the derivatives.
Inspired by the [Octave implementation](https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l196).
"""
from skimage.filters.edges import HSOBEL_WEIGHTS
h1 = np.array(HSOBEL_WEIGHTS)
h1 /= np.sum(abs(h1)) # normalize h1
from scipy.ndimage import convolve
strength2 = np.square(convolve(image, h1.T))
# Note: https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l59
thresh2 = 2 * np.sqrt(np.mean(strength2))
strength2[strength2 <= thresh2] = 0
return _simple_thinning(strength2)
def _simple_thinning(strength):
# type: (numpy.ndarray) -> numpy.ndarray
"""
Perform a very simple thinning.
Inspired by the [Octave implementation](https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l512).
"""
num_rows, num_cols = strength.shape
zero_column = np.zeros((num_rows, 1))
zero_row = np.zeros((1, num_cols))
x = (
(strength > np.c_[zero_column, strength[:, :-1]]) &
(strength > np.c_[strength[:, 1:], zero_column])
)
y = (
(strength > np.r_[zero_row, strength[:-1, :]]) &
(strength > np.r_[strength[1:, :], zero_row])
)
return x | y
# threshold to characterize blocks as edge/non-edge blocks
THRESHOLD = 0.002
# fitting parameter
BETA = 3.6
# block size
BLOCK_HEIGHT, BLOCK_WIDTH = (64, 64)
# just noticeable widths based on the perceptual experiments
WIDTH_JNB = np.concatenate([5*np.ones(51), 3*np.ones(205)])
def compute(image):
# type: (numpy.ndarray) -> float
"""Compute the sharpness metric for the given data."""
# convert the image to double for further processing
image = image.astype(np.float64)
# edge detection using canny and sobel canny edge detection is done to
# classify the blocks as edge or non-edge blocks and sobel edge
# detection is done for the purpose of edge width measurement.
from skimage.feature import canny
canny_edges = canny(image)
sobel_edges = sobel(image)
# edge width calculation
marziliano_widths = marziliano_method(sobel_edges, image)
# sharpness metric calculation
return _calculate_sharpness_metric(image, canny_edges, marziliano_widths)
def marziliano_method(edges, image):
# type: (numpy.ndarray, numpy.ndarray) -> numpy.ndarray
"""
Calculate the widths of the given edges.
:return: A matrix with the same dimensions as the given image with 0's at
non-edge locations and edge-widths at the edge locations.
"""
# `edge_widths` consists of zero and non-zero values. A zero value
# indicates that there is no edge at that position and a non-zero value
# indicates that there is an edge at that position and the value itself
# gives the edge width.
edge_widths = np.zeros(image.shape)
# find the gradient for the image
gradient_y, gradient_x = np.gradient(image)
# dimensions of the image
img_height, img_width = image.shape
# holds the angle information of the edges
edge_angles = np.zeros(image.shape)
# calculate the angle of the edges
for row in range(img_height):
for col in range(img_width):
if gradient_x[row, col] != 0:
edge_angles[row, col] = atan2(gradient_y[row, col], gradient_x[row, col]) * (180 / pi)
elif gradient_x[row, col] == 0 and gradient_y[row, col] == 0:
edge_angles[row,col] = 0
elif gradient_x[row, col] == 0 and gradient_y[row, col] == pi/2:
edge_angles[row, col] = 90
if np.any(edge_angles):
# quantize the angle
quantized_angles = 45 * np.round(edge_angles / 45)
for row in range(1, img_height - 1):
for col in range(1, img_width - 1):
if edges[row, col] == 1:
# gradient angle = 180 or -180
if quantized_angles[row, col] == 180 or quantized_angles[row, col] == -180:
for margin in range(100 + 1):
inner_border = (col - 1) - margin
outer_border = (col - 2) - margin
# outside image or intensity increasing from left to right
if outer_border < 0 or (image[row, outer_border] - image[row, inner_border]) <= 0:
break
width_left = margin + 1
for margin in range(100 + 1):
inner_border = (col + 1) + margin
outer_border = (col + 2) + margin
# outside image or intensity increasing from left to right
if outer_border >= img_width or (image[row, outer_border] - image[row, inner_border]) >= 0:
break
width_right = margin + 1
edge_widths[row, col] = width_left + width_right
# gradient angle = 0
if quantized_angles[row, col] == 0:
for margin in range(100 + 1):
inner_border = (col - 1) - margin
outer_border = (col - 2) - margin
# outside image or intensity decreasing from left to right
if outer_border < 0 or (image[row, outer_border] - image[row, inner_border]) >= 0:
break
width_left = margin + 1
for margin in range(100 + 1):
inner_border = (col + 1) + margin
outer_border = (col + 2) + margin
# outside image or intensity decreasing from left to right
if outer_border >= img_width or (image[row, outer_border] - image[row, inner_border]) <= 0:
break
width_right = margin + 1
edge_widths[row, col] = width_right + width_left
return edge_widths
def _calculate_sharpness_metric(image, edges, edge_widths):
# type: (numpy.array, numpy.array, numpy.array) -> numpy.float64
# get the size of image
img_height, img_width = image.shape
total_num_edges = 0
hist_pblur = np.zeros(101)
# maximum block indices
num_blocks_vertically = int(img_height / BLOCK_HEIGHT)
num_blocks_horizontally = int(img_width / BLOCK_WIDTH)
# loop over the blocks
for i in range(num_blocks_vertically):
for j in range(num_blocks_horizontally):
# get the row and col indices for the block pixel positions
rows = slice(BLOCK_HEIGHT * i, BLOCK_HEIGHT * (i + 1))
cols = slice(BLOCK_WIDTH * j, BLOCK_WIDTH * (j + 1))
if is_edge_block(edges[rows, cols], THRESHOLD):
block_widths = edge_widths[rows, cols]
# rotate block to simulate column-major boolean indexing
block_widths = np.rot90(np.flipud(block_widths), 3)
block_widths = block_widths[block_widths != 0]
block_contrast = get_block_contrast(image[rows, cols])
block_jnb = WIDTH_JNB[block_contrast]
# calculate the probability of blur detection at the edges
# detected in the block
prob_blur_detection = 1 - np.exp(-abs(block_widths/block_jnb) ** BETA)
# update the statistics using the block information
for probability in prob_blur_detection:
bucket = int(round(probability * 100))
hist_pblur[bucket] += 1
total_num_edges += 1
# normalize the pdf
if total_num_edges > 0:
hist_pblur = hist_pblur / total_num_edges
# calculate the sharpness metric
return np.sum(hist_pblur[:64])
def is_edge_block(block, threshold):
# type: (numpy.ndarray, float) -> bool
"""Decide whether the given block is an edge block."""
return np.count_nonzero(block) > (block.size * threshold)
def get_block_contrast(block):
# type: (numpy.ndarray) -> int
return int(np.max(block) - np.min(block))
def estimate_sharpness(image):
if image.ndim == 3:
if image.shape[2] > 1:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
image = image[...,0]
return compute(image)
-87
View File
@@ -1,87 +0,0 @@
import numpy as np
from .blursharpen import LinearMotionBlur
import cv2
def apply_random_rgb_levels(img, mask=None, rnd_state=None):
if rnd_state is None:
rnd_state = np.random
np_rnd = rnd_state.rand
inBlack = np.array([np_rnd()*0.25 , np_rnd()*0.25 , np_rnd()*0.25], dtype=np.float32)
inWhite = np.array([1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25], dtype=np.float32)
inGamma = np.array([0.5+np_rnd(), 0.5+np_rnd(), 0.5+np_rnd()], dtype=np.float32)
outBlack = np.array([np_rnd()*0.25 , np_rnd()*0.25 , np_rnd()*0.25], dtype=np.float32)
outWhite = np.array([1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25], dtype=np.float32)
result = np.clip( (img - inBlack) / (inWhite - inBlack), 0, 1 )
result = ( result ** (1/inGamma) ) * (outWhite - outBlack) + outBlack
result = np.clip(result, 0, 1)
if mask is not None:
result = img*(1-mask) + result*mask
return result
def apply_random_hsv_shift(img, mask=None, rnd_state=None):
if rnd_state is None:
rnd_state = np.random
h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
h = ( h + rnd_state.randint(360) ) % 360
s = np.clip ( s + rnd_state.random()-0.5, 0, 1 )
v = np.clip ( v + rnd_state.random()-0.5, 0, 1 )
result = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR) , 0, 1 )
if mask is not None:
result = img*(1-mask) + result*mask
return result
def apply_random_motion_blur( img, chance, mb_max_size, mask=None, rnd_state=None ):
if rnd_state is None:
rnd_state = np.random
mblur_rnd_kernel = rnd_state.randint(mb_max_size)+1
mblur_rnd_deg = rnd_state.randint(360)
result = img
if rnd_state.randint(100) < np.clip(chance, 0, 100):
result = LinearMotionBlur (result, mblur_rnd_kernel, mblur_rnd_deg )
if mask is not None:
result = img*(1-mask) + result*mask
return result
def apply_random_gaussian_blur( img, chance, kernel_max_size, mask=None, rnd_state=None ):
if rnd_state is None:
rnd_state = np.random
result = img
if rnd_state.randint(100) < np.clip(chance, 0, 100):
gblur_rnd_kernel = rnd_state.randint(kernel_max_size)*2+1
result = cv2.GaussianBlur(result, (gblur_rnd_kernel,)*2 , 0)
if mask is not None:
result = img*(1-mask) + result*mask
return result
def apply_random_bilinear_resize( img, chance, max_size_per, mask=None, rnd_state=None ):
if rnd_state is None:
rnd_state = np.random
result = img
if rnd_state.randint(100) < np.clip(chance, 0, 100):
h,w,c = result.shape
trg = rnd_state.rand()
rw = w - int( trg * int(w*(max_size_per/100.0)) )
rh = h - int( trg * int(h*(max_size_per/100.0)) )
result = cv2.resize (result, (rw,rh), interpolation=cv2.INTER_LINEAR )
result = cv2.resize (result, (w,h), interpolation=cv2.INTER_LINEAR )
if mask is not None:
result = img*(1-mask) + result*mask
return result
-37
View File
@@ -1,37 +0,0 @@
import numpy as np
import cv2
from scipy.spatial import Delaunay
def applyAffineTransform(src, srcTri, dstTri, size) :
warpMat = cv2.getAffineTransform( np.float32(srcTri), np.float32(dstTri) )
return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 )
def morphTriangle(dst_img, src_img, st, dt) :
(h,w,c) = dst_img.shape
sr = np.array( cv2.boundingRect(np.float32(st)) )
dr = np.array( cv2.boundingRect(np.float32(dt)) )
sRect = st - sr[0:2]
dRect = dt - dr[0:2]
d_mask = np.zeros((dr[3], dr[2], c), dtype = np.float32)
cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0);
imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]]
size = (dr[2], dr[3])
warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size)
if c == 1:
warpImage1 = np.expand_dims( warpImage1, -1 )
dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]] = dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]]*(1-d_mask) + warpImage1 * d_mask
def morph_by_points (image, sp, dp):
if sp.shape != dp.shape:
raise ValueError ('morph_by_points() sp.shape != dp.shape')
(h,w,c) = image.shape
result_image = np.zeros(image.shape, dtype = image.dtype)
for tri in Delaunay(dp).simplices:
morphTriangle(result_image, image, sp[tri], dp[tri])
return result_image
-14
View File
@@ -1,14 +0,0 @@
import numpy as np
import cv2
from PIL import Image
#n_colors = [0..256]
def reduce_colors (img_bgr, n_colors):
img_rgb = (img_bgr[...,::-1] * 255.0).astype(np.uint8)
img_rgb_pil = Image.fromarray(img_rgb)
img_rgb_pil_p = img_rgb_pil.convert('P', palette=Image.ADAPTIVE, colors=n_colors)
img_rgb_p = img_rgb_pil_p.convert('RGB')
img_bgr = cv2.cvtColor( np.array(img_rgb_p, dtype=np.float32) / 255.0, cv2.COLOR_RGB2BGR )
return img_bgr
-2
View File
@@ -1,2 +0,0 @@
from .draw import *
from .calc import *
-25
View File
@@ -1,25 +0,0 @@
import numpy as np
import numpy.linalg as npla
def dist_to_edges(pts, pt, is_closed=False):
"""
returns array of dist from pt to edge and projection pt to edges
"""
if is_closed:
a = pts
b = np.concatenate( (pts[1:,:], pts[0:1,:]), axis=0 )
else:
a = pts[:-1,:]
b = pts[1:,:]
pa = pt-a
ba = b-a
div = np.einsum('ij,ij->i', ba, ba)
div[div==0]=1
h = np.clip( np.einsum('ij,ij->i', pa, ba) / div, 0, 1 )
x = npla.norm ( pa - ba*h[...,None], axis=1 )
return x, a+ba*h[...,None]
-44
View File
@@ -1,44 +0,0 @@
"""
Signed distance drawing functions using numpy.
"""
import numpy as np
from numpy import linalg as npla
def circle_faded( hw, center, fade_dists ):
"""
returns drawn circle in [h,w,1] output range [0..1.0] float32
hw = [h,w] resolution
center = [y,x] center of circle
fade_dists = [fade_start, fade_end] fade values
"""
h,w = hw
pts = np.empty( (h,w,2), dtype=np.float32 )
pts[...,1] = np.arange(h)[None,:]
pts[...,0] = np.arange(w)[:,None]
pts = pts.reshape ( (h*w, -1) )
pts_dists = np.abs ( npla.norm(pts-center, axis=-1) )
if fade_dists[1] == 0:
fade_dists[1] = 1
pts_dists = ( pts_dists - fade_dists[0] ) / fade_dists[1]
pts_dists = np.clip( 1-pts_dists, 0, 1)
return pts_dists.reshape ( (h,w,1) ).astype(np.float32)
def random_circle_faded ( hw, rnd_state=None ):
if rnd_state is None:
rnd_state = np.random
h,w = hw
hw_max = max(h,w)
fade_start = rnd_state.randint(hw_max)
fade_end = fade_start + rnd_state.randint(hw_max- fade_start)
return circle_faded (hw, [ rnd_state.randint(h), rnd_state.randint(w) ],
[fade_start, fade_end] )
-64
View File
@@ -1,64 +0,0 @@
import localization
import numpy as np
from PIL import Image, ImageDraw, ImageFont
pil_fonts = {}
def _get_pil_font (font, size):
global pil_fonts
try:
font_str_id = '%s_%d' % (font, size)
if font_str_id not in pil_fonts.keys():
pil_fonts[font_str_id] = ImageFont.truetype(font + ".ttf", size=size, encoding="unic")
pil_font = pil_fonts[font_str_id]
return pil_font
except:
return ImageFont.load_default()
def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None):
h,w,c = shape
try:
pil_font = _get_pil_font( localization.get_default_ttf_font_name() , h-2)
canvas = Image.new('RGB', (w,h) , (0,0,0) )
draw = ImageDraw.Draw(canvas)
offset = ( 0, 0)
draw.text(offset, text, font=pil_font, fill=tuple((np.array(color)*255).astype(np.int)) )
result = np.asarray(canvas) / 255
if c > 3:
result = np.concatenate ( (result, np.ones ((h,w,c-3)) ), axis=-1 )
elif c < 3:
result = result[...,0:c]
return result
except:
return np.zeros ( (h,w,c) )
def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None):
h,w,c = image.shape
l,t,r,b = rect
l = np.clip (l, 0, w-1)
r = np.clip (r, 0, w-1)
t = np.clip (t, 0, h-1)
b = np.clip (b, 0, h-1)
image[t:b, l:r] += get_text_image ( (b-t,r-l,c) , text, color, border, font )
def draw_text_lines (image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
text_lines_len = len(text_lines)
if text_lines_len == 0:
return
l,t,r,b = rect
h = b-t
h_per_line = h // text_lines_len
for i in range(0, text_lines_len):
draw_text (image, (l, i*h_per_line, r, (i+1)*h_per_line), text_lines[i], color, border, font)
def get_draw_text_lines ( image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
image = np.zeros ( image.shape, dtype=np.float )
draw_text_lines ( image, rect, text_lines, color, border, font)
return image
-69
View File
@@ -1,69 +0,0 @@
import numpy as np
import cv2
from core import randomex
def gen_warp_params (w, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], rnd_state=None ):
if rnd_state is None:
rnd_state = np.random
rw = None
if w < 64:
rw = w
w = 64
rotation = rnd_state.uniform( rotation_range[0], rotation_range[1] )
scale = rnd_state.uniform(1 +scale_range[0], 1 +scale_range[1])
tx = rnd_state.uniform( tx_range[0], tx_range[1] )
ty = rnd_state.uniform( ty_range[0], ty_range[1] )
p_flip = flip and rnd_state.randint(10) < 4
#random warp by grid
cell_size = [ w // (2**i) for i in range(1,4) ] [ rnd_state.randint(3) ]
cell_count = w // cell_size + 1
grid_points = np.linspace( 0, w, cell_count)
mapx = np.broadcast_to(grid_points, (cell_count, cell_count)).copy()
mapy = mapx.T
mapx[1:-1,1:-1] = mapx[1:-1,1:-1] + randomex.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
mapy[1:-1,1:-1] = mapy[1:-1,1:-1] + randomex.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
half_cell_size = cell_size // 2
mapx = cv2.resize(mapx, (w+cell_size,)*2 )[half_cell_size:-half_cell_size,half_cell_size:-half_cell_size].astype(np.float32)
mapy = cv2.resize(mapy, (w+cell_size,)*2 )[half_cell_size:-half_cell_size,half_cell_size:-half_cell_size].astype(np.float32)
#random transform
random_transform_mat = cv2.getRotationMatrix2D((w // 2, w // 2), rotation, scale)
random_transform_mat[:, 2] += (tx*w, ty*w)
params = dict()
params['mapx'] = mapx
params['mapy'] = mapy
params['rmat'] = random_transform_mat
params['w'] = w
params['rw'] = rw
params['flip'] = p_flip
return params
def warp_by_params (params, img, can_warp, can_transform, can_flip, border_replicate, cv2_inter=cv2.INTER_CUBIC):
rw = params['rw']
if (can_warp or can_transform) and rw is not None:
img = cv2.resize(img, (64,64), interpolation=cv2_inter)
if can_warp:
img = cv2.remap(img, params['mapx'], params['mapy'], cv2_inter )
if can_transform:
img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2_inter )
if (can_warp or can_transform) and rw is not None:
img = cv2.resize(img, (rw,rw), interpolation=cv2_inter)
if len(img.shape) == 2:
img = img[...,None]
if can_flip and params['flip']:
img = img[:,::-1,...]
return img
-1
View File
@@ -1 +0,0 @@
from .interact import interact
-581
View File
@@ -1,581 +0,0 @@
import multiprocessing
import os
import sys
import threading
import time
import types
import colorama
import cv2
import numpy as np
from tqdm import tqdm
from core import stdex
try:
import IPython #if success we are in colab
from IPython.display import display, clear_output
import PIL
import matplotlib.pyplot as plt
is_colab = True
except:
is_colab = False
yn_str = {True:'y',False:'n'}
class InteractBase(object):
EVENT_LBUTTONDOWN = 1
EVENT_LBUTTONUP = 2
EVENT_MBUTTONDOWN = 3
EVENT_MBUTTONUP = 4
EVENT_RBUTTONDOWN = 5
EVENT_RBUTTONUP = 6
EVENT_MOUSEWHEEL = 10
def __init__(self):
self.named_windows = {}
self.capture_mouse_windows = {}
self.capture_keys_windows = {}
self.mouse_events = {}
self.key_events = {}
self.pg_bar = None
self.focus_wnd_name = None
self.error_log_line_prefix = '/!\\ '
self.process_messages_callbacks = {}
def is_support_windows(self):
return False
def is_colab(self):
return False
def on_destroy_all_windows(self):
raise NotImplemented
def on_create_window (self, wnd_name):
raise NotImplemented
def on_destroy_window (self, wnd_name):
raise NotImplemented
def on_show_image (self, wnd_name, img):
raise NotImplemented
def on_capture_mouse (self, wnd_name):
raise NotImplemented
def on_capture_keys (self, wnd_name):
raise NotImplemented
def on_process_messages(self, sleep_time=0):
raise NotImplemented
def on_wait_any_key(self):
raise NotImplemented
def log_info(self, msg, end='\n'):
if self.pg_bar is not None:
print ("\n")
print (msg, end=end)
def log_err(self, msg, end='\n'):
if self.pg_bar is not None:
print ("\n")
print (f'{self.error_log_line_prefix}{msg}', end=end)
def named_window(self, wnd_name):
if wnd_name not in self.named_windows:
#we will show window only on first show_image
self.named_windows[wnd_name] = 0
self.focus_wnd_name = wnd_name
else: print("named_window: ", wnd_name, " already created.")
def destroy_all_windows(self):
if len( self.named_windows ) != 0:
self.on_destroy_all_windows()
self.named_windows = {}
self.capture_mouse_windows = {}
self.capture_keys_windows = {}
self.mouse_events = {}
self.key_events = {}
self.focus_wnd_name = None
def destroy_window(self, wnd_name):
if wnd_name in self.named_windows:
self.on_destroy_window(wnd_name)
self.named_windows.pop(wnd_name)
if wnd_name == self.focus_wnd_name:
self.focus_wnd_name = list(self.named_windows.keys())[-1] if len( self.named_windows ) != 0 else None
if wnd_name in self.capture_mouse_windows:
self.capture_mouse_windows.pop(wnd_name)
if wnd_name in self.capture_keys_windows:
self.capture_keys_windows.pop(wnd_name)
if wnd_name in self.mouse_events:
self.mouse_events.pop(wnd_name)
if wnd_name in self.key_events:
self.key_events.pop(wnd_name)
def show_image(self, wnd_name, img):
if wnd_name in self.named_windows:
if self.named_windows[wnd_name] == 0:
self.named_windows[wnd_name] = 1
self.on_create_window(wnd_name)
if wnd_name in self.capture_mouse_windows:
self.capture_mouse(wnd_name)
self.on_show_image(wnd_name,img)
else: print("show_image: named_window ", wnd_name, " not found.")
def capture_mouse(self, wnd_name):
if wnd_name in self.named_windows:
self.capture_mouse_windows[wnd_name] = True
if self.named_windows[wnd_name] == 1:
self.on_capture_mouse(wnd_name)
else: print("capture_mouse: named_window ", wnd_name, " not found.")
def capture_keys(self, wnd_name):
if wnd_name in self.named_windows:
if wnd_name not in self.capture_keys_windows:
self.capture_keys_windows[wnd_name] = True
self.on_capture_keys(wnd_name)
else: print("capture_keys: already set for window ", wnd_name)
else: print("capture_keys: named_window ", wnd_name, " not found.")
def progress_bar(self, desc, total, leave=True, initial=0):
if self.pg_bar is None:
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
else: print("progress_bar: already set.")
def progress_bar_inc(self, c):
if self.pg_bar is not None:
self.pg_bar.n += c
self.pg_bar.refresh()
else: print("progress_bar not set.")
def progress_bar_close(self):
if self.pg_bar is not None:
self.pg_bar.close()
self.pg_bar = None
else: print("progress_bar not set.")
def progress_bar_generator(self, data, desc=None, leave=True, initial=0):
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
for x in self.pg_bar:
yield x
self.pg_bar.close()
self.pg_bar = None
def add_process_messages_callback(self, func ):
tid = threading.get_ident()
callbacks = self.process_messages_callbacks.get(tid, None)
if callbacks is None:
callbacks = []
self.process_messages_callbacks[tid] = callbacks
callbacks.append ( func )
def process_messages(self, sleep_time=0):
callbacks = self.process_messages_callbacks.get(threading.get_ident(), None)
if callbacks is not None:
for func in callbacks:
func()
self.on_process_messages(sleep_time)
def wait_any_key(self):
self.on_wait_any_key()
def add_mouse_event(self, wnd_name, x, y, ev, flags):
if wnd_name not in self.mouse_events:
self.mouse_events[wnd_name] = []
self.mouse_events[wnd_name] += [ (x, y, ev, flags) ]
def add_key_event(self, wnd_name, ord_key, ctrl_pressed, alt_pressed, shift_pressed):
if wnd_name not in self.key_events:
self.key_events[wnd_name] = []
self.key_events[wnd_name] += [ (ord_key, chr(ord_key) if ord_key <= 255 else chr(0), ctrl_pressed, alt_pressed, shift_pressed) ]
def get_mouse_events(self, wnd_name):
ar = self.mouse_events.get(wnd_name, [])
self.mouse_events[wnd_name] = []
return ar
def get_key_events(self, wnd_name):
ar = self.key_events.get(wnd_name, [])
self.key_events[wnd_name] = []
return ar
def input(self, s):
return input(s)
def input_number(self, s, default_value, valid_list=None, show_default_value=True, add_info=None, help_message=None):
if show_default_value and default_value is not None:
s = f"[{default_value}] {s}"
if add_info is not None or \
help_message is not None:
s += " ("
if add_info is not None:
s += f" {add_info}"
if help_message is not None:
s += " ?:help"
if add_info is not None or \
help_message is not None:
s += " )"
s += " : "
while True:
try:
inp = input(s)
if len(inp) == 0:
result = default_value
break
if help_message is not None and inp == '?':
print (help_message)
continue
i = float(inp)
if (valid_list is not None) and (i not in valid_list):
result = default_value
break
result = i
break
except:
result = default_value
break
print(result)
return result
def input_int(self, s, default_value, valid_range=None, valid_list=None, add_info=None, show_default_value=True, help_message=None):
if show_default_value:
if len(s) != 0:
s = f"[{default_value}] {s}"
else:
s = f"[{default_value}]"
if add_info is not None or \
valid_range is not None or \
help_message is not None:
s += " ("
if valid_range is not None:
s += f" {valid_range[0]}-{valid_range[1]}"
if add_info is not None:
s += f" {add_info}"
if help_message is not None:
s += " ?:help"
if add_info is not None or \
valid_range is not None or \
help_message is not None:
s += " )"
s += " : "
while True:
try:
inp = input(s)
if len(inp) == 0:
raise ValueError("")
if help_message is not None and inp == '?':
print (help_message)
continue
i = int(inp)
if valid_range is not None:
i = int(np.clip(i, valid_range[0], valid_range[1]))
if (valid_list is not None) and (i not in valid_list):
i = default_value
result = i
break
except:
result = default_value
break
print (result)
return result
def input_bool(self, s, default_value, help_message=None):
s = f"[{yn_str[default_value]}] {s} ( y/n"
if help_message is not None:
s += " ?:help"
s += " ) : "
while True:
try:
inp = input(s)
if len(inp) == 0:
raise ValueError("")
if help_message is not None and inp == '?':
print (help_message)
continue
return bool ( {"y":True,"n":False}.get(inp.lower(), default_value) )
except:
print ( "y" if default_value else "n" )
return default_value
def input_str(self, s, default_value=None, valid_list=None, show_default_value=True, help_message=None):
if show_default_value and default_value is not None:
s = f"[{default_value}] {s}"
if valid_list is not None or \
help_message is not None:
s += " ("
if valid_list is not None:
s += " " + "/".join(valid_list)
if help_message is not None:
s += " ?:help"
if valid_list is not None or \
help_message is not None:
s += " )"
s += " : "
while True:
try:
inp = input(s)
if len(inp) == 0:
if default_value is None:
print("")
return None
result = default_value
break
if help_message is not None and inp == '?':
print(help_message)
continue
if valid_list is not None:
if inp.lower() in valid_list:
result = inp.lower()
break
if inp in valid_list:
result = inp
break
continue
result = inp
break
except:
result = default_value
break
print(result)
return result
def input_process(self, stdin_fd, sq, str):
sys.stdin = os.fdopen(stdin_fd)
try:
inp = input (str)
sq.put (True)
except:
sq.put (False)
def input_in_time (self, str, max_time_sec):
sq = multiprocessing.Queue()
p = multiprocessing.Process(target=self.input_process, args=( sys.stdin.fileno(), sq, str))
p.daemon = True
p.start()
t = time.time()
inp = False
while True:
if not sq.empty():
inp = sq.get()
break
if time.time() - t > max_time_sec:
break
p.terminate()
p.join()
old_stdin = sys.stdin
sys.stdin = os.fdopen( os.dup(sys.stdin.fileno()) )
old_stdin.close()
return inp
def input_process_skip_pending(self, stdin_fd):
sys.stdin = os.fdopen(stdin_fd)
while True:
try:
if sys.stdin.isatty():
sys.stdin.read()
except:
pass
def input_skip_pending(self):
if is_colab:
# currently it does not work on Colab
return
"""
skips unnecessary inputs between the dialogs
"""
p = multiprocessing.Process(target=self.input_process_skip_pending, args=( sys.stdin.fileno(), ))
p.daemon = True
p.start()
time.sleep(0.5)
p.terminate()
p.join()
sys.stdin = os.fdopen( sys.stdin.fileno() )
class InteractDesktop(InteractBase):
def __init__(self):
colorama.init()
super().__init__()
def color_red(self):
pass
def is_support_windows(self):
return True
def on_destroy_all_windows(self):
cv2.destroyAllWindows()
def on_create_window (self, wnd_name):
cv2.namedWindow(wnd_name)
def on_destroy_window (self, wnd_name):
cv2.destroyWindow(wnd_name)
def on_show_image (self, wnd_name, img):
cv2.imshow (wnd_name, img)
def on_capture_mouse (self, wnd_name):
self.last_xy = (0,0)
def onMouse(event, x, y, flags, param):
(inst, wnd_name) = param
if event == cv2.EVENT_LBUTTONDOWN: ev = InteractBase.EVENT_LBUTTONDOWN
elif event == cv2.EVENT_LBUTTONUP: ev = InteractBase.EVENT_LBUTTONUP
elif event == cv2.EVENT_RBUTTONDOWN: ev = InteractBase.EVENT_RBUTTONDOWN
elif event == cv2.EVENT_RBUTTONUP: ev = InteractBase.EVENT_RBUTTONUP
elif event == cv2.EVENT_MBUTTONDOWN: ev = InteractBase.EVENT_MBUTTONDOWN
elif event == cv2.EVENT_MBUTTONUP: ev = InteractBase.EVENT_MBUTTONUP
elif event == cv2.EVENT_MOUSEWHEEL:
ev = InteractBase.EVENT_MOUSEWHEEL
x,y = self.last_xy #fix opencv bug when window size more than screen size
else: ev = 0
self.last_xy = (x,y)
inst.add_mouse_event (wnd_name, x, y, ev, flags)
cv2.setMouseCallback(wnd_name, onMouse, (self,wnd_name) )
def on_capture_keys (self, wnd_name):
pass
def on_process_messages(self, sleep_time=0):
has_windows = False
has_capture_keys = False
if len(self.named_windows) != 0:
has_windows = True
if len(self.capture_keys_windows) != 0:
has_capture_keys = True
if has_windows or has_capture_keys:
wait_key_time = max(1, int(sleep_time*1000) )
ord_key = cv2.waitKeyEx(wait_key_time)
shift_pressed = False
if ord_key != -1:
chr_key = chr(ord_key) if ord_key <= 255 else chr(0)
if chr_key >= 'A' and chr_key <= 'Z':
shift_pressed = True
ord_key += 32
elif chr_key == '?':
shift_pressed = True
ord_key = ord('/')
elif chr_key == '<':
shift_pressed = True
ord_key = ord(',')
elif chr_key == '>':
shift_pressed = True
ord_key = ord('.')
else:
if sleep_time != 0:
time.sleep(sleep_time)
if has_capture_keys and ord_key != -1:
self.add_key_event ( self.focus_wnd_name, ord_key, False, False, shift_pressed)
def on_wait_any_key(self):
cv2.waitKey(0)
class InteractColab(InteractBase):
def is_support_windows(self):
return False
def is_colab(self):
return True
def on_destroy_all_windows(self):
pass
#clear_output()
def on_create_window (self, wnd_name):
pass
#clear_output()
def on_destroy_window (self, wnd_name):
pass
def on_show_image (self, wnd_name, img):
pass
# # cv2 stores colors as BGR; convert to RGB
# if img.ndim == 3:
# if img.shape[2] == 4:
# img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
# else:
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# img = PIL.Image.fromarray(img)
# plt.imshow(img)
# plt.show()
def on_capture_mouse (self, wnd_name):
pass
#print("on_capture_mouse(): Colab does not support")
def on_capture_keys (self, wnd_name):
pass
#print("on_capture_keys(): Colab does not support")
def on_process_messages(self, sleep_time=0):
time.sleep(sleep_time)
def on_wait_any_key(self):
pass
#print("on_wait_any_key(): Colab does not support")
if is_colab:
interact = InteractColab()
else:
interact = InteractDesktop()
-32
View File
@@ -1,32 +0,0 @@
import multiprocessing
from core.interact import interact as io
class MPClassFuncOnDemand():
def __init__(self, class_handle, class_func_name, **class_kwargs):
self.class_handle = class_handle
self.class_func_name = class_func_name
self.class_kwargs = class_kwargs
self.class_func = None
self.s2c = multiprocessing.Queue()
self.c2s = multiprocessing.Queue()
self.lock = multiprocessing.Lock()
io.add_process_messages_callback(self.io_callback)
def io_callback(self):
while not self.c2s.empty():
func_args, func_kwargs = self.c2s.get()
if self.class_func is None:
self.class_func = getattr( self.class_handle(**self.class_kwargs), self.class_func_name)
self.s2c.put ( self.class_func (*func_args, **func_kwargs) )
def __call__(self, *args, **kwargs):
with self.lock:
self.c2s.put ( (args, kwargs) )
return self.s2c.get()
def __getstate__(self):
return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}
-25
View File
@@ -1,25 +0,0 @@
import multiprocessing
from core.interact import interact as io
class MPFunc():
def __init__(self, func):
self.func = func
self.s2c = multiprocessing.Queue()
self.c2s = multiprocessing.Queue()
self.lock = multiprocessing.Lock()
io.add_process_messages_callback(self.io_callback)
def io_callback(self):
while not self.c2s.empty():
func_args, func_kwargs = self.c2s.get()
self.s2c.put ( self.func (*func_args, **func_kwargs) )
def __call__(self, *args, **kwargs):
with self.lock:
self.c2s.put ( (args, kwargs) )
return self.s2c.get()
def __getstate__(self):
return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}
-79
View File
@@ -1,79 +0,0 @@
import multiprocessing
import queue as Queue
import threading
import time
class SubprocessGenerator(object):
@staticmethod
def launch_thread(generator):
generator._start()
@staticmethod
def start_in_parallel( generator_list ):
"""
Start list of generators in parallel
"""
for generator in generator_list:
thread = threading.Thread(target=SubprocessGenerator.launch_thread, args=(generator,) )
thread.daemon = True
thread.start()
while not all ([generator._is_started() for generator in generator_list]):
time.sleep(0.005)
def __init__(self, generator_func, user_param=None, prefetch=2, start_now=True):
super().__init__()
self.prefetch = prefetch
self.generator_func = generator_func
self.user_param = user_param
self.sc_queue = multiprocessing.Queue()
self.cs_queue = multiprocessing.Queue()
self.p = None
if start_now:
self._start()
def _start(self):
if self.p == None:
user_param = self.user_param
self.user_param = None
p = multiprocessing.Process(target=self.process_func, args=(user_param,) )
p.daemon = True
p.start()
self.p = p
def _is_started(self):
return self.p is not None
def process_func(self, user_param):
self.generator_func = self.generator_func(user_param)
while True:
while self.prefetch > -1:
try:
gen_data = next (self.generator_func)
except StopIteration:
self.cs_queue.put (None)
return
self.cs_queue.put (gen_data)
self.prefetch -= 1
self.sc_queue.get()
self.prefetch += 1
def __iter__(self):
return self
def __getstate__(self):
self_dict = self.__dict__.copy()
del self_dict['p']
return self_dict
def __next__(self):
self._start()
gen_data = self.cs_queue.get()
if gen_data is None:
self.p.terminate()
self.p.join()
raise StopIteration()
self.sc_queue.put (1)
return gen_data
-302
View File
@@ -1,302 +0,0 @@
import traceback
import multiprocessing
import time
import sys
from core.interact import interact as io
class Subprocessor(object):
class SilenceException(Exception):
pass
class Cli(object):
def __init__ ( self, client_dict ):
s2c = multiprocessing.Queue()
c2s = multiprocessing.Queue()
self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,s2c,c2s) )
self.s2c = s2c
self.c2s = c2s
self.p.daemon = True
self.p.start()
self.state = None
self.sent_time = None
self.sent_data = None
self.name = None
self.host_dict = None
def kill(self):
self.p.terminate()
self.p.join()
#overridable optional
def on_initialize(self, client_dict):
#initialize your subprocess here using client_dict
pass
#overridable optional
def on_finalize(self):
#finalize your subprocess here
pass
#overridable
def process_data(self, data):
#process 'data' given from host and return result
raise NotImplementedError
#overridable optional
def get_data_name (self, data):
#return string identificator of your 'data'
return "undefined"
def log_info(self, msg): self.c2s.put ( {'op': 'log_info', 'msg':msg } )
def log_err(self, msg): self.c2s.put ( {'op': 'log_err' , 'msg':msg } )
def progress_bar_inc(self, c): self.c2s.put ( {'op': 'progress_bar_inc' , 'c':c } )
def _subprocess_run(self, client_dict, s2c, c2s):
self.c2s = c2s
data = None
is_error = False
try:
self.on_initialize(client_dict)
c2s.put ( {'op': 'init_ok'} )
while True:
msg = s2c.get()
op = msg.get('op','')
if op == 'data':
data = msg['data']
result = self.process_data (data)
c2s.put ( {'op': 'success', 'data' : data, 'result' : result} )
data = None
elif op == 'close':
break
time.sleep(0.001)
self.on_finalize()
c2s.put ( {'op': 'finalized'} )
except Subprocessor.SilenceException as e:
c2s.put ( {'op': 'error', 'data' : data} )
except Exception as e:
err_msg = traceback.format_exc()
c2s.put ( {'op': 'error', 'data' : data, 'err_msg' : err_msg} )
c2s.close()
s2c.close()
self.c2s = None
# disable pickling
def __getstate__(self):
return dict()
def __setstate__(self, d):
self.__dict__.update(d)
#overridable
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=False):
if not issubclass(SubprocessorCli_class, Subprocessor.Cli):
raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli")
self.name = name
self.SubprocessorCli_class = SubprocessorCli_class
self.no_response_time_sec = no_response_time_sec
self.io_loop_sleep_time = io_loop_sleep_time
self.initialize_subprocesses_in_serial = initialize_subprocesses_in_serial
#overridable
def process_info_generator(self):
#yield per process (name, host_dict, client_dict)
raise NotImplementedError
#overridable optional
def on_clients_initialized(self):
#logic when all subprocesses initialized and ready
pass
#overridable optional
def on_clients_finalized(self):
#logic when all subprocess finalized
pass
#overridable
def get_data(self, host_dict):
#return data for processing here
raise NotImplementedError
#overridable
def on_data_return (self, host_dict, data):
#you have to place returned 'data' back to your queue
raise NotImplementedError
#overridable
def on_result (self, host_dict, data, result):
#your logic what to do with 'result' of 'data'
raise NotImplementedError
#overridable
def get_result(self):
#return result that will be returned in func run()
return None
#overridable
def on_tick(self):
#tick in main loop
#return True if system can be finalized when no data in get_data, orelse False
return True
#overridable
def on_check_run(self):
return True
def run(self):
if not self.on_check_run():
return self.get_result()
self.clis = []
def cli_init_dispatcher(cli):
while not cli.c2s.empty():
obj = cli.c2s.get()
op = obj.get('op','')
if op == 'init_ok':
cli.state = 0
elif op == 'log_info':
io.log_info(obj['msg'])
elif op == 'log_err':
io.log_err(obj['msg'])
elif op == 'error':
err_msg = obj.get('err_msg', None)
if err_msg is not None:
io.log_info(f'Error while subprocess initialization: {err_msg}')
cli.kill()
self.clis.remove(cli)
break
#getting info about name of subprocesses, host and client dicts, and spawning them
for name, host_dict, client_dict in self.process_info_generator():
try:
cli = self.SubprocessorCli_class(client_dict)
cli.state = 1
cli.sent_time = 0
cli.sent_data = None
cli.name = name
cli.host_dict = host_dict
self.clis.append (cli)
if self.initialize_subprocesses_in_serial:
while True:
cli_init_dispatcher(cli)
if cli.state == 0:
break
io.process_messages(0.005)
except:
raise Exception (f"Unable to start subprocess {name}. Error: {traceback.format_exc()}")
if len(self.clis) == 0:
raise Exception ("Unable to start Subprocessor '%s' " % (self.name))
#waiting subprocesses their success(or not) initialization
while True:
for cli in self.clis[:]:
cli_init_dispatcher(cli)
if all ([cli.state == 0 for cli in self.clis]):
break
io.process_messages(0.005)
if len(self.clis) == 0:
raise Exception ( "Unable to start subprocesses." )
#ok some processes survived, initialize host logic
self.on_clients_initialized()
#main loop of data processing
while True:
for cli in self.clis[:]:
while not cli.c2s.empty():
obj = cli.c2s.get()
op = obj.get('op','')
if op == 'success':
#success processed data, return data and result to on_result
self.on_result (cli.host_dict, obj['data'], obj['result'])
self.sent_data = None
cli.state = 0
elif op == 'error':
#some error occured while process data, returning chunk to on_data_return
err_msg = obj.get('err_msg', None)
if err_msg is not None:
io.log_info(f'Error while processing data: {err_msg}')
if 'data' in obj.keys():
self.on_data_return (cli.host_dict, obj['data'] )
#and killing process
cli.kill()
self.clis.remove(cli)
elif op == 'log_info':
io.log_info(obj['msg'])
elif op == 'log_err':
io.log_err(obj['msg'])
elif op == 'progress_bar_inc':
io.progress_bar_inc(obj['c'])
for cli in self.clis[:]:
if cli.state == 1:
if cli.sent_time != 0 and self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec:
#subprocess busy too long
print ( '%s doesnt response, terminating it.' % (cli.name) )
self.on_data_return (cli.host_dict, cli.sent_data )
cli.kill()
self.clis.remove(cli)
for cli in self.clis[:]:
if cli.state == 0:
#free state of subprocess, get some data from get_data
data = self.get_data(cli.host_dict)
if data is not None:
#and send it to subprocess
cli.s2c.put ( {'op': 'data', 'data' : data} )
cli.sent_time = time.time()
cli.sent_data = data
cli.state = 1
if self.io_loop_sleep_time != 0:
io.process_messages(self.io_loop_sleep_time)
if self.on_tick() and all ([cli.state == 0 for cli in self.clis]):
#all subprocesses free and no more data available to process, ending loop
break
#gracefully terminating subprocesses
for cli in self.clis[:]:
cli.s2c.put ( {'op': 'close'} )
cli.sent_time = time.time()
while True:
for cli in self.clis[:]:
terminate_it = False
while not cli.c2s.empty():
obj = cli.c2s.get()
obj_op = obj['op']
if obj_op == 'finalized':
terminate_it = True
break
if (time.time() - cli.sent_time) > 30:
terminate_it = True
if terminate_it:
cli.state = 2
cli.kill()
if all ([cli.state == 2 for cli in self.clis]):
break
#finalizing host logic and return result
self.on_clients_finalized()
return self.get_result()
-16
View File
@@ -1,16 +0,0 @@
class ThisThreadGenerator(object):
def __init__(self, generator_func, user_param=None):
super().__init__()
self.generator_func = generator_func
self.user_param = user_param
self.initialized = False
def __iter__(self):
return self
def __next__(self):
if not self.initialized:
self.initialized = True
self.generator_func = self.generator_func(self.user_param)
return next(self.generator_func)
-5
View File
@@ -1,5 +0,0 @@
from .SubprocessorBase import Subprocessor
from .ThisThreadGenerator import ThisThreadGenerator
from .SubprocessGenerator import SubprocessGenerator
from .MPFunc import MPFunc
from .MPClassFuncOnDemand import MPClassFuncOnDemand
-1
View File
@@ -1 +0,0 @@
from .nn import nn
-17
View File
@@ -1,17 +0,0 @@
from core.leras import nn
class ArchiBase():
def __init__(self, *args, name=None, **kwargs):
self.name=name
#overridable
def flow(self, *args, **kwargs):
raise Exception("this archi does not support flow. Use model classes directly.")
#overridable
def get_weights(self):
pass
nn.ArchiBase = ArchiBase
-201
View File
@@ -1,201 +0,0 @@
from core.leras import nn
tf = nn.tf
class DeepFakeArchi(nn.ArchiBase):
"""
resolution
mod None - default
'quick'
"""
def __init__(self, resolution, mod=None, opts=None):
super().__init__()
if opts is None:
opts = ''
if mod is None:
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
super().__init__(*kwargs)
def on_build(self, *args, **kwargs ):
self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME')
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
return x
def get_out_ch(self):
return self.out_ch
class DownscaleBlock(nn.ModelBase):
def on_build(self, in_ch, ch, n_downscales, kernel_size):
self.downs = []
last_ch = in_ch
for i in range(n_downscales):
cur_ch = ch*( min(2**i, 8) )
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size) )
last_ch = self.downs[-1].get_out_ch()
def forward(self, inp):
x = inp
for down in self.downs:
x = down(x)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp + x, 0.2)
return x
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch):
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5)
def forward(self, inp):
return nn.flatten(self.down1(inp))
lowest_dense_res = resolution // (32 if 'd' in opts else 16)
class Inter(nn.ModelBase):
def __init__(self, in_ch, ae_ch, ae_out_ch, **kwargs):
self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
super().__init__(**kwargs)
def on_build(self):
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
if 'u' in opts:
self.dense_norm = nn.DenseNorm()
self.dense1 = nn.Dense( in_ch, ae_ch )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def forward(self, inp):
x = inp
if 'u' in opts:
x = self.dense_norm(x)
x = self.dense1(x)
x = self.dense2(x)
x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
x = self.upscale1(x)
return x
@staticmethod
def get_code_res():
return lowest_dense_res
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch, d_mask_ch ):
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
if 'd' in opts:
self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3)
self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME')
else:
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
def forward(self, inp):
z = inp
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
if 'd' in opts:
x0 = tf.nn.sigmoid(self.out_conv(x))
x0 = nn.upsample2d(x0)
x1 = tf.nn.sigmoid(self.out_conv1(x))
x1 = nn.upsample2d(x1)
x2 = tf.nn.sigmoid(self.out_conv2(x))
x2 = nn.upsample2d(x2)
x3 = tf.nn.sigmoid(self.out_conv3(x))
x3 = nn.upsample2d(x3)
if nn.data_format == "NHWC":
tile_cfg = ( 1, resolution // 2, resolution //2, 1)
else:
tile_cfg = ( 1, 1, resolution // 2, resolution //2 )
z0 = tf.concat ( ( tf.concat ( ( tf.ones ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ),
tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] )
z0 = tf.tile ( z0, tile_cfg )
z1 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.ones ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ),
tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] )
z1 = tf.tile ( z1, tile_cfg )
z2 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ),
tf.concat ( ( tf.ones ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] )
z2 = tf.tile ( z2, tile_cfg )
z3 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ),
tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.ones ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] )
z3 = tf.tile ( z3, tile_cfg )
x = x0*z0 + x1*z1 + x2*z2 + x3*z3
else:
x = tf.nn.sigmoid(self.out_conv(x))
m = self.upscalem0(z)
m = self.upscalem1(m)
m = self.upscalem2(m)
if 'd' in opts:
m = self.upscalem3(m)
m = tf.nn.sigmoid(self.out_convm(m))
return x, m
self.Encoder = Encoder
self.Inter = Inter
self.Decoder = Decoder
nn.DeepFakeArchi = DeepFakeArchi
-2
View File
@@ -1,2 +0,0 @@
from .ArchiBase import *
from .DeepFakeArchi import *
-207
View File
@@ -1,207 +0,0 @@
import sys
import ctypes
import os
class Device(object):
def __init__(self, index, name, total_mem, free_mem, cc=0):
self.index = index
self.name = name
self.cc = cc
self.total_mem = total_mem
self.total_mem_gb = total_mem / 1024**3
self.free_mem = free_mem
self.free_mem_gb = free_mem / 1024**3
def __str__(self):
return f"[{self.index}]:[{self.name}][{self.free_mem_gb:.3}/{self.total_mem_gb :.3}]"
class Devices(object):
all_devices = None
def __init__(self, devices):
self.devices = devices
def __len__(self):
return len(self.devices)
def __getitem__(self, key):
result = self.devices[key]
if isinstance(key, slice):
return Devices(result)
return result
def __iter__(self):
for device in self.devices:
yield device
def get_best_device(self):
result = None
idx_mem = 0
for device in self.devices:
mem = device.total_mem
if mem > idx_mem:
result = device
idx_mem = mem
return result
def get_worst_device(self):
result = None
idx_mem = sys.maxsize
for device in self.devices:
mem = device.total_mem
if mem < idx_mem:
result = device
idx_mem = mem
return result
def get_device_by_index(self, idx):
for device in self.devices:
if device.index == idx:
return device
return None
def get_devices_from_index_list(self, idx_list):
result = []
for device in self.devices:
if device.index in idx_list:
result += [device]
return Devices(result)
def get_equal_devices(self, device):
device_name = device.name
result = []
for device in self.devices:
if device.name == device_name:
result.append (device)
return Devices(result)
def get_devices_at_least_mem(self, totalmemsize_gb):
result = []
for device in self.devices:
if device.total_mem >= totalmemsize_gb*(1024**3):
result.append (device)
return Devices(result)
@staticmethod
def initialize_main_env():
os.environ['NN_DEVICES_INITIALIZED'] = '1'
os.environ['NN_DEVICES_COUNT'] = '0'
min_cc = int(os.environ.get("TF_MIN_REQ_CAP", 35))
libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll')
for libname in libnames:
try:
cuda = ctypes.CDLL(libname)
except:
continue
else:
break
else:
return Devices([])
nGpus = ctypes.c_int()
name = b' ' * 200
cc_major = ctypes.c_int()
cc_minor = ctypes.c_int()
freeMem = ctypes.c_size_t()
totalMem = ctypes.c_size_t()
result = ctypes.c_int()
device = ctypes.c_int()
context = ctypes.c_void_p()
error_str = ctypes.c_char_p()
devices = []
if cuda.cuInit(0) == 0 and \
cuda.cuDeviceGetCount(ctypes.byref(nGpus)) == 0:
for i in range(nGpus.value):
if cuda.cuDeviceGet(ctypes.byref(device), i) != 0 or \
cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device) != 0 or \
cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device) != 0:
continue
if cuda.cuCtxCreate_v2(ctypes.byref(context), 0, device) == 0:
if cuda.cuMemGetInfo_v2(ctypes.byref(freeMem), ctypes.byref(totalMem)) == 0:
cc = cc_major.value * 10 + cc_minor.value
if cc >= min_cc:
devices.append ( {'name' : name.split(b'\0', 1)[0].decode(),
'total_mem' : totalMem.value,
'free_mem' : freeMem.value,
'cc' : cc
})
cuda.cuCtxDetach(context)
os.environ['NN_DEVICES_COUNT'] = str(len(devices))
for i, device in enumerate(devices):
os.environ[f'NN_DEVICE_{i}_NAME'] = device['name']
os.environ[f'NN_DEVICE_{i}_TOTAL_MEM'] = str(device['total_mem'])
os.environ[f'NN_DEVICE_{i}_FREE_MEM'] = str(device['free_mem'])
os.environ[f'NN_DEVICE_{i}_CC'] = str(device['cc'])
@staticmethod
def getDevices():
if Devices.all_devices is None:
if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 1:
raise Exception("nn devices are not initialized. Run initialize_main_env() in main process.")
devices = []
for i in range ( int(os.environ['NN_DEVICES_COUNT']) ):
devices.append ( Device(index=i,
name=os.environ[f'NN_DEVICE_{i}_NAME'],
total_mem=int(os.environ[f'NN_DEVICE_{i}_TOTAL_MEM']),
free_mem=int(os.environ[f'NN_DEVICE_{i}_FREE_MEM']),
cc=int(os.environ[f'NN_DEVICE_{i}_CC']) ))
Devices.all_devices = Devices(devices)
return Devices.all_devices
"""
if Devices.all_devices is None:
min_cc = int(os.environ.get("TF_MIN_REQ_CAP", 35))
libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll')
for libname in libnames:
try:
cuda = ctypes.CDLL(libname)
except:
continue
else:
break
else:
return Devices([])
nGpus = ctypes.c_int()
name = b' ' * 200
cc_major = ctypes.c_int()
cc_minor = ctypes.c_int()
freeMem = ctypes.c_size_t()
totalMem = ctypes.c_size_t()
result = ctypes.c_int()
device = ctypes.c_int()
context = ctypes.c_void_p()
error_str = ctypes.c_char_p()
devices = []
if cuda.cuInit(0) == 0 and \
cuda.cuDeviceGetCount(ctypes.byref(nGpus)) == 0:
for i in range(nGpus.value):
if cuda.cuDeviceGet(ctypes.byref(device), i) != 0 or \
cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device) != 0 or \
cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device) != 0:
continue
if cuda.cuCtxCreate_v2(ctypes.byref(context), 0, device) == 0:
if cuda.cuMemGetInfo_v2(ctypes.byref(freeMem), ctypes.byref(totalMem)) == 0:
cc = cc_major.value * 10 + cc_minor.value
if cc >= min_cc:
devices.append ( Device(index=i,
name=name.split(b'\0', 1)[0].decode(),
total_mem=totalMem.value,
free_mem=freeMem.value,
cc=cc) )
cuda.cuCtxDetach(context)
Devices.all_devices = Devices(devices)
return Devices.all_devices
"""
-82
View File
@@ -1,82 +0,0 @@
import multiprocessing
from core.joblib import Subprocessor
import numpy as np
class CAInitializerSubprocessor(Subprocessor):
@staticmethod
def generate(shape, dtype=np.float32, eps_std=0.05):
"""
Super fast implementation of Convolution Aware Initialization for 4D shapes
Convolution Aware Initialization https://arxiv.org/abs/1702.06295
"""
if len(shape) != 4:
raise ValueError("only shape with rank 4 supported.")
row, column, stack_size, filters_size = shape
fan_in = stack_size * (row * column)
kernel_shape = (row, column)
kernel_fft_shape = np.fft.rfft2(np.zeros(kernel_shape)).shape
basis_size = np.prod(kernel_fft_shape)
if basis_size == 1:
x = np.random.normal( 0.0, eps_std, (filters_size, stack_size, basis_size) )
else:
nbb = stack_size // basis_size + 1
x = np.random.normal(0.0, 1.0, (filters_size, nbb, basis_size, basis_size))
x = x + np.transpose(x, (0,1,3,2) ) * (1-np.eye(basis_size))
u, _, v = np.linalg.svd(x)
x = np.transpose(u, (0,1,3,2) )
x = np.reshape(x, (filters_size, -1, basis_size) )
x = x[:,:stack_size,:]
x = np.reshape(x, ( (filters_size,stack_size,) + kernel_fft_shape ) )
x = np.fft.irfft2( x, kernel_shape ) \
+ np.random.normal(0, eps_std, (filters_size,stack_size,)+kernel_shape)
x = x * np.sqrt( (2/fan_in) / np.var(x) )
x = np.transpose( x, (2, 3, 1, 0) )
return x.astype(dtype)
class Cli(Subprocessor.Cli):
#override
def process_data(self, data):
idx, shape, dtype = data
weights = CAInitializerSubprocessor.generate (shape, dtype)
return idx, weights
#override
def __init__(self, data_list):
self.data_list = data_list
self.data_list_idxs = [*range(len(data_list))]
self.result = [None]*len(data_list)
super().__init__('CAInitializerSubprocessor', CAInitializerSubprocessor.Cli)
#override
def process_info_generator(self):
for i in range( min(multiprocessing.cpu_count(), len(self.data_list)) ):
yield 'CPU%d' % (i), {}, {}
#override
def get_data(self, host_dict):
if len (self.data_list_idxs) > 0:
idx = self.data_list_idxs.pop(0)
shape, dtype = self.data_list[idx]
return idx, shape, dtype
return None
#override
def on_data_return (self, host_dict, data):
self.data_list_idxs.insert(0, data)
#override
def on_result (self, host_dict, data, result):
idx, weights = result
self.result[idx] = weights
#override
def get_result(self):
return self.result
-20
View File
@@ -1,20 +0,0 @@
import numpy as np
from tensorflow.python.ops import init_ops
from core.leras import nn
tf = nn.tf
from .CA import CAInitializerSubprocessor
class initializers():
class ca (init_ops.Initializer):
def __call__(self, shape, dtype=None, partition_info=None):
return tf.zeros( shape, dtype=dtype, name="_cai_")
@staticmethod
def generate_batch( data_list, eps_std=0.05 ):
# list of (shape, np.dtype)
return CAInitializerSubprocessor (data_list).run()
nn.initializers = initializers
-56
View File
@@ -1,56 +0,0 @@
from core.leras import nn
tf = nn.tf
class AdaIN(nn.LayerBase):
"""
"""
def __init__(self, in_ch, mlp_ch, kernel_initializer=None, dtype=None, **kwargs):
self.in_ch = in_ch
self.mlp_ch = mlp_ch
self.kernel_initializer = kernel_initializer
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
kernel_initializer = tf.initializers.he_normal()
self.weight1 = tf.get_variable("weight1", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
self.bias1 = tf.get_variable("bias1", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
self.weight2 = tf.get_variable("weight2", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
self.bias2 = tf.get_variable("bias2", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
def get_weights(self):
return [self.weight1, self.bias1, self.weight2, self.bias2]
def forward(self, inputs):
x, mlp = inputs
gamma = tf.matmul(mlp, self.weight1)
gamma = tf.add(gamma, tf.reshape(self.bias1, (1,self.in_ch) ) )
beta = tf.matmul(mlp, self.weight2)
beta = tf.add(beta, tf.reshape(self.bias2, (1,self.in_ch) ) )
if nn.data_format == "NHWC":
shape = (-1,1,1,self.in_ch)
else:
shape = (-1,self.in_ch,1,1)
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
x = (x - x_mean) / x_std
x *= tf.reshape(gamma, shape)
x += tf.reshape(beta, shape)
return x
nn.AdaIN = AdaIN
-42
View File
@@ -1,42 +0,0 @@
from core.leras import nn
tf = nn.tf
class BatchNorm2D(nn.LayerBase):
"""
currently not for training
"""
def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs):
self.dim = dim
self.eps = eps
self.momentum = momentum
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable("weight", (self.dim,), dtype=self.dtype, initializer=tf.initializers.ones() )
self.bias = tf.get_variable("bias", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros() )
self.running_mean = tf.get_variable("running_mean", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
self.running_var = tf.get_variable("running_var", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
def get_weights(self):
return [self.weight, self.bias, self.running_mean, self.running_var]
def forward(self, x):
if nn.data_format == "NHWC":
shape = (1,1,1,self.dim)
else:
shape = (1,self.dim,1,1)
weight = tf.reshape ( self.weight , shape )
bias = tf.reshape ( self.bias , shape )
running_mean = tf.reshape ( self.running_mean, shape )
running_var = tf.reshape ( self.running_var , shape )
x = (x - running_mean) / tf.sqrt( running_var + self.eps )
x *= weight
x += bias
return x
nn.BatchNorm2D = BatchNorm2D
-50
View File
@@ -1,50 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
class BlurPool(nn.LayerBase):
def __init__(self, filt_size=3, stride=2, **kwargs ):
if nn.data_format == "NHWC":
self.strides = [1,stride,stride,1]
else:
self.strides = [1,1,stride,stride]
self.filt_size = filt_size
pad = [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ]
if nn.data_format == "NHWC":
self.padding = [ [0,0], pad, pad, [0,0] ]
else:
self.padding = [ [0,0], [0,0], pad, pad ]
if(self.filt_size==1):
a = np.array([1.,])
elif(self.filt_size==2):
a = np.array([1., 1.])
elif(self.filt_size==3):
a = np.array([1., 2., 1.])
elif(self.filt_size==4):
a = np.array([1., 3., 3., 1.])
elif(self.filt_size==5):
a = np.array([1., 4., 6., 4., 1.])
elif(self.filt_size==6):
a = np.array([1., 5., 10., 10., 5., 1.])
elif(self.filt_size==7):
a = np.array([1., 6., 15., 20., 15., 6., 1.])
a = a[:,None]*a[None,:]
a = a / np.sum(a)
a = a[:,:,None,None]
self.a = a
super().__init__(**kwargs)
def build_weights(self):
self.k = tf.constant (self.a, dtype=nn.floatx )
def forward(self, x):
k = tf.tile (self.k, (1,1,x.shape[nn.conv2d_ch_axis],1) )
x = tf.pad(x, self.padding )
x = tf.nn.depthwise_conv2d(x, k, self.strides, 'VALID', data_format=nn.data_format)
return x
nn.BlurPool = BlurPool
-112
View File
@@ -1,112 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
class Conv2D(nn.LayerBase):
"""
default kernel_initializer - CA
use_wscale bool enables equalized learning rate, if kernel_initializer is None, it will be forced to random_normal
"""
def __init__(self, in_ch, out_ch, kernel_size, strides=1, padding='SAME', dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
if not isinstance(strides, int):
raise ValueError ("strides must be an int type")
if not isinstance(dilations, int):
raise ValueError ("dilations must be an int type")
kernel_size = int(kernel_size)
if dtype is None:
dtype = nn.floatx
if isinstance(padding, str):
if padding == "SAME":
padding = ( (kernel_size - 1) * dilations + 1 ) // 2
elif padding == "VALID":
padding = 0
else:
raise ValueError ("Wrong padding type. Should be VALID SAME or INT or 4x INTs")
if isinstance(padding, int):
if padding != 0:
if nn.data_format == "NHWC":
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else:
padding = None
if nn.data_format == "NHWC":
strides = [1,strides,strides,1]
else:
strides = [1,1,strides,strides]
if nn.data_format == "NHWC":
dilations = [1,dilations,dilations,1]
else:
dilations = [1,1,dilations,dilations]
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.strides = strides
self.padding = padding
self.dilations = dilations
self.use_bias = use_bias
self.use_wscale = use_wscale
self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer
self.trainable = trainable
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in)
self.wscale = tf.constant(he_std, dtype=self.dtype )
if kernel_initializer is None:
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
if kernel_initializer is None:
kernel_initializer = nn.initializers.ca()
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
if self.use_bias:
bias_initializer = self.bias_initializer
if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self):
weights = [self.weight]
if self.use_bias:
weights += [self.bias]
return weights
def forward(self, x):
weight = self.weight
if self.use_wscale:
weight = weight * self.wscale
if self.padding is not None:
x = tf.pad (x, self.padding, mode='CONSTANT')
x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations, data_format=nn.data_format)
if self.use_bias:
if nn.data_format == "NHWC":
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
else:
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
x = tf.add(x, bias)
return x
def __str__(self):
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
return r
nn.Conv2D = Conv2D
-107
View File
@@ -1,107 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
class Conv2DTranspose(nn.LayerBase):
"""
use_wscale enables weight scale (equalized learning rate)
if kernel_initializer is None, it will be forced to random_normal
"""
def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
if not isinstance(strides, int):
raise ValueError ("strides must be an int type")
kernel_size = int(kernel_size)
if dtype is None:
dtype = nn.floatx
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.strides = strides
self.padding = padding
self.use_bias = use_bias
self.use_wscale = use_wscale
self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer
self.trainable = trainable
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, dtype=self.dtype )
if kernel_initializer is None:
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
if kernel_initializer is None:
kernel_initializer = nn.initializers.ca()
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
if self.use_bias:
bias_initializer = self.bias_initializer
if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self):
weights = [self.weight]
if self.use_bias:
weights += [self.bias]
return weights
def forward(self, x):
shape = x.shape
if nn.data_format == "NHWC":
h,w,c = shape[1], shape[2], shape[3]
output_shape = tf.stack ( (tf.shape(x)[0],
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
self.out_ch) )
strides = [1,self.strides,self.strides,1]
else:
c,h,w = shape[1], shape[2], shape[3]
output_shape = tf.stack ( (tf.shape(x)[0],
self.out_ch,
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
) )
strides = [1,1,self.strides,self.strides]
weight = self.weight
if self.use_wscale:
weight = weight * self.wscale
x = tf.nn.conv2d_transpose(x, weight, output_shape, strides, padding=self.padding, data_format=nn.data_format)
if self.use_bias:
if nn.data_format == "NHWC":
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
else:
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
x = tf.add(x, bias)
return x
def __str__(self):
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
return r
def deconv_length(self, dim_size, stride_size, kernel_size, padding):
assert padding in {'SAME', 'VALID', 'FULL'}
if dim_size is None:
return None
if padding == 'VALID':
dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0)
elif padding == 'FULL':
dim_size = dim_size * stride_size - (stride_size + kernel_size - 2)
elif padding == 'SAME':
dim_size = dim_size * stride_size
return dim_size
nn.Conv2DTranspose = Conv2DTranspose
-76
View File
@@ -1,76 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
class Dense(nn.LayerBase):
def __init__(self, in_ch, out_ch, use_bias=True, use_wscale=False, maxout_ch=0, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
"""
use_wscale enables weight scale (equalized learning rate)
if kernel_initializer is None, it will be forced to random_normal
maxout_ch https://link.springer.com/article/10.1186/s40537-019-0233-0
typical 2-4 if you want to enable DenseMaxout behaviour
"""
self.in_ch = in_ch
self.out_ch = out_ch
self.use_bias = use_bias
self.use_wscale = use_wscale
self.maxout_ch = maxout_ch
self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer
self.trainable = trainable
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
if self.maxout_ch > 1:
weight_shape = (self.in_ch,self.out_ch*self.maxout_ch)
else:
weight_shape = (self.in_ch,self.out_ch)
kernel_initializer = self.kernel_initializer
if self.use_wscale:
gain = 1.0
fan_in = np.prod( weight_shape[:-1] )
he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, dtype=self.dtype )
if kernel_initializer is None:
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
if kernel_initializer is None:
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
if self.use_bias:
bias_initializer = self.bias_initializer
if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self):
weights = [self.weight]
if self.use_bias:
weights += [self.bias]
return weights
def forward(self, x):
weight = self.weight
if self.use_wscale:
weight = weight * self.wscale
x = tf.matmul(x, weight)
if self.maxout_ch > 1:
x = tf.reshape (x, (-1, self.out_ch, self.maxout_ch) )
x = tf.reduce_max(x, axis=-1)
if self.use_bias:
x = tf.add(x, tf.reshape(self.bias, (1,self.out_ch) ) )
return x
nn.Dense = Dense
-16
View File
@@ -1,16 +0,0 @@
from core.leras import nn
tf = nn.tf
class DenseNorm(nn.LayerBase):
def __init__(self, dense=False, eps=1e-06, dtype=None, **kwargs):
self.dense = dense
if dtype is None:
dtype = nn.floatx
self.eps = tf.constant(eps, dtype=dtype, name="epsilon")
super().__init__(**kwargs)
def __call__(self, x):
return x * tf.rsqrt(tf.reduce_mean(tf.square(x), axis=-1, keepdims=True) + self.eps)
nn.DenseNorm = DenseNorm
-110
View File
@@ -1,110 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
class DepthwiseConv2D(nn.LayerBase):
"""
default kernel_initializer - CA
use_wscale bool enables equalized learning rate, if kernel_initializer is None, it will be forced to random_normal
"""
def __init__(self, in_ch, kernel_size, strides=1, padding='SAME', depth_multiplier=1, dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
if not isinstance(strides, int):
raise ValueError ("strides must be an int type")
if not isinstance(dilations, int):
raise ValueError ("dilations must be an int type")
kernel_size = int(kernel_size)
if dtype is None:
dtype = nn.floatx
if isinstance(padding, str):
if padding == "SAME":
padding = ( (kernel_size - 1) * dilations + 1 ) // 2
elif padding == "VALID":
padding = 0
else:
raise ValueError ("Wrong padding type. Should be VALID SAME or INT or 4x INTs")
if isinstance(padding, int):
if padding != 0:
if nn.data_format == "NHWC":
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else:
padding = None
if nn.data_format == "NHWC":
strides = [1,strides,strides,1]
else:
strides = [1,1,strides,strides]
if nn.data_format == "NHWC":
dilations = [1,dilations,dilations,1]
else:
dilations = [1,1,dilations,dilations]
self.in_ch = in_ch
self.depth_multiplier = depth_multiplier
self.kernel_size = kernel_size
self.strides = strides
self.padding = padding
self.dilations = dilations
self.use_bias = use_bias
self.use_wscale = use_wscale
self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer
self.trainable = trainable
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in)
self.wscale = tf.constant(he_std, dtype=self.dtype )
if kernel_initializer is None:
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
if kernel_initializer is None:
kernel_initializer = nn.initializers.ca()
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.depth_multiplier), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
if self.use_bias:
bias_initializer = self.bias_initializer
if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (self.in_ch*self.depth_multiplier,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self):
weights = [self.weight]
if self.use_bias:
weights += [self.bias]
return weights
def forward(self, x):
weight = self.weight
if self.use_wscale:
weight = weight * self.wscale
if self.padding is not None:
x = tf.pad (x, self.padding, mode='CONSTANT')
x = tf.nn.depthwise_conv2d(x, weight, self.strides, 'VALID', data_format=nn.data_format)
if self.use_bias:
if nn.data_format == "NHWC":
bias = tf.reshape (self.bias, (1,1,1,self.in_ch*self.depth_multiplier) )
else:
bias = tf.reshape (self.bias, (1,self.in_ch*self.depth_multiplier,1,1) )
x = tf.add(x, bias)
return x
def __str__(self):
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} depth_multiplier:{self.depth_multiplier} "
return r
nn.DepthwiseConv2D = DepthwiseConv2D
-38
View File
@@ -1,38 +0,0 @@
from core.leras import nn
tf = nn.tf
class FRNorm2D(nn.LayerBase):
"""
Tensorflow implementation of
Filter Response Normalization Layer: Eliminating Batch Dependence in theTraining of Deep Neural Networks
https://arxiv.org/pdf/1911.09737.pdf
"""
def __init__(self, in_ch, dtype=None, **kwargs):
self.in_ch = in_ch
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable("weight", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.ones() )
self.bias = tf.get_variable("bias", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
self.eps = tf.get_variable("eps", (1,), dtype=self.dtype, initializer=tf.initializers.constant(1e-6) )
def get_weights(self):
return [self.weight, self.bias, self.eps]
def forward(self, x):
if nn.data_format == "NHWC":
shape = (1,1,1,self.in_ch)
else:
shape = (1,self.in_ch,1,1)
weight = tf.reshape ( self.weight, shape )
bias = tf.reshape ( self.bias , shape )
nu2 = tf.reduce_mean(tf.square(x), axis=nn.conv2d_spatial_axes, keepdims=True)
x = x * ( 1.0/tf.sqrt(nu2 + tf.abs(self.eps) ) )
return x*weight + bias
nn.FRNorm2D = FRNorm2D
-40
View File
@@ -1,40 +0,0 @@
from core.leras import nn
tf = nn.tf
class InstanceNorm2D(nn.LayerBase):
def __init__(self, in_ch, dtype=None, **kwargs):
self.in_ch = in_ch
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
self.weight = tf.get_variable("weight", (self.in_ch,), dtype=self.dtype, initializer=kernel_initializer )
self.bias = tf.get_variable("bias", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
def get_weights(self):
return [self.weight, self.bias]
def forward(self, x):
if nn.data_format == "NHWC":
shape = (1,1,1,self.in_ch)
else:
shape = (1,self.in_ch,1,1)
weight = tf.reshape ( self.weight , shape )
bias = tf.reshape ( self.bias , shape )
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
x = (x - x_mean) / x_std
x *= weight
x += bias
return x
nn.InstanceNorm2D = InstanceNorm2D
-16
View File
@@ -1,16 +0,0 @@
from core.leras import nn
tf = nn.tf
class LayerBase(nn.Saveable):
#override
def build_weights(self):
pass
#override
def forward(self, *args, **kwargs):
pass
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
nn.LayerBase = LayerBase
-103
View File
@@ -1,103 +0,0 @@
import pickle
from pathlib import Path
from core import pathex
import numpy as np
from core.leras import nn
tf = nn.tf
class Saveable():
def __init__(self, name=None):
self.name = name
#override
def get_weights(self):
#return tf tensors that should be initialized/loaded/saved
return []
#override
def get_weights_np(self):
weights = self.get_weights()
if len(weights) == 0:
return []
return nn.tf_sess.run (weights)
def set_weights(self, new_weights):
weights = self.get_weights()
if len(weights) != len(new_weights):
raise ValueError ('len of lists mismatch')
tuples = []
for w, new_w in zip(weights, new_weights):
if len(w.shape) != new_w.shape:
new_w = new_w.reshape(w.shape)
tuples.append ( (w, new_w) )
nn.batch_set_value (tuples)
def save_weights(self, filename, force_dtype=None):
d = {}
weights = self.get_weights()
if self.name is None:
raise Exception("name must be defined.")
name = self.name
for w, w_val in zip(weights, nn.tf_sess.run (weights)):
w_name_split = w.name.split('/', 1)
if name != w_name_split[0]:
raise Exception("weight first name != Saveable.name")
if force_dtype is not None:
w_val = w_val.astype(force_dtype)
d[ w_name_split[1] ] = w_val
d_dumped = pickle.dumps (d, 4)
pathex.write_bytes_safe ( Path(filename), d_dumped )
def load_weights(self, filename):
"""
returns True if file exists
"""
filepath = Path(filename)
if filepath.exists():
result = True
d_dumped = filepath.read_bytes()
d = pickle.loads(d_dumped)
else:
return False
weights = self.get_weights()
if self.name is None:
raise Exception("name must be defined.")
tuples = []
for w in weights:
w_name_split = w.name.split('/')
if self.name != w_name_split[0]:
raise Exception("weight first name != Saveable.name")
sub_w_name = "/".join(w_name_split[1:])
w_val = d.get(sub_w_name, None)
if w_val is None:
#io.log_err(f"Weight {w.name} was not loaded from file {filename}")
tuples.append ( (w, w.initializer) )
else:
w_val = np.reshape( w_val, w.shape.as_list() )
tuples.append ( (w, w_val) )
nn.batch_set_value(tuples)
return True
def init_weights(self):
nn.init_weights(self.get_weights())
nn.Saveable = Saveable
-31
View File
@@ -1,31 +0,0 @@
from core.leras import nn
tf = nn.tf
class ScaleAdd(nn.LayerBase):
def __init__(self, ch, dtype=None, **kwargs):
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
self.ch = ch
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable("weight",(self.ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
def get_weights(self):
return [self.weight]
def forward(self, inputs):
if nn.data_format == "NHWC":
shape = (1,1,1,self.ch)
else:
shape = (1,self.ch,1,1)
weight = tf.reshape ( self.weight, shape )
x0, x1 = inputs
x = x0 + x1*weight
return x
nn.ScaleAdd = ScaleAdd
-33
View File
@@ -1,33 +0,0 @@
from core.leras import nn
tf = nn.tf
class TLU(nn.LayerBase):
"""
Tensorflow implementation of
Filter Response Normalization Layer: Eliminating Batch Dependence in theTraining of Deep Neural Networks
https://arxiv.org/pdf/1911.09737.pdf
"""
def __init__(self, in_ch, dtype=None, **kwargs):
self.in_ch = in_ch
if dtype is None:
dtype = nn.floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
self.tau = tf.get_variable("tau", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
def get_weights(self):
return [self.tau]
def forward(self, x):
if nn.data_format == "NHWC":
shape = (1,1,1,self.in_ch)
else:
shape = (1,self.in_ch,1,1)
tau = tf.reshape ( self.tau, shape )
return tf.math.maximum(x, tau)
nn.TLU = TLU
-16
View File
@@ -1,16 +0,0 @@
from .Saveable import *
from .LayerBase import *
from .Conv2D import *
from .Conv2DTranspose import *
from .DepthwiseConv2D import *
from .Dense import *
from .BlurPool import *
from .BatchNorm2D import *
from .FRNorm2D import *
from .TLU import *
from .ScaleAdd import *
from .DenseNorm import *
from .AdaIN import *
-22
View File
@@ -1,22 +0,0 @@
from core.leras import nn
tf = nn.tf
class CodeDiscriminator(nn.ModelBase):
def on_build(self, in_ch, code_res, ch=256, conv_kernel_initializer=None):
n_downscales = 1 + code_res // 8
self.convs = []
prev_ch = in_ch
for i in range(n_downscales):
cur_ch = ch * min( (2**i), 8 )
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
prev_ch = cur_ch
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
def forward(self, x):
for conv in self.convs:
x = tf.nn.leaky_relu( conv(x), 0.1 )
return self.out_conv(x)
nn.CodeDiscriminator = CodeDiscriminator
-253
View File
@@ -1,253 +0,0 @@
import types
import numpy as np
from core.interact import interact as io
from core.leras import nn
tf = nn.tf
class ModelBase(nn.Saveable):
def __init__(self, *args, name=None, **kwargs):
super().__init__(name=name)
self.layers = []
self.layers_by_name = {}
self.built = False
self.args = args
self.kwargs = kwargs
self.run_placeholders = None
def _build_sub(self, layer, name):
if isinstance (layer, list):
for i,sublayer in enumerate(layer):
self._build_sub(sublayer, f"{name}_{i}")
elif isinstance (layer, dict):
for subname in layer.keys():
sublayer = layer[subname]
self._build_sub(sublayer, f"{name}_{subname}")
elif isinstance (layer, nn.LayerBase) or \
isinstance (layer, ModelBase):
if layer.name is None:
layer.name = name
if isinstance (layer, nn.LayerBase):
with tf.variable_scope(layer.name):
layer.build_weights()
elif isinstance (layer, ModelBase):
layer.build()
self.layers.append (layer)
self.layers_by_name[layer.name] = layer
def xor_list(self, lst1, lst2):
return [value for value in lst1+lst2 if (value not in lst1) or (value not in lst2) ]
def build(self):
with tf.variable_scope(self.name):
current_vars = []
generator = None
while True:
if generator is None:
generator = self.on_build(*self.args, **self.kwargs)
if not isinstance(generator, types.GeneratorType):
generator = None
if generator is not None:
try:
next(generator)
except StopIteration:
generator = None
v = vars(self)
new_vars = self.xor_list (current_vars, list(v.keys()) )
for name in new_vars:
self._build_sub(v[name],name)
current_vars += new_vars
if generator is None:
break
self.built = True
#override
def get_weights(self):
if not self.built:
self.build()
weights = []
for layer in self.layers:
weights += layer.get_weights()
return weights
def get_layer_by_name(self, name):
return self.layers_by_name.get(name, None)
def get_layers(self):
if not self.built:
self.build()
layers = []
for layer in self.layers:
if isinstance (layer, nn.LayerBase):
layers.append(layer)
else:
layers += layer.get_layers()
return layers
#override
def on_build(self, *args, **kwargs):
"""
init model layers here
return 'yield' if build is not finished
therefore dependency models will be initialized
"""
pass
#override
def forward(self, *args, **kwargs):
#flow layers/models/tensors here
pass
def __call__(self, *args, **kwargs):
if not self.built:
self.build()
return self.forward(*args, **kwargs)
def compute_output_shape(self, shapes):
if not self.built:
self.build()
not_list = False
if not isinstance(shapes, list):
not_list = True
shapes = [shapes]
with tf.device('/CPU:0'):
# CPU tensors will not impact any performance, only slightly RAM "leakage"
phs = []
for dtype,sh in shapes:
phs += [ tf.placeholder(dtype, sh) ]
result = self.__call__(phs[0] if not_list else phs)
if not isinstance(result, list):
result = [result]
result_shapes = []
for t in result:
result_shapes += [ t.shape.as_list() ]
return result_shapes[0] if not_list else result_shapes
def compute_output_channels(self, shapes):
shape = self.compute_output_shape(shapes)
shape_len = len(shape)
if shape_len == 4:
if nn.data_format == "NCHW":
return shape[1]
return shape[-1]
def build_for_run(self, shapes_list):
if not isinstance(shapes_list, list):
raise ValueError("shapes_list must be a list.")
self.run_placeholders = []
for dtype,sh in shapes_list:
self.run_placeholders.append ( tf.placeholder(dtype, sh) )
self.run_output = self.__call__(self.run_placeholders)
def run (self, inputs):
if self.run_placeholders is None:
raise Exception ("Model didn't build for run.")
if len(inputs) != len(self.run_placeholders):
raise ValueError("len(inputs) != self.run_placeholders")
feed_dict = {}
for ph, inp in zip(self.run_placeholders, inputs):
feed_dict[ph] = inp
return nn.tf_sess.run ( self.run_output, feed_dict=feed_dict)
def summary(self):
layers = self.get_layers()
layers_names = []
layers_params = []
max_len_str = 0
max_len_param_str = 0
delim_str = "-"
total_params = 0
#Get layers names and str lenght for delim
for l in layers:
if len(str(l))>max_len_str:
max_len_str = len(str(l))
layers_names+=[str(l).capitalize()]
#Get params for each layer
layers_params = [ int(np.sum(np.prod(w.shape) for w in l.get_weights())) for l in layers ]
total_params = np.sum(layers_params)
#Get str lenght for delim
for p in layers_params:
if len(str(p))>max_len_param_str:
max_len_param_str=len(str(p))
#Set delim
for i in range(max_len_str+max_len_param_str+3):
delim_str += "-"
output = "\n"+delim_str+"\n"
#Format model name str
model_name_str = "| "+self.name.capitalize()
len_model_name_str = len(model_name_str)
for i in range(len(delim_str)-len_model_name_str):
model_name_str+= " " if i!=(len(delim_str)-len_model_name_str-2) else " |"
output += model_name_str +"\n"
output += delim_str +"\n"
#Format layers table
for i in range(len(layers_names)):
output += delim_str +"\n"
l_name = layers_names[i]
l_param = str(layers_params[i])
l_param_str = ""
if len(l_name)<=max_len_str:
for i in range(max_len_str - len(l_name)):
l_name+= " "
if len(l_param)<=max_len_param_str:
for i in range(max_len_param_str - len(l_param)):
l_param_str+= " "
l_param_str += l_param
output +="| "+l_name+"|"+l_param_str+"| \n"
output += delim_str +"\n"
#Format sum of params
total_params_str = "| Total params count: "+str(total_params)
len_total_params_str = len(total_params_str)
for i in range(len(delim_str)-len_total_params_str):
total_params_str+= " " if i!=(len(delim_str)-len_total_params_str-2) else " |"
output += total_params_str +"\n"
output += delim_str +"\n"
io.log_info(output)
nn.ModelBase = ModelBase
-191
View File
@@ -1,191 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
patch_discriminator_kernels = \
{ 1 : (512, [ [1,1] ]),
2 : (512, [ [2,1] ]),
3 : (512, [ [2,1], [2,1] ]),
4 : (512, [ [2,2], [2,2] ]),
5 : (512, [ [3,2], [2,2] ]),
6 : (512, [ [4,2], [2,2] ]),
7 : (512, [ [3,2], [3,2] ]),
8 : (512, [ [4,2], [3,2] ]),
9 : (512, [ [3,2], [4,2] ]),
10 : (512, [ [4,2], [4,2] ]),
11 : (512, [ [3,2], [3,2], [2,1] ]),
12 : (512, [ [4,2], [3,2], [2,1] ]),
13 : (512, [ [3,2], [4,2], [2,1] ]),
14 : (512, [ [4,2], [4,2], [2,1] ]),
15 : (512, [ [3,2], [3,2], [3,1] ]),
16 : (512, [ [4,2], [3,2], [3,1] ]),
17 : (512, [ [3,2], [4,2], [3,1] ]),
18 : (512, [ [4,2], [4,2], [3,1] ]),
19 : (512, [ [3,2], [3,2], [4,1] ]),
20 : (512, [ [4,2], [3,2], [4,1] ]),
21 : (512, [ [3,2], [4,2], [4,1] ]),
22 : (512, [ [4,2], [4,2], [4,1] ]),
23 : (256, [ [3,2], [3,2], [3,2], [2,1] ]),
24 : (256, [ [4,2], [3,2], [3,2], [2,1] ]),
25 : (256, [ [3,2], [4,2], [3,2], [2,1] ]),
26 : (256, [ [4,2], [4,2], [3,2], [2,1] ]),
27 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
28 : (256, [ [4,2], [3,2], [4,2], [2,1] ]),
29 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
30 : (256, [ [4,2], [4,2], [4,2], [2,1] ]),
31 : (256, [ [3,2], [3,2], [3,2], [3,1] ]),
32 : (256, [ [4,2], [3,2], [3,2], [3,1] ]),
33 : (256, [ [3,2], [4,2], [3,2], [3,1] ]),
34 : (256, [ [4,2], [4,2], [3,2], [3,1] ]),
35 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
36 : (256, [ [4,2], [3,2], [4,2], [3,1] ]),
37 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
38 : (256, [ [4,2], [4,2], [4,2], [3,1] ]),
39 : (256, [ [3,2], [3,2], [3,2], [4,1] ]),
40 : (256, [ [4,2], [3,2], [3,2], [4,1] ]),
41 : (256, [ [3,2], [4,2], [3,2], [4,1] ]),
42 : (256, [ [4,2], [4,2], [3,2], [4,1] ]),
43 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
44 : (256, [ [4,2], [3,2], [4,2], [4,1] ]),
45 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
46 : (256, [ [4,2], [4,2], [4,2], [4,1] ]),
}
class PatchDiscriminator(nn.ModelBase):
def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]
if base_ch is None:
base_ch = suggested_base_ch
prev_ch = in_ch
self.convs = []
for i, (kernel_size, strides) in enumerate(kernels_strides):
cur_ch = base_ch * min( (2**i), 8 )
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, padding='SAME', kernel_initializer=conv_kernel_initializer) )
prev_ch = cur_ch
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
def forward(self, x):
for conv in self.convs:
x = tf.nn.leaky_relu( conv(x), 0.1 )
return self.out_conv(x)
nn.PatchDiscriminator = PatchDiscriminator
class UNetPatchDiscriminator(nn.ModelBase):
"""
Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks"
"""
def calc_receptive_field_size(self, layers):
"""
result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html
"""
rf = 0
ts = 1
for i, (k, s) in enumerate(layers):
if i == 0:
rf = k
else:
rf += (k-1)*ts
ts *= s
return rf
def find_archi(self, target_patch_size, max_layers=6):
"""
Find the best configuration of layers using only 3x3 convs for target patch size
"""
s = {}
for layers_count in range(1,max_layers+1):
val = 1 << (layers_count-1)
while True:
val -= 1
layers = []
sum_st = 0
for i in range(layers_count-1):
st = 1 + (1 if val & (1 << i) !=0 else 0 )
layers.append ( [3, st ])
sum_st += st
layers.append ( [3, 2])
sum_st += 2
rf = self.calc_receptive_field_size(layers)
s_rf = s.get(rf, None)
if s_rf is None:
s[rf] = (layers_count, sum_st, layers)
else:
if layers_count < s_rf[0] or \
( layers_count == s_rf[0] and sum_st > s_rf[1] ):
s[rf] = (layers_count, sum_st, layers)
if val == 0:
break
x = sorted(list(s.keys()))
q=x[np.abs(np.array(x)-target_patch_size).argmin()]
return s[q][2]
def on_build(self, patch_size, in_ch):
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp + x, 0.2)
return x
prev_ch = in_ch
self.convs = []
self.res = []
self.upconvs = []
self.upres = []
layers = self.find_archi(patch_size)
base_ch = 16
level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) }
self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID')
for i, (kernel_size, strides) in enumerate(layers):
self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME') )
self.res.append ( ResidualBlock(level_chs[i]) )
self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME') )
self.upres.insert (0, ResidualBlock(level_chs[i-1]*2) )
self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID')
self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID')
self.center_conv = nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID')
def forward(self, x):
x = tf.nn.leaky_relu( self.in_conv(x), 0.1 )
encs = []
for conv, res in zip(self.convs, self.res):
encs.insert(0, x)
x = tf.nn.leaky_relu( conv(x), 0.1 )
x = res(x)
center_out, x = self.center_out(x), self.center_conv(x)
for i, (upconv, enc, upres) in enumerate(zip(self.upconvs, encs, self.upres)):
x = tf.nn.leaky_relu( upconv(x), 0.1 )
x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis)
x = upres(x)
return center_out, self.out_conv(x)
nn.UNetPatchDiscriminator = UNetPatchDiscriminator
-137
View File
@@ -1,137 +0,0 @@
from core.leras import nn
tf = nn.tf
class XSeg(nn.ModelBase):
def on_build (self, in_ch, base_ch, out_ch):
class ConvBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch):
self.conv = nn.Conv2D (in_ch, out_ch, kernel_size=3, padding='SAME')
self.frn = nn.FRNorm2D(out_ch)
self.tlu = nn.TLU(out_ch)
def forward(self, x):
x = self.conv(x)
x = self.frn(x)
x = self.tlu(x)
return x
class UpConvBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch):
self.conv = nn.Conv2DTranspose (in_ch, out_ch, kernel_size=3, padding='SAME')
self.frn = nn.FRNorm2D(out_ch)
self.tlu = nn.TLU(out_ch)
def forward(self, x):
x = self.conv(x)
x = self.frn(x)
x = self.tlu(x)
return x
self.conv01 = ConvBlock(in_ch, base_ch)
self.conv02 = ConvBlock(base_ch, base_ch)
self.bp0 = nn.BlurPool (filt_size=3)
self.conv11 = ConvBlock(base_ch, base_ch*2)
self.conv12 = ConvBlock(base_ch*2, base_ch*2)
self.bp1 = nn.BlurPool (filt_size=3)
self.conv21 = ConvBlock(base_ch*2, base_ch*4)
self.conv22 = ConvBlock(base_ch*4, base_ch*4)
self.conv23 = ConvBlock(base_ch*4, base_ch*4)
self.bp2 = nn.BlurPool (filt_size=3)
self.conv31 = ConvBlock(base_ch*4, base_ch*8)
self.conv32 = ConvBlock(base_ch*8, base_ch*8)
self.conv33 = ConvBlock(base_ch*8, base_ch*8)
self.bp3 = nn.BlurPool (filt_size=3)
self.conv41 = ConvBlock(base_ch*8, base_ch*8)
self.conv42 = ConvBlock(base_ch*8, base_ch*8)
self.conv43 = ConvBlock(base_ch*8, base_ch*8)
self.bp4 = nn.BlurPool (filt_size=3)
self.up4 = UpConvBlock (base_ch*8, base_ch*4)
self.uconv43 = ConvBlock(base_ch*12, base_ch*8)
self.uconv42 = ConvBlock(base_ch*8, base_ch*8)
self.uconv41 = ConvBlock(base_ch*8, base_ch*8)
self.up3 = UpConvBlock (base_ch*8, base_ch*4)
self.uconv33 = ConvBlock(base_ch*12, base_ch*8)
self.uconv32 = ConvBlock(base_ch*8, base_ch*8)
self.uconv31 = ConvBlock(base_ch*8, base_ch*8)
self.up2 = UpConvBlock (base_ch*8, base_ch*4)
self.uconv23 = ConvBlock(base_ch*8, base_ch*4)
self.uconv22 = ConvBlock(base_ch*4, base_ch*4)
self.uconv21 = ConvBlock(base_ch*4, base_ch*4)
self.up1 = UpConvBlock (base_ch*4, base_ch*2)
self.uconv12 = ConvBlock(base_ch*4, base_ch*2)
self.uconv11 = ConvBlock(base_ch*2, base_ch*2)
self.up0 = UpConvBlock (base_ch*2, base_ch)
self.uconv02 = ConvBlock(base_ch*2, base_ch)
self.uconv01 = ConvBlock(base_ch, base_ch)
self.out_conv = nn.Conv2D (base_ch, out_ch, kernel_size=3, padding='SAME')
self.conv_center = ConvBlock(base_ch*8, base_ch*8)
def forward(self, inp):
x = inp
x = self.conv01(x)
x = x0 = self.conv02(x)
x = self.bp0(x)
x = self.conv11(x)
x = x1 = self.conv12(x)
x = self.bp1(x)
x = self.conv21(x)
x = self.conv22(x)
x = x2 = self.conv23(x)
x = self.bp2(x)
x = self.conv31(x)
x = self.conv32(x)
x = x3 = self.conv33(x)
x = self.bp3(x)
x = self.conv41(x)
x = self.conv42(x)
x = x4 = self.conv43(x)
x = self.bp4(x)
x = self.conv_center(x)
x = self.up4(x)
x = self.uconv43(tf.concat([x,x4],axis=nn.conv2d_ch_axis))
x = self.uconv42(x)
x = self.uconv41(x)
x = self.up3(x)
x = self.uconv33(tf.concat([x,x3],axis=nn.conv2d_ch_axis))
x = self.uconv32(x)
x = self.uconv31(x)
x = self.up2(x)
x = self.uconv23(tf.concat([x,x2],axis=nn.conv2d_ch_axis))
x = self.uconv22(x)
x = self.uconv21(x)
x = self.up1(x)
x = self.uconv12(tf.concat([x,x1],axis=nn.conv2d_ch_axis))
x = self.uconv11(x)
x = self.up0(x)
x = self.uconv02(tf.concat([x,x0],axis=nn.conv2d_ch_axis))
x = self.uconv01(x)
logits = self.out_conv(x)
return logits, tf.nn.sigmoid(logits)
nn.XSeg = XSeg
-4
View File
@@ -1,4 +0,0 @@
from .ModelBase import *
from .PatchDiscriminator import *
from .CodeDiscriminator import *
from .XSeg import *
-299
View File
@@ -1,299 +0,0 @@
"""
Leras.
like lighter keras.
This is my lightweight neural network library written from scratch
based on pure tensorflow without keras.
Provides:
+ full freedom of tensorflow operations without keras model's restrictions
+ easy model operations like in PyTorch, but in graph mode (no eager execution)
+ convenient and understandable logic
Reasons why we cannot import tensorflow or any tensorflow.sub modules right here:
1) program is changing env variables based on DeviceConfig before import tensorflow
2) multiprocesses will import tensorflow every spawn
NCHW speed up training for 10-20%.
"""
import os
import sys
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
from pathlib import Path
import numpy as np
from core.interact import interact as io
from .device import Devices
class nn():
current_DeviceConfig = None
tf = None
tf_sess = None
tf_sess_config = None
tf_default_device = None
data_format = None
conv2d_ch_axis = None
conv2d_spatial_axes = None
floatx = None
@staticmethod
def initialize(device_config=None, floatx="float32", data_format="NHWC"):
if nn.tf is None:
if device_config is None:
device_config = nn.getCurrentDeviceConfig()
nn.setCurrentDeviceConfig(device_config)
# Manipulate environment variables before import tensorflow
if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
os.environ.pop('CUDA_VISIBLE_DEVICES')
first_run = False
if len(device_config.devices) != 0:
if sys.platform[0:3] == 'win':
# Windows specific env vars
if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ):
devices_str = "_" + device_config.devices[0].name.replace(' ','_')
else:
devices_str = ""
for device in device_config.devices:
devices_str += "_" + device.name.replace(' ','_')
compute_cache_path = Path(os.environ['APPDATA']) / 'NVIDIA' / ('ComputeCache' + devices_str)
if not compute_cache_path.exists():
first_run = True
os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path)
os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
if first_run:
io.log_info("Caching GPU kernels...")
import tensorflow as tf
nn.tf = tf
import logging
# Disable tensorflow warnings
logging.getLogger('tensorflow').setLevel(logging.ERROR)
# Initialize framework
import core.leras.ops
import core.leras.layers
import core.leras.initializers
import core.leras.optimizers
import core.leras.models
import core.leras.archis
# Configure tensorflow session-config
if len(device_config.devices) == 0:
nn.tf_default_device = "/CPU:0"
config = tf.ConfigProto(device_count={'GPU': 0})
else:
nn.tf_default_device = "/GPU:0"
config = tf.ConfigProto()
config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices])
config.gpu_options.force_gpu_compatible = True
config.gpu_options.allow_growth = True
nn.tf_sess_config = config
if nn.tf_sess is None:
nn.tf_sess = tf.Session(config=nn.tf_sess_config)
if floatx == "float32":
floatx = nn.tf.float32
elif floatx == "float16":
floatx = nn.tf.float16
else:
raise ValueError(f"unsupported floatx {floatx}")
nn.set_floatx(floatx)
nn.set_data_format(data_format)
@staticmethod
def initialize_main_env():
Devices.initialize_main_env()
@staticmethod
def set_floatx(tf_dtype):
"""
set default float type for all layers when dtype is None for them
"""
nn.floatx = tf_dtype
@staticmethod
def set_data_format(data_format):
if data_format != "NHWC" and data_format != "NCHW":
raise ValueError(f"unsupported data_format {data_format}")
nn.data_format = data_format
if data_format == "NHWC":
nn.conv2d_ch_axis = 3
nn.conv2d_spatial_axes = [1,2]
elif data_format == "NCHW":
nn.conv2d_ch_axis = 1
nn.conv2d_spatial_axes = [2,3]
@staticmethod
def get4Dshape ( w, h, c ):
"""
returns 4D shape based on current data_format
"""
if nn.data_format == "NHWC":
return (None,h,w,c)
else:
return (None,c,h,w)
@staticmethod
def to_data_format( x, to_data_format, from_data_format):
if to_data_format == from_data_format:
return x
if to_data_format == "NHWC":
return np.transpose(x, (0,2,3,1) )
elif to_data_format == "NCHW":
return np.transpose(x, (0,3,1,2) )
else:
raise ValueError(f"unsupported to_data_format {to_data_format}")
@staticmethod
def getCurrentDeviceConfig():
if nn.current_DeviceConfig is None:
nn.current_DeviceConfig = DeviceConfig.BestGPU()
return nn.current_DeviceConfig
@staticmethod
def setCurrentDeviceConfig(device_config):
nn.current_DeviceConfig = device_config
@staticmethod
def reset_session():
if nn.tf is not None:
if nn.tf_sess is not None:
nn.tf.reset_default_graph()
nn.tf_sess.close()
nn.tf_sess = nn.tf.Session(config=nn.tf_sess_config)
@staticmethod
def close_session():
if nn.tf_sess is not None:
nn.tf.reset_default_graph()
nn.tf_sess.close()
nn.tf_sess = None
@staticmethod
def get_current_device():
# Undocumented access to last tf.device(...)
objs = nn.tf.get_default_graph()._device_function_stack.peek_objs()
if len(objs) != 0:
return objs[0].display_name
return nn.tf_default_device
@staticmethod
def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False):
devices = Devices.getDevices()
if len(devices) == 0:
return []
all_devices_indexes = [device.index for device in devices]
if choose_only_one:
suggest_best_multi_gpu = False
suggest_all_gpu = False
if suggest_all_gpu:
best_device_indexes = all_devices_indexes
elif suggest_best_multi_gpu:
best_device_indexes = [device.index for device in devices.get_equal_devices(devices.get_best_device()) ]
else:
best_device_indexes = [ devices.get_best_device().index ]
best_device_indexes = ",".join([str(x) for x in best_device_indexes])
io.log_info ("")
if choose_only_one:
io.log_info ("Choose one GPU idx.")
else:
io.log_info ("Choose one or several GPU idxs (separated by comma).")
io.log_info ("")
if allow_cpu:
io.log_info ("[CPU] : CPU")
for device in devices:
io.log_info (f" [{device.index}] : {device.name}")
io.log_info ("")
while True:
try:
if choose_only_one:
choosed_idxs = io.input_str("Which GPU index to choose?", best_device_indexes)
else:
choosed_idxs = io.input_str("Which GPU indexes to choose?", best_device_indexes)
if allow_cpu and choosed_idxs.lower() == "cpu":
choosed_idxs = []
break
choosed_idxs = [ int(x) for x in choosed_idxs.split(',') ]
if choose_only_one:
if len(choosed_idxs) == 1:
break
else:
if all( [idx in all_devices_indexes for idx in choosed_idxs] ):
break
except:
pass
io.log_info ("")
return choosed_idxs
class DeviceConfig():
@staticmethod
def ask_choose_device(*args, **kwargs):
return nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(*args,**kwargs) )
def __init__ (self, devices=None):
devices = devices or []
if not isinstance(devices, Devices):
devices = Devices(devices)
self.devices = devices
self.cpu_only = len(devices) == 0
@staticmethod
def BestGPU():
devices = Devices.getDevices()
if len(devices) == 0:
return nn.DeviceConfig.CPU()
return nn.DeviceConfig([devices.get_best_device()])
@staticmethod
def WorstGPU():
devices = Devices.getDevices()
if len(devices) == 0:
return nn.DeviceConfig.CPU()
return nn.DeviceConfig([devices.get_worst_device()])
@staticmethod
def GPUIndexes(indexes):
if len(indexes) != 0:
devices = Devices.getDevices().get_devices_from_index_list(indexes)
else:
devices = []
return nn.DeviceConfig(devices)
@staticmethod
def CPU():
return nn.DeviceConfig([])
-378
View File
@@ -1,378 +0,0 @@
import numpy as np
from core.leras import nn
tf = nn.tf
from tensorflow.python.ops import array_ops, random_ops, math_ops, sparse_ops, gradients
from tensorflow.python.framework import sparse_tensor
def tf_get_value(tensor):
return nn.tf_sess.run (tensor)
nn.tf_get_value = tf_get_value
def batch_set_value(tuples):
if len(tuples) != 0:
with nn.tf.device('/CPU:0'):
assign_ops = []
feed_dict = {}
for x, value in tuples:
if isinstance(value, nn.tf.Operation) or \
isinstance(value, nn.tf.Variable):
assign_ops.append(value)
else:
value = np.asarray(value, dtype=x.dtype.as_numpy_dtype)
assign_placeholder = nn.tf.placeholder( x.dtype.base_dtype, shape=[None]*value.ndim )
assign_op = nn.tf.assign (x, assign_placeholder )
assign_ops.append(assign_op)
feed_dict[assign_placeholder] = value
nn.tf_sess.run(assign_ops, feed_dict=feed_dict)
nn.batch_set_value = batch_set_value
def init_weights(weights):
ops = []
ca_tuples_w = []
ca_tuples = []
for w in weights:
initializer = w.initializer
for input in initializer.inputs:
if "_cai_" in input.name:
ca_tuples_w.append (w)
ca_tuples.append ( (w.shape.as_list(), w.dtype.as_numpy_dtype) )
break
else:
ops.append (initializer)
if len(ops) != 0:
nn.tf_sess.run (ops)
if len(ca_tuples) != 0:
nn.batch_set_value( [*zip(ca_tuples_w, nn.initializers.ca.generate_batch (ca_tuples))] )
nn.init_weights = init_weights
def tf_gradients ( loss, vars ):
grads = gradients.gradients(loss, vars, colocate_gradients_with_ops=True )
gv = [*zip(grads,vars)]
for g,v in gv:
if g is None:
raise Exception(f"Variable {v.name} is declared as trainable, but no tensors flow through it.")
return gv
nn.gradients = tf_gradients
def average_gv_list(grad_var_list, tf_device_string=None):
if len(grad_var_list) == 1:
return grad_var_list[0]
e = tf.device(tf_device_string) if tf_device_string is not None else None
if e is not None: e.__enter__()
result = []
for i, (gv) in enumerate(grad_var_list):
for j,(g,v) in enumerate(gv):
g = tf.expand_dims(g, 0)
if i == 0:
result += [ [[g], v] ]
else:
result[j][0] += [g]
for i,(gs,v) in enumerate(result):
result[i] = ( tf.reduce_mean( tf.concat (gs, 0), 0 ), v )
if e is not None: e.__exit__(None,None,None)
return result
nn.average_gv_list = average_gv_list
def average_tensor_list(tensors_list, tf_device_string=None):
if len(tensors_list) == 1:
return tensors_list[0]
e = tf.device(tf_device_string) if tf_device_string is not None else None
if e is not None: e.__enter__()
result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0)
if e is not None: e.__exit__(None,None,None)
return result
nn.average_tensor_list = average_tensor_list
def concat (tensors_list, axis):
"""
Better version.
"""
if len(tensors_list) == 1:
return tensors_list[0]
return tf.concat(tensors_list, axis)
nn.concat = concat
def gelu(x):
cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
return x * cdf
nn.gelu = gelu
def upsample2d(x, size=2):
if nn.data_format == "NCHW":
b,c,h,w = x.shape.as_list()
x = tf.reshape (x, (-1,c,h,1,w,1) )
x = tf.tile(x, (1,1,1,size,1,size) )
x = tf.reshape (x, (-1,c,h*size,w*size) )
return x
else:
return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) )
nn.upsample2d = upsample2d
def resize2d_bilinear(x, size=2):
h = x.shape[nn.conv2d_spatial_axes[0]].value
w = x.shape[nn.conv2d_spatial_axes[1]].value
if nn.data_format == "NCHW":
x = tf.transpose(x, (0,2,3,1))
if size > 0:
new_size = (h*size,w*size)
else:
new_size = (h//-size,w//-size)
x = tf.image.resize(x, new_size, method=tf.image.ResizeMethod.BILINEAR)
if nn.data_format == "NCHW":
x = tf.transpose(x, (0,3,1,2))
return x
nn.resize2d_bilinear = resize2d_bilinear
def resize2d_nearest(x, size=2):
if size in [-1,0,1]:
return x
if size > 0:
raise Exception("")
else:
if nn.data_format == "NCHW":
x = x[:,:,::-size,::-size]
else:
x = x[:,::-size,::-size,:]
return x
h = x.shape[nn.conv2d_spatial_axes[0]].value
w = x.shape[nn.conv2d_spatial_axes[1]].value
if nn.data_format == "NCHW":
x = tf.transpose(x, (0,2,3,1))
if size > 0:
new_size = (h*size,w*size)
else:
new_size = (h//-size,w//-size)
x = tf.image.resize(x, new_size, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
if nn.data_format == "NCHW":
x = tf.transpose(x, (0,3,1,2))
return x
nn.resize2d_nearest = resize2d_nearest
def flatten(x):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
x = tf.transpose(x, (0,3,1,2) )
return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
nn.flatten = flatten
def max_pool(x, kernel_size=2, strides=2):
if nn.data_format == "NHWC":
return tf.nn.max_pool(x, [1,kernel_size,kernel_size,1], [1,strides,strides,1], 'SAME', data_format=nn.data_format)
else:
return tf.nn.max_pool(x, [1,1,kernel_size,kernel_size], [1,1,strides,strides], 'SAME', data_format=nn.data_format)
nn.max_pool = max_pool
def reshape_4D(x, w,h,c):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
x = tf.reshape (x, (-1,c,h,w))
x = tf.transpose(x, (0,2,3,1) )
return x
else:
return tf.reshape (x, (-1,c,h,w))
nn.reshape_4D = reshape_4D
def random_binomial(shape, p=0.0, dtype=None, seed=None):
if dtype is None:
dtype=tf.float32
if seed is None:
seed = np.random.randint(10e6)
return array_ops.where(
random_ops.random_uniform(shape, dtype=tf.float16, seed=seed) < p,
array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype))
nn.random_binomial = random_binomial
def gaussian_blur(input, radius=2.0):
def gaussian(x, mu, sigma):
return np.exp(-(float(x) - float(mu)) ** 2 / (2 * sigma ** 2))
def make_kernel(sigma):
kernel_size = max(3, int(2 * 2 * sigma + 1))
mean = np.floor(0.5 * kernel_size)
kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
kernel = np_kernel / np.sum(np_kernel)
return kernel, kernel_size
gauss_kernel, kernel_size = make_kernel(radius)
padding = kernel_size//2
if padding != 0:
if nn.data_format == "NHWC":
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else:
padding = None
gauss_kernel = gauss_kernel[:,:,None,None]
x = input
k = tf.tile (gauss_kernel, (1,1,x.shape[nn.conv2d_ch_axis],1) )
x = tf.pad(x, padding )
x = tf.nn.depthwise_conv2d(x, k, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
return x
nn.gaussian_blur = gaussian_blur
def style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
def sd(content, style, loss_weight):
content_nc = content.shape[ nn.conv2d_ch_axis ]
style_nc = style.shape[nn.conv2d_ch_axis]
if content_nc != style_nc:
raise Exception("style_loss() content_nc != style_nc")
c_mean, c_var = tf.nn.moments(content, axes=nn.conv2d_spatial_axes, keep_dims=True)
s_mean, s_var = tf.nn.moments(style, axes=nn.conv2d_spatial_axes, keep_dims=True)
c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5)
mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3])
std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3])
return (mean_loss + std_loss) * ( loss_weight / content_nc.value )
if gaussian_blur_radius > 0.0:
target = gaussian_blur(target, gaussian_blur_radius)
style = gaussian_blur(style, gaussian_blur_radius)
return sd( target, style, loss_weight=loss_weight )
nn.style_loss = style_loss
def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
if img1.dtype != img2.dtype:
raise ValueError("img1.dtype != img2.dtype")
not_float32 = img1.dtype != tf.float32
if not_float32:
img_dtype = img1.dtype
img1 = tf.cast(img1, tf.float32)
img2 = tf.cast(img2, tf.float32)
filter_size = max(1, filter_size)
kernel = np.arange(0, filter_size, dtype=np.float32)
kernel -= (filter_size - 1 ) / 2.0
kernel = kernel**2
kernel *= ( -0.5 / (filter_sigma**2) )
kernel = np.reshape (kernel, (1,-1)) + np.reshape(kernel, (-1,1) )
kernel = tf.constant ( np.reshape (kernel, (1,-1)), dtype=tf.float32 )
kernel = tf.nn.softmax(kernel)
kernel = tf.reshape (kernel, (filter_size, filter_size, 1, 1))
kernel = tf.tile (kernel, (1,1, img1.shape[ nn.conv2d_ch_axis ] ,1))
def reducer(x):
return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
c1 = (k1 * max_val) ** 2
c2 = (k2 * max_val) ** 2
mean0 = reducer(img1)
mean1 = reducer(img2)
num0 = mean0 * mean1 * 2.0
den0 = tf.square(mean0) + tf.square(mean1)
luminance = (num0 + c1) / (den0 + c1)
num1 = reducer(img1 * img2) * 2.0
den1 = reducer(tf.square(img1) + tf.square(img2))
c2 *= 1.0 #compensation factor
cs = (num1 - num0 + c2) / (den1 - den0 + c2)
ssim_val = tf.reduce_mean(luminance * cs, axis=nn.conv2d_spatial_axes )
dssim = (1.0 - ssim_val ) / 2.0
if not_float32:
dssim = tf.cast(dssim, img_dtype)
return dssim
nn.dssim = dssim
def space_to_depth(x, size):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
b,h,w,c = x.shape.as_list()
oh, ow = h // size, w // size
x = tf.reshape(x, (-1, size, oh, size, ow, c))
x = tf.transpose(x, (0, 2, 4, 1, 3, 5))
x = tf.reshape(x, (-1, oh, ow, size* size* c ))
return x
else:
return tf.space_to_depth(x, size, data_format=nn.data_format)
nn.space_to_depth = space_to_depth
def depth_to_space(x, size):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
b,h,w,c = x.shape.as_list()
oh, ow = h * size, w * size
oc = c // (size * size)
x = tf.reshape(x, (-1, h, w, size, size, oc, ) )
x = tf.transpose(x, (0, 1, 3, 2, 4, 5))
x = tf.reshape(x, (-1, oh, ow, oc, ))
return x
else:
return tf.depth_to_space(x, size, data_format=nn.data_format)
nn.depth_to_space = depth_to_space
def rgb_to_lab(srgb):
srgb_pixels = tf.reshape(srgb, [-1, 3])
linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)
exponential_mask = tf.cast(srgb_pixels > 0.04045, dtype=tf.float32)
rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (((srgb_pixels + 0.055) / 1.055) ** 2.4) * exponential_mask
rgb_to_xyz = tf.constant([
# X Y Z
[0.412453, 0.212671, 0.019334], # R
[0.357580, 0.715160, 0.119193], # G
[0.180423, 0.072169, 0.950227], # B
])
xyz_pixels = tf.matmul(rgb_pixels, rgb_to_xyz)
xyz_normalized_pixels = tf.multiply(xyz_pixels, [1/0.950456, 1.0, 1/1.088754])
epsilon = 6/29
linear_mask = tf.cast(xyz_normalized_pixels <= (epsilon**3), dtype=tf.float32)
exponential_mask = tf.cast(xyz_normalized_pixels > (epsilon**3), dtype=tf.float32)
fxfyfz_pixels = (xyz_normalized_pixels / (3 * epsilon**2) + 4/29) * linear_mask + (xyz_normalized_pixels ** (1/3)) * exponential_mask
fxfyfz_to_lab = tf.constant([
# l a b
[ 0.0, 500.0, 0.0], # fx
[116.0, -500.0, 200.0], # fy
[ 0.0, 0.0, -200.0], # fz
])
lab_pixels = tf.matmul(fxfyfz_pixels, fxfyfz_to_lab) + tf.constant([-16.0, 0.0, 0.0])
return tf.reshape(lab_pixels, tf.shape(srgb))
nn.rgb_to_lab = rgb_to_lab
"""
def tf_suppress_lower_mean(t, eps=0.00001):
if t.shape.ndims != 1:
raise ValueError("tf_suppress_lower_mean: t rank must be 1")
t_mean_eps = tf.reduce_mean(t) - eps
q = tf.clip_by_value(t, t_mean_eps, tf.reduce_max(t) )
q = tf.clip_by_value(q-t_mean_eps, 0, eps)
q = q * (t/eps)
return q
"""

Some files were not shown because too many files have changed in this diff Show More