Commit fba8bde8 authored by bailuo's avatar bailuo
Browse files

update

parents
Pipeline #1808 failed with stages
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
class ImageMask(gr.components.Image):
"""
Sets: source="canvas", tool="sketch"
"""
is_template = True
def __init__(self, **kwargs):
super().__init__(source="upload",
tool="sketch",
interactive=False,
**kwargs)
def preprocess(self, x):
if x is None:
return x
if self.tool == "sketch" and self.source in ["upload", "webcam"
] and type(x) != dict:
decode_image = gr.processing_utils.decode_base64_to_image(x)
width, height = decode_image.size
mask = np.ones((height, width, 4), dtype=np.uint8)
mask[..., -1] = 255
mask = self.postprocess(mask)
x = {'image': x, 'mask': mask}
return super().preprocess(x)
def get_valid_mask(mask: np.ndarray):
"""Convert mask from gr.Image(0 to 255, RGBA) to binary mask.
"""
if mask.ndim == 3:
mask_pil = Image.fromarray(mask).convert('L')
mask = np.array(mask_pil)
if mask.max() == 255:
mask = mask / 255
return mask
def draw_points_on_image(image,
points,
curr_point=None,
highlight_all=True,
radius_scale=0.01):
overlay_rgba = Image.new("RGBA", image.size, 0)
overlay_draw = ImageDraw.Draw(overlay_rgba)
for point_key, point in points.items():
if ((curr_point is not None and curr_point == point_key)
or highlight_all):
p_color = (255, 0, 0)
t_color = (0, 0, 255)
else:
p_color = (255, 0, 0, 35)
t_color = (0, 0, 255, 35)
rad_draw = int(image.size[0] * radius_scale)
p_start = point.get("start_temp", point["start"])
p_target = point["target"]
if p_start is not None and p_target is not None:
p_draw = int(p_start[0]), int(p_start[1])
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.line(
(p_draw[0], p_draw[1], t_draw[0], t_draw[1]),
fill=(255, 255, 0),
width=2,
)
if p_start is not None:
p_draw = int(p_start[0]), int(p_start[1])
overlay_draw.ellipse(
(
p_draw[0] - rad_draw,
p_draw[1] - rad_draw,
p_draw[0] + rad_draw,
p_draw[1] + rad_draw,
),
fill=p_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(p_draw, "p", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(p_draw, "p", align="center", fill=(0, 0, 0))
if p_target is not None:
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.ellipse(
(
t_draw[0] - rad_draw,
t_draw[1] - rad_draw,
t_draw[0] + rad_draw,
t_draw[1] + rad_draw,
),
fill=t_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(t_draw, "t", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(t_draw, "t", align="center", fill=(0, 0, 0))
return Image.alpha_composite(image.convert("RGBA"),
overlay_rgba).convert("RGB")
def draw_mask_on_image(image, mask):
im_mask = np.uint8(mask * 255)
im_mask_rgba = np.concatenate(
(
np.tile(im_mask[..., None], [1, 1, 3]),
45 * np.ones(
(im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8),
),
axis=-1,
)
im_mask_rgba = Image.fromarray(im_mask_rgba).convert("RGBA")
return Image.alpha_composite(image.convert("RGBA"),
im_mask_rgba).convert("RGB")
def on_change_single_global_state(keys,
value,
global_state,
map_transform=None):
if map_transform is not None:
value = map_transform(value)
curr_state = global_state
if isinstance(keys, str):
last_key = keys
else:
for k in keys[:-1]:
curr_state = curr_state[k]
last_key = keys[-1]
curr_state[last_key] = value
return global_state
def get_latest_points_pair(points_dict):
if not points_dict:
return None
point_idx = list(points_dict.keys())
latest_point_idx = max(point_idx)
return latest_point_idx
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
from gradio.events import Dependency
class ImageMask(gr.components.Image):
"""
Sets: source="canvas", tool="sketch"
"""
is_template = True
def __init__(self, **kwargs):
super().__init__(source="upload",
tool="sketch",
interactive=False,
**kwargs)
def preprocess(self, x):
if x is None:
return x
if self.tool == "sketch" and self.source in ["upload", "webcam"
] and type(x) != dict:
decode_image = gr.processing_utils.decode_base64_to_image(x)
width, height = decode_image.size
mask = np.ones((height, width, 4), dtype=np.uint8)
mask[..., -1] = 255
mask = self.postprocess(mask)
x = {'image': x, 'mask': mask}
return super().preprocess(x)
from typing import Callable, Literal, Sequence, Any, TYPE_CHECKING
from gradio.blocks import Block
if TYPE_CHECKING:
from gradio.components import Timer
def get_valid_mask(mask: np.ndarray):
"""Convert mask from gr.Image(0 to 255, RGBA) to binary mask.
"""
if mask.ndim == 3:
mask_pil = Image.fromarray(mask).convert('L')
mask = np.array(mask_pil)
if mask.max() == 255:
mask = mask / 255
return mask
def draw_points_on_image(image,
points,
curr_point=None,
highlight_all=True,
radius_scale=0.01):
overlay_rgba = Image.new("RGBA", image.size, 0)
overlay_draw = ImageDraw.Draw(overlay_rgba)
for point_key, point in points.items():
if ((curr_point is not None and curr_point == point_key)
or highlight_all):
p_color = (255, 0, 0)
t_color = (0, 0, 255)
else:
p_color = (255, 0, 0, 35)
t_color = (0, 0, 255, 35)
rad_draw = int(image.size[0] * radius_scale)
p_start = point.get("start_temp", point["start"])
p_target = point["target"]
if p_start is not None and p_target is not None:
p_draw = int(p_start[0]), int(p_start[1])
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.line(
(p_draw[0], p_draw[1], t_draw[0], t_draw[1]),
fill=(255, 255, 0),
width=2,
)
if p_start is not None:
p_draw = int(p_start[0]), int(p_start[1])
overlay_draw.ellipse(
(
p_draw[0] - rad_draw,
p_draw[1] - rad_draw,
p_draw[0] + rad_draw,
p_draw[1] + rad_draw,
),
fill=p_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(p_draw, "p", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(p_draw, "p", align="center", fill=(0, 0, 0))
if p_target is not None:
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.ellipse(
(
t_draw[0] - rad_draw,
t_draw[1] - rad_draw,
t_draw[0] + rad_draw,
t_draw[1] + rad_draw,
),
fill=t_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(t_draw, "t", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(t_draw, "t", align="center", fill=(0, 0, 0))
return Image.alpha_composite(image.convert("RGBA"),
overlay_rgba).convert("RGB")
def draw_mask_on_image(image, mask):
im_mask = np.uint8(mask * 255)
im_mask_rgba = np.concatenate(
(
np.tile(im_mask[..., None], [1, 1, 3]),
45 * np.ones(
(im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8),
),
axis=-1,
)
im_mask_rgba = Image.fromarray(im_mask_rgba).convert("RGBA")
return Image.alpha_composite(image.convert("RGBA"),
im_mask_rgba).convert("RGB")
def on_change_single_global_state(keys,
value,
global_state,
map_transform=None):
if map_transform is not None:
value = map_transform(value)
curr_state = global_state
if isinstance(keys, str):
last_key = keys
else:
for k in keys[:-1]:
curr_state = curr_state[k]
last_key = keys[-1]
curr_state[last_key] = value
return global_state
def get_latest_points_pair(points_dict):
if not points_dict:
return None
point_idx = list(points_dict.keys())
latest_point_idx = max(point_idx)
return latest_point_idx
\ No newline at end of file
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
# empty
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import math
import os
import functools
import contextlib
import numpy as np
import OpenGL.GL as gl
import OpenGL.GL.ARB.texture_float
import dnnlib
#----------------------------------------------------------------------------
def init_egl():
assert os.environ['PYOPENGL_PLATFORM'] == 'egl' # Must be set before importing OpenGL.
import OpenGL.EGL as egl
import ctypes
# Initialize EGL.
display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY)
assert display != egl.EGL_NO_DISPLAY
major = ctypes.c_int32()
minor = ctypes.c_int32()
ok = egl.eglInitialize(display, major, minor)
assert ok
assert major.value * 10 + minor.value >= 14
# Choose config.
config_attribs = [
egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT,
egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT,
egl.EGL_NONE
]
configs = (ctypes.c_int32 * 1)()
num_configs = ctypes.c_int32()
ok = egl.eglChooseConfig(display, config_attribs, configs, 1, num_configs)
assert ok
assert num_configs.value == 1
config = configs[0]
# Create dummy pbuffer surface.
surface_attribs = [
egl.EGL_WIDTH, 1,
egl.EGL_HEIGHT, 1,
egl.EGL_NONE
]
surface = egl.eglCreatePbufferSurface(display, config, surface_attribs)
assert surface != egl.EGL_NO_SURFACE
# Setup GL context.
ok = egl.eglBindAPI(egl.EGL_OPENGL_API)
assert ok
context = egl.eglCreateContext(display, config, egl.EGL_NO_CONTEXT, None)
assert context != egl.EGL_NO_CONTEXT
ok = egl.eglMakeCurrent(display, surface, surface, context)
assert ok
#----------------------------------------------------------------------------
_texture_formats = {
('uint8', 1): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE, internalformat=gl.GL_LUMINANCE8),
('uint8', 2): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE_ALPHA, internalformat=gl.GL_LUMINANCE8_ALPHA8),
('uint8', 3): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGB, internalformat=gl.GL_RGB8),
('uint8', 4): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGBA, internalformat=gl.GL_RGBA8),
('float32', 1): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE32F_ARB),
('float32', 2): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE_ALPHA, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE_ALPHA32F_ARB),
('float32', 3): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGB, internalformat=gl.GL_RGB32F),
('float32', 4): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGBA, internalformat=gl.GL_RGBA32F),
}
def get_texture_format(dtype, channels):
return _texture_formats[(np.dtype(dtype).name, int(channels))]
#----------------------------------------------------------------------------
def prepare_texture_data(image):
image = np.asarray(image)
if image.ndim == 2:
image = image[:, :, np.newaxis]
if image.dtype.name == 'float64':
image = image.astype('float32')
return image
#----------------------------------------------------------------------------
def draw_pixels(image, *, pos=0, zoom=1, align=0, rint=True):
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
align = np.broadcast_to(np.asarray(align, dtype='float32'), [2])
image = prepare_texture_data(image)
height, width, channels = image.shape
size = zoom * [width, height]
pos = pos - size * align
if rint:
pos = np.rint(pos)
fmt = get_texture_format(image.dtype, channels)
gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_PIXEL_MODE_BIT)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glRasterPos2f(pos[0], pos[1])
gl.glPixelZoom(zoom[0], -zoom[1])
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
gl.glDrawPixels(width, height, fmt.format, fmt.type, image)
gl.glPopClientAttrib()
gl.glPopAttrib()
#----------------------------------------------------------------------------
def read_pixels(width, height, *, pos=0, dtype='uint8', channels=3):
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
dtype = np.dtype(dtype)
fmt = get_texture_format(dtype, channels)
image = np.empty([height, width, channels], dtype=dtype)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)
gl.glReadPixels(int(np.round(pos[0])), int(np.round(pos[1])), width, height, fmt.format, fmt.type, image)
gl.glPopClientAttrib()
return np.flipud(image)
#----------------------------------------------------------------------------
class Texture:
def __init__(self, *, image=None, width=None, height=None, channels=None, dtype=None, bilinear=True, mipmap=True):
self.gl_id = None
self.bilinear = bilinear
self.mipmap = mipmap
# Determine size and dtype.
if image is not None:
image = prepare_texture_data(image)
self.height, self.width, self.channels = image.shape
self.dtype = image.dtype
else:
assert width is not None and height is not None
self.width = width
self.height = height
self.channels = channels if channels is not None else 3
self.dtype = np.dtype(dtype) if dtype is not None else np.uint8
# Validate size and dtype.
assert isinstance(self.width, int) and self.width >= 0
assert isinstance(self.height, int) and self.height >= 0
assert isinstance(self.channels, int) and self.channels >= 1
assert self.is_compatible(width=width, height=height, channels=channels, dtype=dtype)
# Create texture object.
self.gl_id = gl.glGenTextures(1)
with self.bind():
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR if self.bilinear else gl.GL_NEAREST)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR if self.mipmap else gl.GL_NEAREST)
self.update(image)
def delete(self):
if self.gl_id is not None:
gl.glDeleteTextures([self.gl_id])
self.gl_id = None
def __del__(self):
try:
self.delete()
except:
pass
@contextlib.contextmanager
def bind(self):
prev_id = gl.glGetInteger(gl.GL_TEXTURE_BINDING_2D)
gl.glBindTexture(gl.GL_TEXTURE_2D, self.gl_id)
yield
gl.glBindTexture(gl.GL_TEXTURE_2D, prev_id)
def update(self, image):
if image is not None:
image = prepare_texture_data(image)
assert self.is_compatible(image=image)
with self.bind():
fmt = get_texture_format(self.dtype, self.channels)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, fmt.internalformat, self.width, self.height, 0, fmt.format, fmt.type, image)
if self.mipmap:
gl.glGenerateMipmap(gl.GL_TEXTURE_2D)
gl.glPopClientAttrib()
def draw(self, *, pos=0, zoom=1, align=0, rint=False, color=1, alpha=1, rounding=0):
zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
size = zoom * [self.width, self.height]
with self.bind():
gl.glPushAttrib(gl.GL_ENABLE_BIT)
gl.glEnable(gl.GL_TEXTURE_2D)
draw_rect(pos=pos, size=size, align=align, rint=rint, color=color, alpha=alpha, rounding=rounding)
gl.glPopAttrib()
def is_compatible(self, *, image=None, width=None, height=None, channels=None, dtype=None): # pylint: disable=too-many-return-statements
if image is not None:
if image.ndim != 3:
return False
ih, iw, ic = image.shape
if not self.is_compatible(width=iw, height=ih, channels=ic, dtype=image.dtype):
return False
if width is not None and self.width != width:
return False
if height is not None and self.height != height:
return False
if channels is not None and self.channels != channels:
return False
if dtype is not None and self.dtype != dtype:
return False
return True
#----------------------------------------------------------------------------
class Framebuffer:
def __init__(self, *, texture=None, width=None, height=None, channels=None, dtype=None, msaa=0):
self.texture = texture
self.gl_id = None
self.gl_color = None
self.gl_depth_stencil = None
self.msaa = msaa
# Determine size and dtype.
if texture is not None:
assert isinstance(self.texture, Texture)
self.width = texture.width
self.height = texture.height
self.channels = texture.channels
self.dtype = texture.dtype
else:
assert width is not None and height is not None
self.width = width
self.height = height
self.channels = channels if channels is not None else 4
self.dtype = np.dtype(dtype) if dtype is not None else np.float32
# Validate size and dtype.
assert isinstance(self.width, int) and self.width >= 0
assert isinstance(self.height, int) and self.height >= 0
assert isinstance(self.channels, int) and self.channels >= 1
assert width is None or width == self.width
assert height is None or height == self.height
assert channels is None or channels == self.channels
assert dtype is None or dtype == self.dtype
# Create framebuffer object.
self.gl_id = gl.glGenFramebuffers(1)
with self.bind():
# Setup color buffer.
if self.texture is not None:
assert self.msaa == 0
gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self.texture.gl_id, 0)
else:
fmt = get_texture_format(self.dtype, self.channels)
self.gl_color = gl.glGenRenderbuffers(1)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_color)
gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, fmt.internalformat, self.width, self.height)
gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_RENDERBUFFER, self.gl_color)
# Setup depth/stencil buffer.
self.gl_depth_stencil = gl.glGenRenderbuffers(1)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_depth_stencil)
gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, gl.GL_DEPTH24_STENCIL8, self.width, self.height)
gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_DEPTH_STENCIL_ATTACHMENT, gl.GL_RENDERBUFFER, self.gl_depth_stencil)
def delete(self):
if self.gl_id is not None:
gl.glDeleteFramebuffers([self.gl_id])
self.gl_id = None
if self.gl_color is not None:
gl.glDeleteRenderbuffers(1, [self.gl_color])
self.gl_color = None
if self.gl_depth_stencil is not None:
gl.glDeleteRenderbuffers(1, [self.gl_depth_stencil])
self.gl_depth_stencil = None
def __del__(self):
try:
self.delete()
except:
pass
@contextlib.contextmanager
def bind(self):
prev_fbo = gl.glGetInteger(gl.GL_FRAMEBUFFER_BINDING)
prev_rbo = gl.glGetInteger(gl.GL_RENDERBUFFER_BINDING)
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.gl_id)
if self.width is not None and self.height is not None:
gl.glViewport(0, 0, self.width, self.height)
yield
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, prev_fbo)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, prev_rbo)
def blit(self, dst=None):
assert dst is None or isinstance(dst, Framebuffer)
with self.bind():
gl.glBindFramebuffer(gl.GL_DRAW_FRAMEBUFFER, 0 if dst is None else dst.fbo)
gl.glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, gl.GL_COLOR_BUFFER_BIT, gl.GL_NEAREST)
#----------------------------------------------------------------------------
def draw_shape(vertices, *, mode=gl.GL_TRIANGLE_FAN, pos=0, size=1, color=1, alpha=1):
assert vertices.ndim == 2 and vertices.shape[1] == 2
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
size = np.broadcast_to(np.asarray(size, dtype='float32'), [2])
color = np.broadcast_to(np.asarray(color, dtype='float32'), [3])
alpha = np.clip(np.broadcast_to(np.asarray(alpha, dtype='float32'), []), 0, 1)
gl.glPushClientAttrib(gl.GL_CLIENT_VERTEX_ARRAY_BIT)
gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_TRANSFORM_BIT)
gl.glMatrixMode(gl.GL_MODELVIEW)
gl.glPushMatrix()
gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
gl.glEnableClientState(gl.GL_TEXTURE_COORD_ARRAY)
gl.glVertexPointer(2, gl.GL_FLOAT, 0, vertices)
gl.glTexCoordPointer(2, gl.GL_FLOAT, 0, vertices)
gl.glTranslate(pos[0], pos[1], 0)
gl.glScale(size[0], size[1], 1)
gl.glColor4f(color[0] * alpha, color[1] * alpha, color[2] * alpha, alpha)
gl.glDrawArrays(mode, 0, vertices.shape[0])
gl.glPopMatrix()
gl.glPopAttrib()
gl.glPopClientAttrib()
#----------------------------------------------------------------------------
def draw_arrow(x1, y1, x2, y2, l=10, width=1.0):
# Compute the length and angle of the arrow
dx = x2 - x1
dy = y2 - y1
length = math.sqrt(dx**2 + dy**2)
if length < l:
return
angle = math.atan2(dy, dx)
# Save the current modelview matrix
gl.glPushMatrix()
# Translate and rotate the coordinate system
gl.glTranslatef(x1, y1, 0.0)
gl.glRotatef(angle * 180.0 / math.pi, 0.0, 0.0, 1.0)
# Set the line width
gl.glLineWidth(width)
# gl.glColor3f(0.75, 0.75, 0.75)
# Begin drawing lines
gl.glBegin(gl.GL_LINES)
# Draw the shaft of the arrow
gl.glVertex2f(0.0, 0.0)
gl.glVertex2f(length, 0.0)
# Draw the head of the arrow
gl.glVertex2f(length, 0.0)
gl.glVertex2f(length - 2 * l, l)
gl.glVertex2f(length, 0.0)
gl.glVertex2f(length - 2 * l, -l)
# End drawing lines
gl.glEnd()
# Restore the modelview matrix
gl.glPopMatrix()
#----------------------------------------------------------------------------
def draw_rect(*, pos=0, pos2=None, size=None, align=0, rint=False, color=1, alpha=1, rounding=0):
assert pos2 is None or size is None
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
pos2 = np.broadcast_to(np.asarray(pos2, dtype='float32'), [2]) if pos2 is not None else None
size = np.broadcast_to(np.asarray(size, dtype='float32'), [2]) if size is not None else None
size = size if size is not None else pos2 - pos if pos2 is not None else np.array([1, 1], dtype='float32')
pos = pos - size * align
if rint:
pos = np.rint(pos)
rounding = np.broadcast_to(np.asarray(rounding, dtype='float32'), [2])
rounding = np.minimum(np.abs(rounding) / np.maximum(np.abs(size), 1e-8), 0.5)
if np.min(rounding) == 0:
rounding *= 0
vertices = _setup_rect(float(rounding[0]), float(rounding[1]))
draw_shape(vertices, mode=gl.GL_TRIANGLE_FAN, pos=pos, size=size, color=color, alpha=alpha)
@functools.lru_cache(maxsize=10000)
def _setup_rect(rx, ry):
t = np.linspace(0, np.pi / 2, 1 if max(rx, ry) == 0 else 64)
s = 1 - np.sin(t); c = 1 - np.cos(t)
x = [c * rx, 1 - s * rx, 1 - c * rx, s * rx]
y = [s * ry, c * ry, 1 - s * ry, 1 - c * ry]
v = np.stack([x, y], axis=-1).reshape(-1, 2)
return v.astype('float32')
#----------------------------------------------------------------------------
def draw_circle(*, center=0, radius=100, hole=0, color=1, alpha=1):
hole = np.broadcast_to(np.asarray(hole, dtype='float32'), [])
vertices = _setup_circle(float(hole))
draw_shape(vertices, mode=gl.GL_TRIANGLE_STRIP, pos=center, size=radius, color=color, alpha=alpha)
@functools.lru_cache(maxsize=10000)
def _setup_circle(hole):
t = np.linspace(0, np.pi * 2, 128)
s = np.sin(t); c = np.cos(t)
v = np.stack([c, s, c * hole, s * hole], axis=-1).reshape(-1, 2)
return v.astype('float32')
#----------------------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import time
import glfw
import OpenGL.GL as gl
from . import gl_utils
#----------------------------------------------------------------------------
class GlfwWindow: # pylint: disable=too-many-public-methods
def __init__(self, *, title='GlfwWindow', window_width=1920, window_height=1080, deferred_show=True, close_on_esc=True):
self._glfw_window = None
self._drawing_frame = False
self._frame_start_time = None
self._frame_delta = 0
self._fps_limit = None
self._vsync = None
self._skip_frames = 0
self._deferred_show = deferred_show
self._close_on_esc = close_on_esc
self._esc_pressed = False
self._drag_and_drop_paths = None
self._capture_next_frame = False
self._captured_frame = None
# Create window.
glfw.init()
glfw.window_hint(glfw.VISIBLE, False)
self._glfw_window = glfw.create_window(width=window_width, height=window_height, title=title, monitor=None, share=None)
self._attach_glfw_callbacks()
self.make_context_current()
# Adjust window.
self.set_vsync(False)
self.set_window_size(window_width, window_height)
if not self._deferred_show:
glfw.show_window(self._glfw_window)
def close(self):
if self._drawing_frame:
self.end_frame()
if self._glfw_window is not None:
glfw.destroy_window(self._glfw_window)
self._glfw_window = None
#glfw.terminate() # Commented out to play it nice with other glfw clients.
def __del__(self):
try:
self.close()
except:
pass
@property
def window_width(self):
return self.content_width
@property
def window_height(self):
return self.content_height + self.title_bar_height
@property
def content_width(self):
width, _height = glfw.get_window_size(self._glfw_window)
return width
@property
def content_height(self):
_width, height = glfw.get_window_size(self._glfw_window)
return height
@property
def title_bar_height(self):
_left, top, _right, _bottom = glfw.get_window_frame_size(self._glfw_window)
return top
@property
def monitor_width(self):
_, _, width, _height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
return width
@property
def monitor_height(self):
_, _, _width, height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
return height
@property
def frame_delta(self):
return self._frame_delta
def set_title(self, title):
glfw.set_window_title(self._glfw_window, title)
def set_window_size(self, width, height):
width = min(width, self.monitor_width)
height = min(height, self.monitor_height)
glfw.set_window_size(self._glfw_window, width, max(height - self.title_bar_height, 0))
if width == self.monitor_width and height == self.monitor_height:
self.maximize()
def set_content_size(self, width, height):
self.set_window_size(width, height + self.title_bar_height)
def maximize(self):
glfw.maximize_window(self._glfw_window)
def set_position(self, x, y):
glfw.set_window_pos(self._glfw_window, x, y + self.title_bar_height)
def center(self):
self.set_position((self.monitor_width - self.window_width) // 2, (self.monitor_height - self.window_height) // 2)
def set_vsync(self, vsync):
vsync = bool(vsync)
if vsync != self._vsync:
glfw.swap_interval(1 if vsync else 0)
self._vsync = vsync
def set_fps_limit(self, fps_limit):
self._fps_limit = int(fps_limit)
def should_close(self):
return glfw.window_should_close(self._glfw_window) or (self._close_on_esc and self._esc_pressed)
def skip_frame(self):
self.skip_frames(1)
def skip_frames(self, num): # Do not update window for the next N frames.
self._skip_frames = max(self._skip_frames, int(num))
def is_skipping_frames(self):
return self._skip_frames > 0
def capture_next_frame(self):
self._capture_next_frame = True
def pop_captured_frame(self):
frame = self._captured_frame
self._captured_frame = None
return frame
def pop_drag_and_drop_paths(self):
paths = self._drag_and_drop_paths
self._drag_and_drop_paths = None
return paths
def draw_frame(self): # To be overridden by subclass.
self.begin_frame()
# Rendering code goes here.
self.end_frame()
def make_context_current(self):
if self._glfw_window is not None:
glfw.make_context_current(self._glfw_window)
def begin_frame(self):
# End previous frame.
if self._drawing_frame:
self.end_frame()
# Apply FPS limit.
if self._frame_start_time is not None and self._fps_limit is not None:
delay = self._frame_start_time - time.perf_counter() + 1 / self._fps_limit
if delay > 0:
time.sleep(delay)
cur_time = time.perf_counter()
if self._frame_start_time is not None:
self._frame_delta = cur_time - self._frame_start_time
self._frame_start_time = cur_time
# Process events.
glfw.poll_events()
# Begin frame.
self._drawing_frame = True
self.make_context_current()
# Initialize GL state.
gl.glViewport(0, 0, self.content_width, self.content_height)
gl.glMatrixMode(gl.GL_PROJECTION)
gl.glLoadIdentity()
gl.glTranslate(-1, 1, 0)
gl.glScale(2 / max(self.content_width, 1), -2 / max(self.content_height, 1), 1)
gl.glMatrixMode(gl.GL_MODELVIEW)
gl.glLoadIdentity()
gl.glEnable(gl.GL_BLEND)
gl.glBlendFunc(gl.GL_ONE, gl.GL_ONE_MINUS_SRC_ALPHA) # Pre-multiplied alpha.
# Clear.
gl.glClearColor(0, 0, 0, 1)
gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
def end_frame(self):
assert self._drawing_frame
self._drawing_frame = False
# Skip frames if requested.
if self._skip_frames > 0:
self._skip_frames -= 1
return
# Capture frame if requested.
if self._capture_next_frame:
self._captured_frame = gl_utils.read_pixels(self.content_width, self.content_height)
self._capture_next_frame = False
# Update window.
if self._deferred_show:
glfw.show_window(self._glfw_window)
self._deferred_show = False
glfw.swap_buffers(self._glfw_window)
def _attach_glfw_callbacks(self):
glfw.set_key_callback(self._glfw_window, self._glfw_key_callback)
glfw.set_drop_callback(self._glfw_window, self._glfw_drop_callback)
def _glfw_key_callback(self, _window, key, _scancode, action, _mods):
if action == glfw.PRESS and key == glfw.KEY_ESCAPE:
self._esc_pressed = True
def _glfw_drop_callback(self, _window, paths):
self._drag_and_drop_paths = paths
#----------------------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import contextlib
import imgui
#----------------------------------------------------------------------------
def set_default_style(color_scheme='dark', spacing=9, indent=23, scrollbar=27):
s = imgui.get_style()
s.window_padding = [spacing, spacing]
s.item_spacing = [spacing, spacing]
s.item_inner_spacing = [spacing, spacing]
s.columns_min_spacing = spacing
s.indent_spacing = indent
s.scrollbar_size = scrollbar
s.frame_padding = [4, 3]
s.window_border_size = 1
s.child_border_size = 1
s.popup_border_size = 1
s.frame_border_size = 1
s.window_rounding = 0
s.child_rounding = 0
s.popup_rounding = 3
s.frame_rounding = 3
s.scrollbar_rounding = 3
s.grab_rounding = 3
getattr(imgui, f'style_colors_{color_scheme}')(s)
c0 = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
c1 = s.colors[imgui.COLOR_FRAME_BACKGROUND]
s.colors[imgui.COLOR_POPUP_BACKGROUND] = [x * 0.7 + y * 0.3 for x, y in zip(c0, c1)][:3] + [1]
#----------------------------------------------------------------------------
@contextlib.contextmanager
def grayed_out(cond=True):
if cond:
s = imgui.get_style()
text = s.colors[imgui.COLOR_TEXT_DISABLED]
grab = s.colors[imgui.COLOR_SCROLLBAR_GRAB]
back = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
imgui.push_style_color(imgui.COLOR_TEXT, *text)
imgui.push_style_color(imgui.COLOR_CHECK_MARK, *grab)
imgui.push_style_color(imgui.COLOR_SLIDER_GRAB, *grab)
imgui.push_style_color(imgui.COLOR_SLIDER_GRAB_ACTIVE, *grab)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND, *back)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_BUTTON, *back)
imgui.push_style_color(imgui.COLOR_BUTTON_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_BUTTON_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_HEADER, *back)
imgui.push_style_color(imgui.COLOR_HEADER_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_HEADER_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_POPUP_BACKGROUND, *back)
yield
imgui.pop_style_color(14)
else:
yield
#----------------------------------------------------------------------------
@contextlib.contextmanager
def item_width(width=None):
if width is not None:
imgui.push_item_width(width)
yield
imgui.pop_item_width()
else:
yield
#----------------------------------------------------------------------------
def scoped_by_object_id(method):
def decorator(self, *args, **kwargs):
imgui.push_id(str(id(self)))
res = method(self, *args, **kwargs)
imgui.pop_id()
return res
return decorator
#----------------------------------------------------------------------------
def button(label, width=0, enabled=True):
with grayed_out(not enabled):
clicked = imgui.button(label, width=width)
clicked = clicked and enabled
return clicked
#----------------------------------------------------------------------------
def collapsing_header(text, visible=None, flags=0, default=False, enabled=True, show=True):
expanded = False
if show:
if default:
flags |= imgui.TREE_NODE_DEFAULT_OPEN
if not enabled:
flags |= imgui.TREE_NODE_LEAF
with grayed_out(not enabled):
expanded, visible = imgui.collapsing_header(text, visible=visible, flags=flags)
expanded = expanded and enabled
return expanded, visible
#----------------------------------------------------------------------------
def popup_button(label, width=0, enabled=True):
if button(label, width, enabled):
imgui.open_popup(label)
opened = imgui.begin_popup(label)
return opened
#----------------------------------------------------------------------------
def input_text(label, value, buffer_length, flags, width=None, help_text=''):
old_value = value
color = list(imgui.get_style().colors[imgui.COLOR_TEXT])
if value == '':
color[-1] *= 0.5
with item_width(width):
imgui.push_style_color(imgui.COLOR_TEXT, *color)
value = value if value != '' else help_text
changed, value = imgui.input_text(label, value, buffer_length, flags)
value = value if value != help_text else ''
imgui.pop_style_color(1)
if not flags & imgui.INPUT_TEXT_ENTER_RETURNS_TRUE:
changed = (value != old_value)
return changed, value
#----------------------------------------------------------------------------
def drag_previous_control(enabled=True):
dragging = False
dx = 0
dy = 0
if imgui.begin_drag_drop_source(imgui.DRAG_DROP_SOURCE_NO_PREVIEW_TOOLTIP):
if enabled:
dragging = True
dx, dy = imgui.get_mouse_drag_delta()
imgui.reset_mouse_drag_delta()
imgui.end_drag_drop_source()
return dragging, dx, dy
#----------------------------------------------------------------------------
def drag_button(label, width=0, enabled=True):
clicked = button(label, width=width, enabled=enabled)
dragging, dx, dy = drag_previous_control(enabled=enabled)
return clicked, dragging, dx, dy
#----------------------------------------------------------------------------
def drag_hidden_window(label, x, y, width, height, enabled=True):
imgui.push_style_color(imgui.COLOR_WINDOW_BACKGROUND, 0, 0, 0, 0)
imgui.push_style_color(imgui.COLOR_BORDER, 0, 0, 0, 0)
imgui.set_next_window_position(x, y)
imgui.set_next_window_size(width, height)
imgui.begin(label, closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
dragging, dx, dy = drag_previous_control(enabled=enabled)
imgui.end()
imgui.pop_style_color(2)
return dragging, dx, dy
#----------------------------------------------------------------------------
def click_hidden_window(label, x, y, width, height, img_w, img_h, enabled=True):
imgui.push_style_color(imgui.COLOR_WINDOW_BACKGROUND, 0, 0, 0, 0)
imgui.push_style_color(imgui.COLOR_BORDER, 0, 0, 0, 0)
imgui.set_next_window_position(x, y)
imgui.set_next_window_size(width, height)
imgui.begin(label, closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
clicked, down = False, False
img_x, img_y = 0, 0
if imgui.is_mouse_down():
posx, posy = imgui.get_mouse_pos()
if posx >= x and posx < x + width and posy >= y and posy < y + height:
if imgui.is_mouse_clicked():
clicked = True
down = True
img_x = round((posx - x) / (width - 1) * (img_w - 1))
img_y = round((posy - y) / (height - 1) * (img_h - 1))
imgui.end()
imgui.pop_style_color(2)
return clicked, down, img_x, img_y
#----------------------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import os
import imgui
import imgui.integrations.glfw
from . import glfw_window
from . import imgui_utils
from . import text_utils
#----------------------------------------------------------------------------
class ImguiWindow(glfw_window.GlfwWindow):
def __init__(self, *, title='ImguiWindow', font=None, font_sizes=range(14,24), **glfw_kwargs):
if font is None:
font = text_utils.get_default_font()
font_sizes = {int(size) for size in font_sizes}
super().__init__(title=title, **glfw_kwargs)
# Init fields.
self._imgui_context = None
self._imgui_renderer = None
self._imgui_fonts = None
self._cur_font_size = max(font_sizes)
# Delete leftover imgui.ini to avoid unexpected behavior.
if os.path.isfile('imgui.ini'):
os.remove('imgui.ini')
# Init ImGui.
self._imgui_context = imgui.create_context()
self._imgui_renderer = _GlfwRenderer(self._glfw_window)
self._attach_glfw_callbacks()
imgui.get_io().ini_saving_rate = 0 # Disable creating imgui.ini at runtime.
imgui.get_io().mouse_drag_threshold = 0 # Improve behavior with imgui_utils.drag_custom().
self._imgui_fonts = {size: imgui.get_io().fonts.add_font_from_file_ttf(font, size) for size in font_sizes}
self._imgui_renderer.refresh_font_texture()
def close(self):
self.make_context_current()
self._imgui_fonts = None
if self._imgui_renderer is not None:
self._imgui_renderer.shutdown()
self._imgui_renderer = None
if self._imgui_context is not None:
#imgui.destroy_context(self._imgui_context) # Commented out to avoid creating imgui.ini at the end.
self._imgui_context = None
super().close()
def _glfw_key_callback(self, *args):
super()._glfw_key_callback(*args)
self._imgui_renderer.keyboard_callback(*args)
@property
def font_size(self):
return self._cur_font_size
@property
def spacing(self):
return round(self._cur_font_size * 0.4)
def set_font_size(self, target): # Applied on next frame.
self._cur_font_size = min((abs(key - target), key) for key in self._imgui_fonts.keys())[1]
def begin_frame(self):
# Begin glfw frame.
super().begin_frame()
# Process imgui events.
self._imgui_renderer.mouse_wheel_multiplier = self._cur_font_size / 10
if self.content_width > 0 and self.content_height > 0:
self._imgui_renderer.process_inputs()
# Begin imgui frame.
imgui.new_frame()
imgui.push_font(self._imgui_fonts[self._cur_font_size])
imgui_utils.set_default_style(spacing=self.spacing, indent=self.font_size, scrollbar=self.font_size+4)
def end_frame(self):
imgui.pop_font()
imgui.render()
imgui.end_frame()
self._imgui_renderer.render(imgui.get_draw_data())
super().end_frame()
#----------------------------------------------------------------------------
# Wrapper class for GlfwRenderer to fix a mouse wheel bug on Linux.
class _GlfwRenderer(imgui.integrations.glfw.GlfwRenderer):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.mouse_wheel_multiplier = 1
def scroll_callback(self, window, x_offset, y_offset):
self.io.mouse_wheel += y_offset * self.mouse_wheel_multiplier
#----------------------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import functools
from typing import Optional
import dnnlib
import numpy as np
import PIL.Image
import PIL.ImageFont
import scipy.ndimage
from . import gl_utils
#----------------------------------------------------------------------------
def get_default_font():
url = 'http://fonts.gstatic.com/s/opensans/v17/mem8YaGs126MiZpBA-U1UpcaXcl0Aw.ttf' # Open Sans regular
return dnnlib.util.open_url(url, return_filename=True)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=None)
def get_pil_font(font=None, size=32):
if font is None:
font = get_default_font()
return PIL.ImageFont.truetype(font=font, size=size)
#----------------------------------------------------------------------------
def get_array(string, *, dropshadow_radius: int=None, **kwargs):
if dropshadow_radius is not None:
offset_x = int(np.ceil(dropshadow_radius*2/3))
offset_y = int(np.ceil(dropshadow_radius*2/3))
return _get_array_priv(string, dropshadow_radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
else:
return _get_array_priv(string, **kwargs)
@functools.lru_cache(maxsize=10000)
def _get_array_priv(
string: str, *,
size: int = 32,
max_width: Optional[int]=None,
max_height: Optional[int]=None,
min_size=10,
shrink_coef=0.8,
dropshadow_radius: int=None,
offset_x: int=None,
offset_y: int=None,
**kwargs
):
cur_size = size
array = None
while True:
if dropshadow_radius is not None:
# separate implementation for dropshadow text rendering
array = _get_array_impl_dropshadow(string, size=cur_size, radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
else:
array = _get_array_impl(string, size=cur_size, **kwargs)
height, width, _ = array.shape
if (max_width is None or width <= max_width) and (max_height is None or height <= max_height) or (cur_size <= min_size):
break
cur_size = max(int(cur_size * shrink_coef), min_size)
return array
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def _get_array_impl(string, *, font=None, size=32, outline=0, outline_pad=3, outline_coef=3, outline_exp=2, line_pad: int=None):
pil_font = get_pil_font(font=font, size=size)
lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
width = max(line.shape[1] for line in lines)
lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
line_spacing = line_pad if line_pad is not None else size // 2
lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
mask = np.concatenate(lines, axis=0)
alpha = mask
if outline > 0:
mask = np.pad(mask, int(np.ceil(outline * outline_pad)), mode='constant', constant_values=0)
alpha = mask.astype(np.float32) / 255
alpha = scipy.ndimage.gaussian_filter(alpha, outline)
alpha = 1 - np.maximum(1 - alpha * outline_coef, 0) ** outline_exp
alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
alpha = np.maximum(alpha, mask)
return np.stack([mask, alpha], axis=-1)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def _get_array_impl_dropshadow(string, *, font=None, size=32, radius: int, offset_x: int, offset_y: int, line_pad: int=None, **kwargs):
assert (offset_x > 0) and (offset_y > 0)
pil_font = get_pil_font(font=font, size=size)
lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
width = max(line.shape[1] for line in lines)
lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
line_spacing = line_pad if line_pad is not None else size // 2
lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
mask = np.concatenate(lines, axis=0)
alpha = mask
mask = np.pad(mask, 2*radius + max(abs(offset_x), abs(offset_y)), mode='constant', constant_values=0)
alpha = mask.astype(np.float32) / 255
alpha = scipy.ndimage.gaussian_filter(alpha, radius)
alpha = 1 - np.maximum(1 - alpha * 1.5, 0) ** 1.4
alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
alpha = np.pad(alpha, [(offset_y, 0), (offset_x, 0)], mode='constant')[:-offset_y, :-offset_x]
alpha = np.maximum(alpha, mask)
return np.stack([mask, alpha], axis=-1)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def get_texture(string, bilinear=True, mipmap=True, **kwargs):
return gl_utils.Texture(image=get_array(string, **kwargs), bilinear=bilinear, mipmap=mipmap)
#----------------------------------------------------------------------------
icon.png

68.4 KB

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
"""Converting legacy network pickle into the new format."""
import click
import pickle
import re
import copy
import numpy as np
import torch
import dnnlib
from torch_utils import misc
#----------------------------------------------------------------------------
def load_network_pkl(f, force_fp16=False):
data = _LegacyUnpickler(f).load()
# Legacy TensorFlow pickle => convert.
if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):
tf_G, tf_D, tf_Gs = data
G = convert_tf_generator(tf_G)
D = convert_tf_discriminator(tf_D)
G_ema = convert_tf_generator(tf_Gs)
data = dict(G=G, D=D, G_ema=G_ema)
# Add missing fields.
if 'training_set_kwargs' not in data:
data['training_set_kwargs'] = None
if 'augment_pipe' not in data:
data['augment_pipe'] = None
# Validate contents.
assert isinstance(data['G'], torch.nn.Module)
assert isinstance(data['D'], torch.nn.Module)
assert isinstance(data['G_ema'], torch.nn.Module)
assert isinstance(data['training_set_kwargs'], (dict, type(None)))
assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))
# Force FP16.
if force_fp16:
for key in ['G', 'D', 'G_ema']:
old = data[key]
kwargs = copy.deepcopy(old.init_kwargs)
fp16_kwargs = kwargs.get('synthesis_kwargs', kwargs)
fp16_kwargs.num_fp16_res = 4
fp16_kwargs.conv_clamp = 256
if kwargs != old.init_kwargs:
new = type(old)(**kwargs).eval().requires_grad_(False)
misc.copy_params_and_buffers(old, new, require_all=True)
data[key] = new
return data
#----------------------------------------------------------------------------
class _TFNetworkStub(dnnlib.EasyDict):
pass
class _LegacyUnpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'dnnlib.tflib.network' and name == 'Network':
return _TFNetworkStub
return super().find_class(module, name)
#----------------------------------------------------------------------------
def _collect_tf_params(tf_net):
# pylint: disable=protected-access
tf_params = dict()
def recurse(prefix, tf_net):
for name, value in tf_net.variables:
tf_params[prefix + name] = value
for name, comp in tf_net.components.items():
recurse(prefix + name + '/', comp)
recurse('', tf_net)
return tf_params
#----------------------------------------------------------------------------
def _populate_module_params(module, *patterns):
for name, tensor in misc.named_params_and_buffers(module):
found = False
value = None
for pattern, value_fn in zip(patterns[0::2], patterns[1::2]):
match = re.fullmatch(pattern, name)
if match:
found = True
if value_fn is not None:
value = value_fn(*match.groups())
break
try:
assert found
if value is not None:
tensor.copy_(torch.from_numpy(np.array(value)))
except:
print(name, list(tensor.shape))
raise
#----------------------------------------------------------------------------
def convert_tf_generator(tf_G):
if tf_G.version < 4:
raise ValueError('TensorFlow pickle version too low')
# Collect kwargs.
tf_kwargs = tf_G.static_kwargs
known_kwargs = set()
def kwarg(tf_name, default=None, none=None):
known_kwargs.add(tf_name)
val = tf_kwargs.get(tf_name, default)
return val if val is not None else none
# Convert kwargs.
from training import networks_stylegan2
network_class = networks_stylegan2.Generator
kwargs = dnnlib.EasyDict(
z_dim = kwarg('latent_size', 512),
c_dim = kwarg('label_size', 0),
w_dim = kwarg('dlatent_size', 512),
img_resolution = kwarg('resolution', 1024),
img_channels = kwarg('num_channels', 3),
channel_base = kwarg('fmap_base', 16384) * 2,
channel_max = kwarg('fmap_max', 512),
num_fp16_res = kwarg('num_fp16_res', 0),
conv_clamp = kwarg('conv_clamp', None),
architecture = kwarg('architecture', 'skip'),
resample_filter = kwarg('resample_kernel', [1,3,3,1]),
use_noise = kwarg('use_noise', True),
activation = kwarg('nonlinearity', 'lrelu'),
mapping_kwargs = dnnlib.EasyDict(
num_layers = kwarg('mapping_layers', 8),
embed_features = kwarg('label_fmaps', None),
layer_features = kwarg('mapping_fmaps', None),
activation = kwarg('mapping_nonlinearity', 'lrelu'),
lr_multiplier = kwarg('mapping_lrmul', 0.01),
w_avg_beta = kwarg('w_avg_beta', 0.995, none=1),
),
)
# Check for unknown kwargs.
kwarg('truncation_psi')
kwarg('truncation_cutoff')
kwarg('style_mixing_prob')
kwarg('structure')
kwarg('conditioning')
kwarg('fused_modconv')
unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
if len(unknown_kwargs) > 0:
raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
# Collect params.
tf_params = _collect_tf_params(tf_G)
for name, value in list(tf_params.items()):
match = re.fullmatch(r'ToRGB_lod(\d+)/(.*)', name)
if match:
r = kwargs.img_resolution // (2 ** int(match.group(1)))
tf_params[f'{r}x{r}/ToRGB/{match.group(2)}'] = value
kwargs.synthesis.kwargs.architecture = 'orig'
#for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
# Convert params.
G = network_class(**kwargs).eval().requires_grad_(False)
# pylint: disable=unnecessary-lambda
# pylint: disable=f-string-without-interpolation
_populate_module_params(G,
r'mapping\.w_avg', lambda: tf_params[f'dlatent_avg'],
r'mapping\.embed\.weight', lambda: tf_params[f'mapping/LabelEmbed/weight'].transpose(),
r'mapping\.embed\.bias', lambda: tf_params[f'mapping/LabelEmbed/bias'],
r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'mapping/Dense{i}/weight'].transpose(),
r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'mapping/Dense{i}/bias'],
r'synthesis\.b4\.const', lambda: tf_params[f'synthesis/4x4/Const/const'][0],
r'synthesis\.b4\.conv1\.weight', lambda: tf_params[f'synthesis/4x4/Conv/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b4\.conv1\.bias', lambda: tf_params[f'synthesis/4x4/Conv/bias'],
r'synthesis\.b4\.conv1\.noise_const', lambda: tf_params[f'synthesis/noise0'][0, 0],
r'synthesis\.b4\.conv1\.noise_strength', lambda: tf_params[f'synthesis/4x4/Conv/noise_strength'],
r'synthesis\.b4\.conv1\.affine\.weight', lambda: tf_params[f'synthesis/4x4/Conv/mod_weight'].transpose(),
r'synthesis\.b4\.conv1\.affine\.bias', lambda: tf_params[f'synthesis/4x4/Conv/mod_bias'] + 1,
r'synthesis\.b(\d+)\.conv0\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.conv0\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/bias'],
r'synthesis\.b(\d+)\.conv0\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-5}'][0, 0],
r'synthesis\.b(\d+)\.conv0\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/noise_strength'],
r'synthesis\.b(\d+)\.conv0\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.conv0\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_bias'] + 1,
r'synthesis\.b(\d+)\.conv1\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.conv1\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/bias'],
r'synthesis\.b(\d+)\.conv1\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-4}'][0, 0],
r'synthesis\.b(\d+)\.conv1\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/noise_strength'],
r'synthesis\.b(\d+)\.conv1\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.conv1\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_bias'] + 1,
r'synthesis\.b(\d+)\.torgb\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.torgb\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/bias'],
r'synthesis\.b(\d+)\.torgb\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.torgb\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_bias'] + 1,
r'synthesis\.b(\d+)\.skip\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Skip/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
r'.*\.resample_filter', None,
r'.*\.act_filter', None,
)
return G
#----------------------------------------------------------------------------
def convert_tf_discriminator(tf_D):
if tf_D.version < 4:
raise ValueError('TensorFlow pickle version too low')
# Collect kwargs.
tf_kwargs = tf_D.static_kwargs
known_kwargs = set()
def kwarg(tf_name, default=None):
known_kwargs.add(tf_name)
return tf_kwargs.get(tf_name, default)
# Convert kwargs.
kwargs = dnnlib.EasyDict(
c_dim = kwarg('label_size', 0),
img_resolution = kwarg('resolution', 1024),
img_channels = kwarg('num_channels', 3),
architecture = kwarg('architecture', 'resnet'),
channel_base = kwarg('fmap_base', 16384) * 2,
channel_max = kwarg('fmap_max', 512),
num_fp16_res = kwarg('num_fp16_res', 0),
conv_clamp = kwarg('conv_clamp', None),
cmap_dim = kwarg('mapping_fmaps', None),
block_kwargs = dnnlib.EasyDict(
activation = kwarg('nonlinearity', 'lrelu'),
resample_filter = kwarg('resample_kernel', [1,3,3,1]),
freeze_layers = kwarg('freeze_layers', 0),
),
mapping_kwargs = dnnlib.EasyDict(
num_layers = kwarg('mapping_layers', 0),
embed_features = kwarg('mapping_fmaps', None),
layer_features = kwarg('mapping_fmaps', None),
activation = kwarg('nonlinearity', 'lrelu'),
lr_multiplier = kwarg('mapping_lrmul', 0.1),
),
epilogue_kwargs = dnnlib.EasyDict(
mbstd_group_size = kwarg('mbstd_group_size', None),
mbstd_num_channels = kwarg('mbstd_num_features', 1),
activation = kwarg('nonlinearity', 'lrelu'),
),
)
# Check for unknown kwargs.
kwarg('structure')
kwarg('conditioning')
unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
if len(unknown_kwargs) > 0:
raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
# Collect params.
tf_params = _collect_tf_params(tf_D)
for name, value in list(tf_params.items()):
match = re.fullmatch(r'FromRGB_lod(\d+)/(.*)', name)
if match:
r = kwargs.img_resolution // (2 ** int(match.group(1)))
tf_params[f'{r}x{r}/FromRGB/{match.group(2)}'] = value
kwargs.architecture = 'orig'
#for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
# Convert params.
from training import networks_stylegan2
D = networks_stylegan2.Discriminator(**kwargs).eval().requires_grad_(False)
# pylint: disable=unnecessary-lambda
# pylint: disable=f-string-without-interpolation
_populate_module_params(D,
r'b(\d+)\.fromrgb\.weight', lambda r: tf_params[f'{r}x{r}/FromRGB/weight'].transpose(3, 2, 0, 1),
r'b(\d+)\.fromrgb\.bias', lambda r: tf_params[f'{r}x{r}/FromRGB/bias'],
r'b(\d+)\.conv(\d+)\.weight', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/weight'].transpose(3, 2, 0, 1),
r'b(\d+)\.conv(\d+)\.bias', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/bias'],
r'b(\d+)\.skip\.weight', lambda r: tf_params[f'{r}x{r}/Skip/weight'].transpose(3, 2, 0, 1),
r'mapping\.embed\.weight', lambda: tf_params[f'LabelEmbed/weight'].transpose(),
r'mapping\.embed\.bias', lambda: tf_params[f'LabelEmbed/bias'],
r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'Mapping{i}/weight'].transpose(),
r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'Mapping{i}/bias'],
r'b4\.conv\.weight', lambda: tf_params[f'4x4/Conv/weight'].transpose(3, 2, 0, 1),
r'b4\.conv\.bias', lambda: tf_params[f'4x4/Conv/bias'],
r'b4\.fc\.weight', lambda: tf_params[f'4x4/Dense0/weight'].transpose(),
r'b4\.fc\.bias', lambda: tf_params[f'4x4/Dense0/bias'],
r'b4\.out\.weight', lambda: tf_params[f'Output/weight'].transpose(),
r'b4\.out\.bias', lambda: tf_params[f'Output/bias'],
r'.*\.resample_filter', None,
)
return D
#----------------------------------------------------------------------------
@click.command()
@click.option('--source', help='Input pickle', required=True, metavar='PATH')
@click.option('--dest', help='Output pickle', required=True, metavar='PATH')
@click.option('--force-fp16', help='Force the networks to use FP16', type=bool, default=False, metavar='BOOL', show_default=True)
def convert_network_pickle(source, dest, force_fp16):
"""Convert legacy network pickle into the native PyTorch format.
The tool is able to load the main network configurations exported using the TensorFlow version of StyleGAN2 or StyleGAN2-ADA.
It does not support e.g. StyleGAN2-ADA comparison methods, StyleGAN2 configs A-D, or StyleGAN1 networks.
Example:
\b
python legacy.py \\
--source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \\
--dest=stylegan2-cat-config-f.pkl
"""
print(f'Loading "{source}"...')
with dnnlib.util.open_url(source) as f:
data = load_network_pkl(f, force_fp16=force_fp16)
print(f'Saving "{dest}"...')
with open(dest, 'wb') as f:
pickle.dump(data, f)
print('Done.')
#----------------------------------------------------------------------------
if __name__ == "__main__":
convert_network_pickle() # pylint: disable=no-value-for-parameter
#----------------------------------------------------------------------------
# 模型唯一标识
modelCode = 1062
# 模型名称
modelName=draggan_pytorch
# 模型描述
modelDescription=通过拖拽编辑生成图像。
# 应用场景
appScenario=AIGC,零售,制造,电商,医疗,教育
# 框架类型
frameType=pytorch
# torch>=2.0.0
scipy>=1.11.1
Ninja==1.10.2
gradio==3.36.1
imageio-ffmpeg>=0.4.3
huggingface_hub
hf_transfer
pyopengl
imgui
glfw==2.6.1
pillow==9.5.0
# torchvision>=0.15.2
imageio>=2.9.0
import os
import sys
import json
import requests
from tqdm import tqdm
def download_file(url: str, filename: str, download_dir: str):
"""Download a file if it does not already exist."""
try:
filepath = os.path.join(download_dir, filename)
content_length = int(requests.head(url).headers.get("content-length", 0))
# If file already exists and size matches, skip download
if os.path.isfile(filepath) and os.path.getsize(filepath) == content_length:
print(f"{filepath} already exists. Skipping download.")
return
if os.path.isfile(filepath) and os.path.getsize(filepath) != content_length:
print(f"{filepath} already exists but size does not match. Redownloading.")
else:
print(f"Downloading {filename} from {url}")
# Start download, stream=True allows for progress tracking
response = requests.get(url, stream=True)
# Check if request was successful
response.raise_for_status()
# Create progress bar
total_size = int(response.headers.get('content-length', 0))
progress_bar = tqdm(
total=total_size,
unit='iB',
unit_scale=True,
ncols=70,
file=sys.stdout
)
# Write response content to file
with open(filepath, 'wb') as f:
for data in response.iter_content(chunk_size=1024):
f.write(data)
progress_bar.update(len(data)) # Update progress bar
# Close progress bar
progress_bar.close()
# Error handling for incomplete downloads
if total_size != 0 and progress_bar.n != total_size:
print("ERROR, something went wrong while downloading")
raise Exception()
except Exception as e:
print(f"An error occurred: {e}")
def main():
"""Main function to download files from URLs in a config file."""
# Get JSON config file path
script_dir = os.path.dirname(os.path.realpath(__file__))
config_file_path = os.path.join(script_dir, "download_models.json")
# Set download directory
download_dir = "checkpoints"
os.makedirs(download_dir, exist_ok=True)
# Load URL and filenames from JSON
with open(config_file_path, "r") as f:
config = json.load(f)
# Download each file specified in config
for url, filename in config.items():
download_file(url, filename, download_dir)
if __name__ == "__main__":
main()
{
"https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl": "stylegan2_lions_512_pytorch.pkl",
"https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl": "stylegan2_dogs_1024_pytorch.pkl",
"https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl": "stylegan2_horses_256_pytorch.pkl",
"https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl": "stylegan2_elephants_512_pytorch.pkl",
"https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl": "stylegan2-ffhq-512x512.pkl",
"https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl": "stylegan2-afhqcat-512x512.pkl",
"http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl": "stylegan2-car-config-f.pkl",
"http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl": "stylegan2-cat-config-f.pkl"
}
@echo off
python visualizer_drag.py ^
checkpoints/stylegan2_lions_512_pytorch.pkl ^
checkpoints/stylegan2-ffhq-512x512.pkl ^
checkpoints/stylegan2-afhqcat-512x512.pkl ^
checkpoints/stylegan2-car-config-f.pkl ^
checkpoints/stylegan2_dogs_1024_pytorch.pkl ^
checkpoints/stylegan2_horses_256_pytorch.pkl ^
checkpoints/stylegan2-cat-config-f.pkl ^
checkpoints/stylegan2_elephants_512_pytorch.pkl ^
checkpoints/stylegan_human_v2_512.pkl ^
checkpoints/stylegan2-lhq-256x256.pkl
python visualizer_drag.py \
checkpoints/stylegan2_lions_512_pytorch.pkl \
checkpoints/stylegan2-ffhq-512x512.pkl \
checkpoints/stylegan2-afhqcat-512x512.pkl \
checkpoints/stylegan2-car-config-f.pkl \
checkpoints/stylegan2_dogs_1024_pytorch.pkl \
checkpoints/stylegan2_horses_256_pytorch.pkl \
checkpoints/stylegan2-cat-config-f.pkl \
checkpoints/stylegan2_elephants_512_pytorch.pkl \
checkpoints/stylegan_human_v2_512.pkl \
checkpoints/stylegan2-lhq-256x256.pkl
.DS_Store
__pycache__
*.pt
*.pth
*.pdparams
*.pdiparams
*.pdmodel
*.pkl
*.info
*.yaml
\ No newline at end of file
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import codecs
import os
import time
import yaml
import numpy as np
import cv2
import paddle
import paddleseg.transforms as T
from paddle.inference import create_predictor, PrecisionType
from paddle.inference import Config as PredictConfig
from paddleseg.core.infer import reverse_transform
from paddleseg.cvlibs import manager
from paddleseg.utils import TimeAverager
from ..scripts.optic_flow_process import optic_flow_process
class DeployConfig:
def __init__(self, path):
with codecs.open(path, 'r', 'utf-8') as file:
self.dic = yaml.load(file, Loader=yaml.FullLoader)
self._transforms = self._load_transforms(self.dic['Deploy'][
'transforms'])
self._dir = os.path.dirname(path)
@property
def transforms(self):
return self._transforms
@property
def model(self):
return os.path.join(self._dir, self.dic['Deploy']['model'])
@property
def params(self):
return os.path.join(self._dir, self.dic['Deploy']['params'])
def _load_transforms(self, t_list):
com = manager.TRANSFORMS
transforms = []
for t in t_list:
ctype = t.pop('type')
transforms.append(com[ctype](**t))
return transforms
class Predictor:
def __init__(self, args):
self.cfg = DeployConfig(args.cfg)
self.args = args
self.compose = T.Compose(self.cfg.transforms)
resize_h, resize_w = args.input_shape
self.disflow = cv2.DISOpticalFlow_create(
cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
self.prev_gray = np.zeros((resize_h, resize_w), np.uint8)
self.prev_cfd = np.zeros((resize_h, resize_w), np.float32)
self.is_init = True
pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
pred_cfg.disable_glog_info()
if self.args.use_gpu:
pred_cfg.enable_use_gpu(100, 0)
self.predictor = create_predictor(pred_cfg)
if self.args.test_speed:
self.cost_averager = TimeAverager()
def preprocess(self, img):
ori_shapes = []
processed_imgs = []
processed_img = self.compose(img)[0]
processed_imgs.append(processed_img)
ori_shapes.append(img.shape)
return processed_imgs, ori_shapes
def run(self, img, bg):
input_names = self.predictor.get_input_names()
input_handle = self.predictor.get_input_handle(input_names[0])
processed_imgs, ori_shapes = self.preprocess(img)
data = np.array(processed_imgs)
input_handle.reshape(data.shape)
input_handle.copy_from_cpu(data)
if self.args.test_speed:
start = time.time()
self.predictor.run()
if self.args.test_speed:
self.cost_averager.record(time.time() - start)
output_names = self.predictor.get_output_names()
output_handle = self.predictor.get_output_handle(output_names[0])
output = output_handle.copy_to_cpu()
return self.postprocess(output, img, ori_shapes[0], bg)
def postprocess(self, pred, img, ori_shape, bg):
if not os.path.exists(self.args.save_dir):
os.makedirs(self.args.save_dir)
resize_w = pred.shape[-1]
resize_h = pred.shape[-2]
if self.args.soft_predict:
if self.args.use_optic_flow:
score_map = pred[:, 1, :, :].squeeze(0)
score_map = 255 * score_map
cur_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
cur_gray = cv2.resize(cur_gray, (resize_w, resize_h))
optflow_map = optic_flow_process(cur_gray, score_map, self.prev_gray, self.prev_cfd, \
self.disflow, self.is_init)
self.prev_gray = cur_gray.copy()
self.prev_cfd = optflow_map.copy()
self.is_init = False
score_map = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2)
score_map = np.transpose(score_map, [2, 0, 1])[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0),
[1, 2, 0]) / 255
else:
score_map = pred[:, 1, :, :]
score_map = score_map[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0), [1, 2, 0])
else:
if pred.ndim == 3:
pred = pred[:, np.newaxis, ...]
result = reverse_transform(
paddle.to_tensor(
pred, dtype='float32'),
ori_shape,
self.cfg.transforms,
mode='bilinear')
result = np.array(result)
if self.args.add_argmax:
result = np.argmax(result, axis=1)
else:
result = result.squeeze(1)
alpha = np.transpose(result, [1, 2, 0])
# background replace
h, w, _ = img.shape
if bg is None:
bg = np.ones_like(img)*255
else:
bg = cv2.resize(bg, (w, h))
if bg.ndim == 2:
bg = bg[..., np.newaxis]
comb = (alpha * img + (1 - alpha) * bg).astype(np.uint8)
return comb, alpha, bg, img
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment