update
This commit is contained in:
smthem
@@ -0,0 +1,4 @@
|
||||
|
||||
from .stable_makeup_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Face Morpher module init code
|
||||
"""
|
||||
from .morpher import morpher, list_imgpaths
|
||||
from .averager import averager
|
||||
|
||||
__all__ = ['list_imgpaths',
|
||||
'morpher',
|
||||
'averager']
|
||||
@@ -0,0 +1,133 @@
|
||||
"""
|
||||
Align face and image sizes
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def positive_cap(num):
|
||||
""" Cap a number to ensure positivity
|
||||
|
||||
:param num: positive or negative number
|
||||
:returns: (overflow, capped_number)
|
||||
"""
|
||||
if num < 0:
|
||||
return 0, abs(num)
|
||||
else:
|
||||
return num, 0
|
||||
|
||||
def roi_coordinates(rect, size, scale):
|
||||
""" Align the rectangle into the center and return the top-left coordinates
|
||||
within the new size. If rect is smaller, we add borders.
|
||||
|
||||
:param rect: (x, y, w, h) bounding rectangle of the face
|
||||
:param size: (width, height) are the desired dimensions
|
||||
:param scale: scaling factor of the rectangle to be resized
|
||||
:returns: 4 numbers. Top-left coordinates of the aligned ROI.
|
||||
(x, y, border_x, border_y). All values are > 0.
|
||||
"""
|
||||
rectx, recty, rectw, recth = rect
|
||||
new_height, new_width = size
|
||||
mid_x = int((rectx + rectw/2) * scale)
|
||||
mid_y = int((recty + recth/2) * scale)
|
||||
roi_x = mid_x - int(new_width/2)
|
||||
roi_y = mid_y - int(new_height/2)
|
||||
|
||||
roi_x, border_x = positive_cap(roi_x)
|
||||
roi_y, border_y = positive_cap(roi_y)
|
||||
return roi_x, roi_y, border_x, border_y
|
||||
|
||||
def scaling_factor(rect, size):
|
||||
""" Calculate the scaling factor for the current image to be
|
||||
resized to the new dimensions
|
||||
|
||||
:param rect: (x, y, w, h) bounding rectangle of the face
|
||||
:param size: (width, height) are the desired dimensions
|
||||
:returns: floating point scaling factor
|
||||
"""
|
||||
new_height, new_width = size
|
||||
rect_h, rect_w = rect[2:]
|
||||
height_ratio = rect_h / new_height
|
||||
width_ratio = rect_w / new_width
|
||||
scale = 1
|
||||
if height_ratio > width_ratio:
|
||||
new_recth = 0.8 * new_height
|
||||
scale = new_recth / rect_h
|
||||
else:
|
||||
new_rectw = 0.8 * new_width
|
||||
scale = new_rectw / rect_w
|
||||
return scale
|
||||
|
||||
def resize_image(img, scale):
|
||||
""" Resize image with the provided scaling factor
|
||||
|
||||
:param img: image to be resized
|
||||
:param scale: scaling factor for resizing the image
|
||||
"""
|
||||
cur_height, cur_width = img.shape[:2]
|
||||
new_scaled_height = int(scale * cur_height)
|
||||
new_scaled_width = int(scale * cur_width)
|
||||
|
||||
return cv2.resize(img, (new_scaled_width, new_scaled_height))
|
||||
|
||||
|
||||
def de_resize_image(img, scale):
|
||||
""" Resize image with the provided scaling factor
|
||||
|
||||
:param img: image to be resized
|
||||
:param scale: scaling factor for resizing the image
|
||||
"""
|
||||
cur_height, cur_width = img.shape[:2]
|
||||
new_scaled_height = int(cur_height / scale)
|
||||
new_scaled_width = int(cur_width / scale)
|
||||
|
||||
return cv2.resize(img, (new_scaled_width, new_scaled_height))
|
||||
|
||||
def resize_align(img, points, size):
|
||||
""" Resize image and associated points, align face to the center
|
||||
and crop to the desired size
|
||||
|
||||
:param img: image to be resized
|
||||
:param points: *m* x 2 array of points
|
||||
:param size: (height, width) tuple of new desired size
|
||||
"""
|
||||
new_height, new_width = size # 910 910
|
||||
|
||||
# Resize image based on bounding rectangle
|
||||
rect = cv2.boundingRect(np.array([points], np.int32)) # face rect
|
||||
scale = scaling_factor(rect, size)
|
||||
img = resize_image(img, scale)
|
||||
|
||||
# Align bounding rect to center
|
||||
cur_height, cur_width = img.shape[:2]
|
||||
roi_x, roi_y, border_x, border_y = roi_coordinates(rect, size, scale)
|
||||
roi_h = np.min([new_height-border_y, cur_height-roi_y])
|
||||
roi_w = np.min([new_width-border_x, cur_width-roi_x])
|
||||
|
||||
# Crop to supplied size
|
||||
crop = np.zeros((new_height, new_width, 3), img.dtype)
|
||||
crop[border_y:border_y+roi_h, border_x:border_x+roi_w] = (
|
||||
img[roi_y:roi_y+roi_h, roi_x:roi_x+roi_w])
|
||||
|
||||
# Scale and align face points to the crop
|
||||
points[:, 0] = (points[:, 0] * scale) + (border_x - roi_x)
|
||||
points[:, 1] = (points[:, 1] * scale) + (border_y - roi_y)
|
||||
|
||||
return (crop, points)
|
||||
|
||||
def restore_image(img, points, size, crop):
|
||||
|
||||
new_height, new_width = size # 910 910
|
||||
# Resize image based on bounding rectangle
|
||||
rect = cv2.boundingRect(np.array([points], np.int32)) # face rect
|
||||
scale = scaling_factor(rect, size)
|
||||
img = resize_image(img, scale)
|
||||
|
||||
# Align bounding rect to center
|
||||
cur_height, cur_width = img.shape[:2]
|
||||
roi_x, roi_y, border_x, border_y = roi_coordinates(rect, size, scale)
|
||||
roi_h = np.min([new_height-border_y, cur_height-roi_y])
|
||||
roi_w = np.min([new_width-border_x, cur_width-roi_x])
|
||||
|
||||
img[roi_y:roi_y + roi_h, roi_x:roi_x + roi_w] = crop[border_y:border_y + roi_h, border_x:border_x + roi_w]
|
||||
img = de_resize_image(img, scale)
|
||||
return img
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
::
|
||||
|
||||
Face averager
|
||||
|
||||
Usage:
|
||||
averager.py --images=<images_folder> [--blur] [--plot]
|
||||
[--background=(black|transparent|average)]
|
||||
[--width=<width>] [--height=<height>]
|
||||
[--out=<filename>] [--destimg=<filename>]
|
||||
|
||||
Options:
|
||||
-h, --help Show this screen.
|
||||
--images=<folder> Folder to images (.jpg, .jpeg, .png)
|
||||
--blur Flag to blur edges of image [default: False]
|
||||
--width=<width> Custom width of the images/video [default: 500]
|
||||
--height=<height> Custom height of the images/video [default: 600]
|
||||
--out=<filename> Filename to save the average face [default: result.png]
|
||||
--destimg=<filename> Destination face image to overlay average face
|
||||
--plot Flag to display the average face [default: False]
|
||||
--background=<bg> Background of image to be one of (black|transparent|average) [default: black]
|
||||
--version Show version.
|
||||
"""
|
||||
|
||||
from docopt import docopt
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.image as mpimg
|
||||
|
||||
from facemorpher import locator
|
||||
from facemorpher import aligner
|
||||
from facemorpher import warper
|
||||
from facemorpher import blender
|
||||
from facemorpher import plotter
|
||||
|
||||
# import locator
|
||||
# import aligner
|
||||
# import warper
|
||||
# import blender
|
||||
# import plotter
|
||||
|
||||
|
||||
def list_imgpaths(imgfolder):
|
||||
for fname in os.listdir(imgfolder):
|
||||
if (fname.lower().endswith('.jpg') or
|
||||
fname.lower().endswith('.png') or
|
||||
fname.lower().endswith('.jpeg')):
|
||||
yield os.path.join(imgfolder, fname)
|
||||
|
||||
def sharpen(img):
|
||||
blured = cv2.GaussianBlur(img, (0, 0), 2.5)
|
||||
return cv2.addWeighted(img, 1.4, blured, -0.4, 0)
|
||||
|
||||
def load_image_points(path, size):
|
||||
img = cv2.imread(path)
|
||||
points = locator.face_points(img)
|
||||
|
||||
if len(points) == 0:
|
||||
print('No face in %s' % path)
|
||||
return None, None
|
||||
else:
|
||||
return aligner.resize_align(img, points, size)
|
||||
|
||||
def averager(imgpaths, dest_filename=None, width=500, height=600, background='black',
|
||||
blur_edges=False, out_filename='result.png', plot=False):
|
||||
|
||||
size = (height, width)
|
||||
|
||||
images = []
|
||||
point_set = []
|
||||
for path in imgpaths:
|
||||
img, points = load_image_points(path, size)
|
||||
if img is not None:
|
||||
images.append(img)
|
||||
point_set.append(points)
|
||||
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError('Could not find any valid images.' +
|
||||
' Supported formats are .jpg, .png, .jpeg')
|
||||
|
||||
if dest_filename is not None:
|
||||
dest_img, dest_points = load_image_points(dest_filename, size)
|
||||
if dest_img is None or dest_points is None:
|
||||
raise Exception('No face or detected face points in dest img: ' + dest_filename)
|
||||
else:
|
||||
dest_img = np.zeros(images[0].shape, np.uint8)
|
||||
dest_points = locator.average_points(point_set)
|
||||
|
||||
num_images = len(images)
|
||||
result_images = np.zeros(images[0].shape, np.float32)
|
||||
for i in range(num_images):
|
||||
result_images += warper.warp_image(images[i], point_set[i],
|
||||
dest_points, size, np.float32)
|
||||
|
||||
result_image = np.uint8(result_images / num_images)
|
||||
face_indexes = np.nonzero(result_image)
|
||||
dest_img[face_indexes] = result_image[face_indexes]
|
||||
|
||||
mask = blender.mask_from_points(size, dest_points)
|
||||
if blur_edges:
|
||||
blur_radius = 10
|
||||
mask = cv2.blur(mask, (blur_radius, blur_radius))
|
||||
|
||||
if background in ('transparent', 'average'):
|
||||
dest_img = np.dstack((dest_img, mask))
|
||||
|
||||
if background == 'average':
|
||||
average_background = locator.average_points(images)
|
||||
dest_img = blender.overlay_image(dest_img, mask, average_background)
|
||||
|
||||
print('Averaged {} images'.format(num_images))
|
||||
cv2.imwrite(out_filename, dest_img)
|
||||
|
||||
def main():
|
||||
list_imgpaths = ["/share2/zhangyuxuan/project/reenactment_diffusion/FFHQ_format/test_face/warp.png", "/share2/zhangyuxuan/project/reenactment_diffusion/FFHQ_format/test_face/warp.png"]
|
||||
destimg = "/share2/zhangyuxuan/project/reenactment_diffusion/FFHQ_format/test_face/warp.png"
|
||||
width = 910
|
||||
height = 910
|
||||
background = 'transparent'
|
||||
blur = True
|
||||
out = "/share2/zhangyuxuan/project/reenactment_diffusion/FFHQ_format/test_face/1.jpg"
|
||||
plot = False
|
||||
|
||||
try:
|
||||
averager(list_imgpaths, destimg,
|
||||
width, height,
|
||||
background, blur, out, plot)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,133 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import scipy.sparse
|
||||
|
||||
def mask_from_points(size, points):
|
||||
""" Create a mask of supplied size from supplied points
|
||||
:param size: tuple of output mask size
|
||||
:param points: array of [x, y] points
|
||||
:returns: mask of values 0 and 255 where
|
||||
255 indicates the convex hull containing the points
|
||||
"""
|
||||
radius = 10 # kernel size
|
||||
kernel = np.ones((radius, radius), np.uint8)
|
||||
|
||||
mask = np.zeros(size, np.uint8)
|
||||
cv2.fillConvexPoly(mask, cv2.convexHull(points), 255)
|
||||
mask = cv2.erode(mask, kernel)
|
||||
|
||||
return mask
|
||||
|
||||
def overlay_image(foreground_image, mask, background_image):
|
||||
""" Overlay foreground image onto the background given a mask
|
||||
:param foreground_image: foreground image points
|
||||
:param mask: [0-255] values in mask
|
||||
:param background_image: background image points
|
||||
:returns: image with foreground where mask > 0 overlaid on background image
|
||||
"""
|
||||
foreground_pixels = mask > 0
|
||||
background_image[..., :3][foreground_pixels] = foreground_image[..., :3][foreground_pixels]
|
||||
return background_image
|
||||
|
||||
def apply_mask(img, mask):
|
||||
""" Apply mask to supplied image
|
||||
:param img: max 3 channel image
|
||||
:param mask: [0-255] values in mask
|
||||
:returns: new image with mask applied
|
||||
"""
|
||||
masked_img = np.copy(img)
|
||||
num_channels = 3
|
||||
for c in range(num_channels):
|
||||
masked_img[..., c] = img[..., c] * (mask / 255)
|
||||
|
||||
return masked_img
|
||||
|
||||
def weighted_average(img1, img2, percent=0.5):
|
||||
if percent <= 0:
|
||||
return img2
|
||||
elif percent >= 1:
|
||||
return img1
|
||||
else:
|
||||
return cv2.addWeighted(img1, percent, img2, 1-percent, 0)
|
||||
|
||||
def alpha_feathering(src_img, dest_img, img_mask, blur_radius=15):
|
||||
mask = cv2.blur(img_mask, (blur_radius, blur_radius))
|
||||
mask = mask / 255.0
|
||||
|
||||
result_img = np.empty(src_img.shape, np.uint8)
|
||||
for i in range(3):
|
||||
result_img[..., i] = src_img[..., i] * mask + dest_img[..., i] * (1-mask)
|
||||
|
||||
return result_img
|
||||
|
||||
def poisson_blend(img_source, dest_img, img_mask, offset=(0, 0)):
|
||||
# http://opencv.jp/opencv2-x-samples/poisson-blending
|
||||
img_target = np.copy(dest_img)
|
||||
import pyamg
|
||||
# compute regions to be blended
|
||||
region_source = (
|
||||
max(-offset[0], 0),
|
||||
max(-offset[1], 0),
|
||||
min(img_target.shape[0] - offset[0], img_source.shape[0]),
|
||||
min(img_target.shape[1] - offset[1], img_source.shape[1]))
|
||||
region_target = (
|
||||
max(offset[0], 0),
|
||||
max(offset[1], 0),
|
||||
min(img_target.shape[0], img_source.shape[0] + offset[0]),
|
||||
min(img_target.shape[1], img_source.shape[1] + offset[1]))
|
||||
region_size = (region_source[2] - region_source[0],
|
||||
region_source[3] - region_source[1])
|
||||
|
||||
# clip and normalize mask image
|
||||
img_mask = img_mask[region_source[0]:region_source[2],
|
||||
region_source[1]:region_source[3]]
|
||||
|
||||
# create coefficient matrix
|
||||
coff_mat = scipy.sparse.identity(np.prod(region_size), format='lil')
|
||||
for y in range(region_size[0]):
|
||||
for x in range(region_size[1]):
|
||||
if img_mask[y, x]:
|
||||
index = x + y * region_size[1]
|
||||
coff_mat[index, index] = 4
|
||||
if index + 1 < np.prod(region_size):
|
||||
coff_mat[index, index + 1] = -1
|
||||
if index - 1 >= 0:
|
||||
coff_mat[index, index - 1] = -1
|
||||
if index + region_size[1] < np.prod(region_size):
|
||||
coff_mat[index, index + region_size[1]] = -1
|
||||
if index - region_size[1] >= 0:
|
||||
coff_mat[index, index - region_size[1]] = -1
|
||||
coff_mat = coff_mat.tocsr()
|
||||
|
||||
# create poisson matrix for b
|
||||
poisson_mat = pyamg.gallery.poisson(img_mask.shape)
|
||||
# for each layer (ex. RGB)
|
||||
for num_layer in range(img_target.shape[2]):
|
||||
# get subimages
|
||||
t = img_target[region_target[0]:region_target[2],
|
||||
region_target[1]:region_target[3], num_layer]
|
||||
s = img_source[region_source[0]:region_source[2],
|
||||
region_source[1]:region_source[3], num_layer]
|
||||
t = t.flatten()
|
||||
s = s.flatten()
|
||||
|
||||
# create b
|
||||
b = poisson_mat * s
|
||||
for y in range(region_size[0]):
|
||||
for x in range(region_size[1]):
|
||||
if not img_mask[y, x]:
|
||||
index = x + y * region_size[1]
|
||||
b[index] = t[index]
|
||||
|
||||
# solve Ax = b
|
||||
x = pyamg.solve(coff_mat, b, verb=False, tol=1e-10)
|
||||
|
||||
# assign x to target image
|
||||
x = np.reshape(x, region_size)
|
||||
x[x > 255] = 255
|
||||
x[x < 0] = 0
|
||||
x = np.array(x, img_target.dtype)
|
||||
img_target[region_target[0]:region_target[2],
|
||||
region_target[1]:region_target[3], num_layer] = x
|
||||
|
||||
return img_target
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
Locate face points
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os.path as path
|
||||
import dlib
|
||||
import os
|
||||
|
||||
|
||||
DATA_DIR = "/share2/zhangyuxuan/project/reenactment_diffusion/FFHQ_format"
|
||||
# DATA_DIR = "/mnt/nj-aigc/usr/zhangyuxuan1/project/FFHQ_format"
|
||||
|
||||
dlib_detector = dlib.get_frontal_face_detector()
|
||||
dlib_predictor = dlib.shape_predictor(path.join(DATA_DIR, 'shape_predictor_68_face_landmarks.dat'))
|
||||
|
||||
def boundary_points(points, width_percent=0.1, height_percent=0.1):
|
||||
""" Produce additional boundary points
|
||||
:param points: *m* x 2 array of x,y points
|
||||
:param width_percent: [-1, 1] percentage of width to taper inwards. Negative for opposite direction
|
||||
:param height_percent: [-1, 1] percentage of height to taper downwards. Negative for opposite direction
|
||||
:returns: 2 additional points at the top corners
|
||||
"""
|
||||
x, y, w, h = cv2.boundingRect(np.array([points], np.int32))
|
||||
spacerw = int(w * width_percent)
|
||||
spacerh = int(h * height_percent)
|
||||
return [[x+spacerw, y+spacerh],
|
||||
[x+w-spacerw, y+spacerh]]
|
||||
|
||||
|
||||
def face_points(img, add_boundary_points=True):
|
||||
return face_points_dlib(img, add_boundary_points)
|
||||
|
||||
def face_points_dlib(img, add_boundary_points=True):
|
||||
""" Locates 68 face points using dlib (http://dlib.net)
|
||||
Requires shape_predictor_68_face_landmarks.dat to be in face_morpher/data
|
||||
Download at: http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
|
||||
:param img: an image array
|
||||
:param add_boundary_points: bool to add additional boundary points
|
||||
:returns: Array of x,y face points. Empty array if no face found
|
||||
"""
|
||||
try:
|
||||
points = []
|
||||
rgbimg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
rects = dlib_detector(rgbimg, 1)
|
||||
|
||||
if rects and len(rects) > 0:
|
||||
# We only take the first found face
|
||||
shapes = dlib_predictor(rgbimg, rects[0])
|
||||
points = np.array([(shapes.part(i).x, shapes.part(i).y) for i in range(68)], np.int32)
|
||||
|
||||
if add_boundary_points:
|
||||
# Add more points inwards and upwards as dlib only detects up to eyebrows
|
||||
points = np.vstack([
|
||||
points,
|
||||
boundary_points(points, 0.1, -0.03),
|
||||
boundary_points(points, 0.13, -0.05),
|
||||
boundary_points(points, 0.15, -0.08),
|
||||
boundary_points(points, 0.33, -0.12)])
|
||||
|
||||
return points
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return []
|
||||
|
||||
def face_points_stasm(img, add_boundary_points=True):
|
||||
import stasm
|
||||
""" Locates 77 face points using stasm (http://www.milbo.users.sonic.net/stasm)
|
||||
|
||||
:param img: an image array
|
||||
:param add_boundary_points: bool to add 2 additional points
|
||||
:returns: Array of x,y face points. Empty array if no face found
|
||||
"""
|
||||
try:
|
||||
points = stasm.search_single(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
|
||||
except Exception as e:
|
||||
print('Failed finding face points: ', e)
|
||||
return []
|
||||
|
||||
points = points.astype(np.int32)
|
||||
if len(points) == 0:
|
||||
return points
|
||||
|
||||
if add_boundary_points:
|
||||
return np.vstack([points, boundary_points(points)])
|
||||
|
||||
return points
|
||||
|
||||
def average_points(point_set):
|
||||
""" Averages a set of face points from images
|
||||
|
||||
:param point_set: *n* x *m* x 2 array of face points. \\
|
||||
*n* = number of images. *m* = number of face points per image
|
||||
"""
|
||||
return np.mean(point_set, 0).astype(np.int32)
|
||||
|
||||
def weighted_average_points(start_points, end_points, percent=0.5):
|
||||
""" Weighted average of two sets of supplied points
|
||||
|
||||
:param start_points: *m* x 2 array of start face points.
|
||||
:param end_points: *m* x 2 array of end face points.
|
||||
:param percent: [0, 1] percentage weight on start_points
|
||||
:returns: *m* x 2 array of weighted average points
|
||||
"""
|
||||
if percent <= 0:
|
||||
return end_points
|
||||
elif percent >= 1:
|
||||
return start_points
|
||||
else:
|
||||
return np.asarray(start_points*percent + end_points*(1-percent), np.int32)
|
||||
@@ -0,0 +1,153 @@
|
||||
"""
|
||||
::
|
||||
|
||||
Morph from source to destination face or
|
||||
Morph through all images in a folder
|
||||
|
||||
Usage:
|
||||
morpher.py (--src=<src_path> --dest=<dest_path> | --images=<folder>)
|
||||
[--width=<width>] [--height=<height>]
|
||||
[--num=<num_frames>] [--fps=<frames_per_second>]
|
||||
[--out_frames=<folder>] [--out_video=<filename>]
|
||||
[--plot] [--background=(black|transparent|average)]
|
||||
|
||||
Options:
|
||||
-h, --help Show this screen.
|
||||
--src=<src_imgpath> Filepath to source image (.jpg, .jpeg, .png)
|
||||
--dest=<dest_imgpath> Filepath to destination image (.jpg, .jpeg, .png)
|
||||
--images=<folder> Folderpath to images
|
||||
--width=<width> Custom width of the images/video [default: 500]
|
||||
--height=<height> Custom height of the images/video [default: 600]
|
||||
--num=<num_frames> Number of morph frames [default: 20]
|
||||
--fps=<fps> Number frames per second for the video [default: 10]
|
||||
--out_frames=<folder> Folder path to save all image frames
|
||||
--out_video=<filename> Filename to save a video
|
||||
--plot Flag to plot images to result.png [default: False]
|
||||
--background=<bg> Background of images to be one of (black|transparent|average) [default: black]
|
||||
--version Show version.
|
||||
"""
|
||||
from docopt import docopt
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from facemorpher import locator
|
||||
from facemorpher import aligner
|
||||
from facemorpher import warper
|
||||
from facemorpher import blender
|
||||
from facemorpher import plotter
|
||||
from facemorpher import videoer
|
||||
|
||||
def verify_args(args):
|
||||
if args['--images'] is None:
|
||||
valid = os.path.isfile(args['--src']) & os.path.isfile(args['--dest'])
|
||||
if not valid:
|
||||
print('--src=%s or --dest=%s file does not exist. Double check the supplied paths' % (
|
||||
args['--src'], args['--dest']))
|
||||
exit(1)
|
||||
else:
|
||||
valid = os.path.isdir(args['--images'])
|
||||
if not valid:
|
||||
print('--images=%s is not a valid directory' % args['--images'])
|
||||
exit(1)
|
||||
|
||||
def load_image_points(path, size):
|
||||
img = cv2.imread(path)
|
||||
points = locator.face_points(img)
|
||||
|
||||
if len(points) == 0:
|
||||
print('No face in %s' % path)
|
||||
return None, None
|
||||
else:
|
||||
return aligner.resize_align(img, points, size)
|
||||
|
||||
def load_valid_image_points(imgpaths, size):
|
||||
for path in imgpaths:
|
||||
img, points = load_image_points(path, size)
|
||||
if img is not None:
|
||||
print(path)
|
||||
yield (img, points)
|
||||
|
||||
def list_imgpaths(images_folder=None, src_image=None, dest_image=None):
|
||||
if images_folder is None:
|
||||
yield src_image
|
||||
yield dest_image
|
||||
else:
|
||||
for fname in os.listdir(images_folder):
|
||||
if (fname.lower().endswith('.jpg') or
|
||||
fname.lower().endswith('.png') or
|
||||
fname.lower().endswith('.jpeg')):
|
||||
yield os.path.join(images_folder, fname)
|
||||
|
||||
def morph(src_img, src_points, dest_img, dest_points,
|
||||
video, width=500, height=600, num_frames=20, fps=10,
|
||||
out_frames=None, out_video=None, plot=False, background='black'):
|
||||
"""
|
||||
Create a morph sequence from source to destination image
|
||||
|
||||
:param src_img: ndarray source image
|
||||
:param src_points: source image array of x,y face points
|
||||
:param dest_img: ndarray destination image
|
||||
:param dest_points: destination image array of x,y face points
|
||||
:param video: facemorpher.videoer.Video object
|
||||
"""
|
||||
size = (height, width)
|
||||
stall_frames = np.clip(int(fps*0.15), 1, fps) # Show first & last longer
|
||||
plt = plotter.Plotter(plot, num_images=num_frames, out_folder=out_frames)
|
||||
num_frames -= (stall_frames * 2) # No need to process src and dest image
|
||||
|
||||
plt.plot_one(src_img)
|
||||
video.write(src_img, 1)
|
||||
|
||||
# Produce morph frames!
|
||||
for percent in np.linspace(1, 0, num=num_frames):
|
||||
points = locator.weighted_average_points(src_points, dest_points, percent)
|
||||
src_face = warper.warp_image(src_img, src_points, points, size)
|
||||
end_face = warper.warp_image(dest_img, dest_points, points, size)
|
||||
average_face = blender.weighted_average(src_face, end_face, percent)
|
||||
|
||||
if background in ('transparent', 'average'):
|
||||
mask = blender.mask_from_points(average_face.shape[:2], points)
|
||||
average_face = np.dstack((average_face, mask))
|
||||
|
||||
if background == 'average':
|
||||
average_background = blender.weighted_average(src_img, dest_img, percent)
|
||||
average_face = blender.overlay_image(average_face, mask, average_background)
|
||||
|
||||
plt.plot_one(average_face)
|
||||
plt.save(average_face)
|
||||
video.write(average_face)
|
||||
|
||||
plt.plot_one(dest_img)
|
||||
video.write(dest_img, stall_frames)
|
||||
plt.show()
|
||||
|
||||
def morpher(imgpaths, width=500, height=600, num_frames=20, fps=10,
|
||||
out_frames=None, out_video=None, plot=False, background='black'):
|
||||
"""
|
||||
Create a morph sequence from multiple images in imgpaths
|
||||
|
||||
:param imgpaths: array or generator of image paths
|
||||
"""
|
||||
video = videoer.Video(out_video, fps, width, height)
|
||||
images_points_gen = load_valid_image_points(imgpaths, (height, width))
|
||||
src_img, src_points = next(images_points_gen)
|
||||
for dest_img, dest_points in images_points_gen:
|
||||
morph(src_img, src_points, dest_img, dest_points, video,
|
||||
width, height, num_frames, fps, out_frames, out_video, plot, background)
|
||||
src_img, src_points = dest_img, dest_points
|
||||
video.end()
|
||||
|
||||
def main():
|
||||
args = docopt(__doc__, version='Face Morpher 1.0')
|
||||
verify_args(args)
|
||||
|
||||
morpher(list_imgpaths(args['--images'], args['--src'], args['--dest']),
|
||||
int(args['--width']), int(args['--height']),
|
||||
int(args['--num']), int(args['--fps']),
|
||||
args['--out_frames'], args['--out_video'],
|
||||
args['--plot'], args['--background'])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
Plot and save images
|
||||
"""
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.image as mpimg
|
||||
import os.path
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def bgr2rgb(img):
|
||||
# OpenCV's BGR to RGB
|
||||
rgb = np.copy(img)
|
||||
rgb[..., 0], rgb[..., 2] = img[..., 2], img[..., 0]
|
||||
return rgb
|
||||
|
||||
def check_do_plot(func):
|
||||
def inner(self, *args, **kwargs):
|
||||
if self.do_plot:
|
||||
func(self, *args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
def check_do_save(func):
|
||||
def inner(self, *args, **kwargs):
|
||||
if self.do_save:
|
||||
func(self, *args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
class Plotter(object):
|
||||
def __init__(self, plot=True, rows=0, cols=0, num_images=0, out_folder=None, out_filename=None):
|
||||
self.save_counter = 1
|
||||
self.plot_counter = 1
|
||||
self.do_plot = plot
|
||||
self.do_save = out_filename is not None
|
||||
self.out_filename = out_filename
|
||||
self.set_filepath(out_folder)
|
||||
|
||||
if (rows + cols) == 0 and num_images > 0:
|
||||
# Auto-calculate the number of rows and cols for the figure
|
||||
self.rows = np.ceil(np.sqrt(num_images / 2.0))
|
||||
self.cols = np.ceil(num_images / self.rows)
|
||||
else:
|
||||
self.rows = rows
|
||||
self.cols = cols
|
||||
|
||||
def set_filepath(self, folder):
|
||||
if folder is None:
|
||||
self.filepath = None
|
||||
return
|
||||
|
||||
if not os.path.exists(folder):
|
||||
os.makedirs(folder)
|
||||
self.filepath = os.path.join(folder, 'frame{0:03d}.png')
|
||||
self.do_save = True
|
||||
|
||||
@check_do_save
|
||||
def save(self, img, filename=None):
|
||||
if self.filepath:
|
||||
filename = self.filepath.format(self.save_counter)
|
||||
self.save_counter += 1
|
||||
elif filename is None:
|
||||
filename = self.out_filename
|
||||
|
||||
mpimg.imsave(filename, bgr2rgb(img))
|
||||
print(filename + ' saved')
|
||||
|
||||
@check_do_plot
|
||||
def plot_one(self, img):
|
||||
p = plt.subplot(self.rows, self.cols, self.plot_counter)
|
||||
p.axes.get_xaxis().set_visible(False)
|
||||
p.axes.get_yaxis().set_visible(False)
|
||||
plt.imshow(bgr2rgb(img))
|
||||
self.plot_counter += 1
|
||||
|
||||
@check_do_plot
|
||||
def show(self):
|
||||
plt.gcf().subplots_adjust(hspace=0.05, wspace=0,
|
||||
left=0, bottom=0, right=1, top=0.98)
|
||||
plt.axis('off')
|
||||
#plt.show()
|
||||
plt.savefig('result.png')
|
||||
|
||||
@check_do_plot
|
||||
def plot_mesh(self, points, tri, color='k'):
|
||||
""" plot triangles """
|
||||
for tri_indices in tri.simplices:
|
||||
t_ext = [tri_indices[0], tri_indices[1], tri_indices[2], tri_indices[0]]
|
||||
plt.plot(points[t_ext, 0], points[t_ext, 1], color)
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Create a video with image frames
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def check_write_video(func):
|
||||
def inner(self, *args, **kwargs):
|
||||
if self.video:
|
||||
return func(self, *args, **kwargs)
|
||||
else:
|
||||
pass
|
||||
return inner
|
||||
|
||||
|
||||
class Video(object):
|
||||
def __init__(self, filename, fps, w, h):
|
||||
self.filename = filename
|
||||
|
||||
if filename is None:
|
||||
self.video = None
|
||||
else:
|
||||
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
|
||||
self.video = cv2.VideoWriter(filename, fourcc, fps, (w, h), True)
|
||||
|
||||
@check_write_video
|
||||
def write(self, img, num_times=1):
|
||||
for i in range(num_times):
|
||||
self.video.write(img[..., :3])
|
||||
|
||||
@check_write_video
|
||||
def end(self):
|
||||
print(self.filename + ' saved')
|
||||
self.video.release()
|
||||
@@ -0,0 +1,139 @@
|
||||
import numpy as np
|
||||
import scipy.spatial as spatial
|
||||
|
||||
def bilinear_interpolate(img, coords):
|
||||
""" Interpolates over every image channel
|
||||
http://en.wikipedia.org/wiki/Bilinear_interpolation
|
||||
|
||||
:param img: max 3 channel image
|
||||
:param coords: 2 x _m_ array. 1st row = xcoords, 2nd row = ycoords
|
||||
:returns: array of interpolated pixels with same shape as coords
|
||||
"""
|
||||
int_coords = np.int32(coords)
|
||||
x0, y0 = int_coords
|
||||
dx, dy = coords - int_coords
|
||||
|
||||
# 4 Neighour pixels
|
||||
q11 = img[y0, x0]
|
||||
q21 = img[y0, x0+1]
|
||||
q12 = img[y0+1, x0]
|
||||
q22 = img[y0+1, x0+1]
|
||||
|
||||
btm = q21.T * dx + q11.T * (1 - dx)
|
||||
top = q22.T * dx + q12.T * (1 - dx)
|
||||
inter_pixel = top * dy + btm * (1 - dy)
|
||||
|
||||
return inter_pixel.T
|
||||
|
||||
def grid_coordinates(points):
|
||||
""" x,y grid coordinates within the ROI of supplied points
|
||||
|
||||
:param points: points to generate grid coordinates
|
||||
:returns: array of (x, y) coordinates
|
||||
"""
|
||||
xmin = np.min(points[:, 0])
|
||||
xmax = np.max(points[:, 0]) + 1
|
||||
ymin = np.min(points[:, 1])
|
||||
ymax = np.max(points[:, 1]) + 1
|
||||
return np.asarray([(x, y) for y in range(ymin, ymax)
|
||||
for x in range(xmin, xmax)], np.uint32)
|
||||
|
||||
def process_warp(src_img, result_img, tri_affines, dst_points, delaunay):
|
||||
"""
|
||||
Warp each triangle from the src_image only within the
|
||||
ROI of the destination image (points in dst_points).
|
||||
"""
|
||||
roi_coords = grid_coordinates(dst_points)
|
||||
# indices to vertices. -1 if pixel is not in any triangle
|
||||
roi_tri_indices = delaunay.find_simplex(roi_coords)
|
||||
|
||||
for simplex_index in range(len(delaunay.simplices)):
|
||||
coords = roi_coords[roi_tri_indices == simplex_index]
|
||||
num_coords = len(coords)
|
||||
out_coords = np.dot(tri_affines[simplex_index],
|
||||
np.vstack((coords.T, np.ones(num_coords))))
|
||||
x, y = coords.T
|
||||
result_img[y, x] = bilinear_interpolate(src_img, out_coords)
|
||||
|
||||
return None
|
||||
|
||||
def triangular_affine_matrices(vertices, src_points, dest_points):
|
||||
"""
|
||||
Calculate the affine transformation matrix for each
|
||||
triangle (x,y) vertex from dest_points to src_points
|
||||
|
||||
:param vertices: array of triplet indices to corners of triangle
|
||||
:param src_points: array of [x, y] points to landmarks for source image
|
||||
:param dest_points: array of [x, y] points to landmarks for destination image
|
||||
:returns: 2 x 3 affine matrix transformation for a triangle
|
||||
"""
|
||||
ones = [1, 1, 1]
|
||||
for tri_indices in vertices:
|
||||
src_tri = np.vstack((src_points[tri_indices, :].T, ones))
|
||||
dst_tri = np.vstack((dest_points[tri_indices, :].T, ones))
|
||||
mat = np.dot(src_tri, np.linalg.inv(dst_tri))[:2, :]
|
||||
yield mat
|
||||
|
||||
def warp_image(src_img, src_points, dest_points, dest_shape, dtype=np.uint8):
|
||||
# Resultant image will not have an alpha channel
|
||||
num_chans = 3
|
||||
src_img = src_img[:, :, :3]
|
||||
|
||||
rows, cols = dest_shape[:2]
|
||||
result_img = np.zeros((rows, cols, num_chans), dtype)
|
||||
|
||||
delaunay = spatial.Delaunay(dest_points)
|
||||
tri_affines = np.asarray(list(triangular_affine_matrices(
|
||||
delaunay.simplices, src_points, dest_points)))
|
||||
|
||||
process_warp(src_img, result_img, tri_affines, dest_points, delaunay)
|
||||
|
||||
return result_img
|
||||
|
||||
def test_local():
|
||||
from functools import partial
|
||||
import cv2
|
||||
import scipy.misc
|
||||
import locator
|
||||
import aligner
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
# Load source image
|
||||
face_points_func = partial(locator.face_points, '../data')
|
||||
base_path = '../females/Screenshot 2015-03-04 17.11.12.png'
|
||||
src_path = '../females/BlDmB5QCYAAY8iw.jpg'
|
||||
src_img = cv2.imread(src_path)
|
||||
|
||||
# Define control points for warps
|
||||
src_points = face_points_func(src_path)
|
||||
base_img = cv2.imread(base_path)
|
||||
base_points = face_points_func(base_path)
|
||||
|
||||
size = (600, 500)
|
||||
src_img, src_points = aligner.resize_align(src_img, src_points, size)
|
||||
base_img, base_points = aligner.resize_align(base_img, base_points, size)
|
||||
result_points = locator.weighted_average_points(src_points, base_points, 0.2)
|
||||
|
||||
# Perform transform
|
||||
dst_img1 = warp_image(src_img, src_points, result_points, size)
|
||||
dst_img2 = warp_image(base_img, base_points, result_points, size)
|
||||
|
||||
import blender
|
||||
ave = blender.weighted_average(dst_img1, dst_img2, 0.6)
|
||||
mask = blender.mask_from_points(size, result_points)
|
||||
blended_img = blender.poisson_blend(dst_img1, dst_img2, mask)
|
||||
|
||||
plt.subplot(2, 2, 1)
|
||||
plt.imshow(ave)
|
||||
plt.subplot(2, 2, 2)
|
||||
plt.imshow(dst_img1)
|
||||
plt.subplot(2, 2, 3)
|
||||
plt.imshow(dst_img2)
|
||||
plt.subplot(2, 2, 4)
|
||||
|
||||
plt.imshow(blended_img)
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_local()
|
||||
@@ -0,0 +1,26 @@
|
||||
###############
|
||||
# conststants #
|
||||
###############
|
||||
DEFAULT_TARGET_GUIDANCE_SCALE = 15
|
||||
DEFAULT_SOURCE_GUIDANCE_SCALE = 3.5
|
||||
DEFAULT_DIFFUSION_STEPS = 100
|
||||
DEFAULT_SKIP_STEPS = 36
|
||||
DEFAULT_SEED = 0
|
||||
|
||||
|
||||
DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE = 7
|
||||
DEFAULT_WARMUP_STEPS = 2
|
||||
DEFAULT_THRESHOLD = 0.95
|
||||
DEFAULT_NEGATIVE_GUIDANCE=False
|
||||
|
||||
STYLE_SEGA_CONCEPT_GUIDANCE_SCALE = 8
|
||||
STYLE_WARMUP_STEPS = 2
|
||||
STYLE_THRESHOLD = 0.92
|
||||
|
||||
FACE_SEGA_CONCEPT_GUIDANCE_SCALE = 5
|
||||
FACE_WARMUP_STEPS = 2
|
||||
FACE_THRESHOLD = 0.95
|
||||
|
||||
OBJECT_SEGA_CONCEPT_GUIDANCE_SCALE = 15
|
||||
OBJECT_WARMUP_STEPS = 5
|
||||
OBJECT_THRESHOLD = 0.95
|
||||
@@ -0,0 +1,275 @@
|
||||
import torch
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
from PIL import Image, ImageDraw ,ImageFont
|
||||
from matplotlib import pyplot as plt
|
||||
import torchvision.transforms as T
|
||||
import os
|
||||
import yaml
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_512(image_path, left=0, right=0, top=0, bottom=0, device=None):
|
||||
if type(image_path) is str:
|
||||
image = np.array(Image.open(image_path).convert('RGB'))[:, :, :3]
|
||||
else:
|
||||
image = image_path
|
||||
h, w, c = image.shape
|
||||
left = min(left, w-1)
|
||||
right = min(right, w - left - 1)
|
||||
top = min(top, h - left - 1)
|
||||
bottom = min(bottom, h - top - 1)
|
||||
image = image[top:h-bottom, left:w-right]
|
||||
h, w, c = image.shape
|
||||
if h < w:
|
||||
offset = (w - h) // 2
|
||||
image = image[:, offset:offset + h]
|
||||
elif w < h:
|
||||
offset = (h - w) // 2
|
||||
image = image[offset:offset + w]
|
||||
image = np.array(Image.fromarray(image).resize((512, 512)))
|
||||
image = torch.from_numpy(image).float() / 127.5 - 1
|
||||
image = image.permute(2, 0, 1).unsqueeze(0).to(device, dtype =torch.float32)
|
||||
|
||||
return image
|
||||
|
||||
|
||||
|
||||
def mu_tilde(model, xt,x0, timestep):
|
||||
"mu_tilde(x_t, x_0) DDPM paper eq. 7"
|
||||
prev_timestep = timestep - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps
|
||||
alpha_prod_t_prev = model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else model.scheduler.final_alpha_cumprod
|
||||
alpha_t = model.scheduler.alphas[timestep]
|
||||
beta_t = 1 - alpha_t
|
||||
alpha_bar = model.scheduler.alphas_cumprod[timestep]
|
||||
return ((alpha_prod_t_prev ** 0.5 * beta_t) / (1-alpha_bar)) * x0 + ((alpha_t**0.5 *(1-alpha_prod_t_prev)) / (1- alpha_bar))*xt
|
||||
|
||||
def sample_xts_from_x0(model, x0, num_inference_steps=50):
|
||||
"""
|
||||
Samples from P(x_1:T|x_0)
|
||||
"""
|
||||
# torch.manual_seed(43256465436)
|
||||
alpha_bar = model.scheduler.alphas_cumprod
|
||||
sqrt_one_minus_alpha_bar = (1-alpha_bar) ** 0.5
|
||||
alphas = model.scheduler.alphas
|
||||
betas = 1 - alphas
|
||||
variance_noise_shape = (
|
||||
num_inference_steps,
|
||||
model.unet.in_channels,
|
||||
model.unet.sample_size,
|
||||
model.unet.sample_size)
|
||||
|
||||
timesteps = model.scheduler.timesteps.to(model.device)
|
||||
t_to_idx = {int(v):k for k,v in enumerate(timesteps)}
|
||||
xts = torch.zeros(variance_noise_shape).to(x0.device, dtype =torch.float16)
|
||||
for t in reversed(timesteps):
|
||||
idx = t_to_idx[int(t)]
|
||||
xts[idx] = x0 * (alpha_bar[t] ** 0.5) + torch.randn_like(x0, dtype =torch.float16) * sqrt_one_minus_alpha_bar[t]
|
||||
xts = torch.cat([xts, x0 ],dim = 0)
|
||||
|
||||
return xts
|
||||
|
||||
def encode_text(model, prompts):
|
||||
text_input = model.tokenizer(
|
||||
prompts,
|
||||
padding="max_length",
|
||||
max_length=model.tokenizer.model_max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
with torch.no_grad():
|
||||
text_encoding = model.text_encoder(text_input.input_ids.to(model.device))[0]
|
||||
return text_encoding
|
||||
|
||||
def forward_step(model, model_output, timestep, sample):
|
||||
next_timestep = min(model.scheduler.config.num_train_timesteps - 2,
|
||||
timestep + model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps)
|
||||
|
||||
# 2. compute alphas, betas
|
||||
alpha_prod_t = model.scheduler.alphas_cumprod[timestep]
|
||||
# alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep] if next_ltimestep >= 0 else self.scheduler.final_alpha_cumprod
|
||||
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
|
||||
# 3. compute predicted original sample from predicted noise also called
|
||||
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
|
||||
|
||||
# 5. TODO: simple noising implementatiom
|
||||
next_sample = model.scheduler.add_noise(pred_original_sample,
|
||||
model_output,
|
||||
torch.LongTensor([next_timestep]))
|
||||
return next_sample
|
||||
|
||||
|
||||
def get_variance(model, timestep): #, prev_timestep):
|
||||
prev_timestep = timestep - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps
|
||||
alpha_prod_t = model.scheduler.alphas_cumprod[timestep]
|
||||
alpha_prod_t_prev = model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else model.scheduler.final_alpha_cumprod
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
beta_prod_t_prev = 1 - alpha_prod_t_prev
|
||||
variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
|
||||
return variance
|
||||
|
||||
def inversion_forward_process(model, x0,
|
||||
etas = None,
|
||||
prog_bar = False,
|
||||
prompt = "",
|
||||
cfg_scale = 3.5,
|
||||
num_inference_steps=50, eps = None):
|
||||
|
||||
if not prompt=="":
|
||||
text_embeddings = encode_text(model, prompt)
|
||||
uncond_embedding = encode_text(model, "")
|
||||
timesteps = model.scheduler.timesteps.to(model.device)
|
||||
variance_noise_shape = (
|
||||
num_inference_steps,
|
||||
model.unet.in_channels,
|
||||
model.unet.sample_size,
|
||||
model.unet.sample_size)
|
||||
if etas is None or (type(etas) in [int, float] and etas == 0):
|
||||
eta_is_zero = True
|
||||
zs = None
|
||||
else:
|
||||
eta_is_zero = False
|
||||
if type(etas) in [int, float]: etas = [etas]*model.scheduler.num_inference_steps
|
||||
xts = sample_xts_from_x0(model, x0, num_inference_steps=num_inference_steps)
|
||||
alpha_bar = model.scheduler.alphas_cumprod
|
||||
zs = torch.zeros(size=variance_noise_shape, device=model.device, dtype =torch.float16)
|
||||
|
||||
t_to_idx = {int(v):k for k,v in enumerate(timesteps)}
|
||||
xt = x0
|
||||
op = tqdm(reversed(timesteps), desc= "Inverting...") if prog_bar else reversed(timesteps)
|
||||
|
||||
for t in op:
|
||||
idx = t_to_idx[int(t)]
|
||||
# 1. predict noise residual
|
||||
if not eta_is_zero:
|
||||
xt = xts[idx][None]
|
||||
|
||||
with torch.no_grad():
|
||||
out = model.unet.forward(xt, timestep = t, encoder_hidden_states = uncond_embedding)
|
||||
if not prompt=="":
|
||||
cond_out = model.unet.forward(xt, timestep=t, encoder_hidden_states = text_embeddings)
|
||||
|
||||
if not prompt=="":
|
||||
## classifier free guidance
|
||||
noise_pred = out.sample + cfg_scale * (cond_out.sample - out.sample)
|
||||
else:
|
||||
noise_pred = out.sample
|
||||
|
||||
if eta_is_zero:
|
||||
# 2. compute more noisy image and set x_t -> x_t+1
|
||||
xt = forward_step(model, noise_pred, t, xt)
|
||||
|
||||
else:
|
||||
xtm1 = xts[idx+1][None]
|
||||
# pred of x0
|
||||
pred_original_sample = (xt - (1-alpha_bar[t]) ** 0.5 * noise_pred ) / alpha_bar[t] ** 0.5
|
||||
|
||||
# direction to xt
|
||||
prev_timestep = t - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps
|
||||
alpha_prod_t_prev = model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else model.scheduler.final_alpha_cumprod
|
||||
|
||||
variance = get_variance(model, t)
|
||||
pred_sample_direction = (1 - alpha_prod_t_prev - etas[idx] * variance ) ** (0.5) * noise_pred
|
||||
|
||||
mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
|
||||
|
||||
z = (xtm1 - mu_xt ) / ( etas[idx] * variance ** 0.5 )
|
||||
zs[idx] = z
|
||||
|
||||
# correction to avoid error accumulation
|
||||
xtm1 = mu_xt + ( etas[idx] * variance ** 0.5 )*z
|
||||
xts[idx+1] = xtm1
|
||||
|
||||
if not zs is None:
|
||||
zs[-1] = torch.zeros_like(zs[-1])
|
||||
|
||||
return xt, zs, xts
|
||||
|
||||
|
||||
def reverse_step(model, model_output, timestep, sample, eta = 0, variance_noise=None):
|
||||
# 1. get previous step value (=t-1)
|
||||
prev_timestep = timestep - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps
|
||||
# 2. compute alphas, betas
|
||||
alpha_prod_t = model.scheduler.alphas_cumprod[timestep]
|
||||
alpha_prod_t_prev = model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else model.scheduler.final_alpha_cumprod
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
# 3. compute predicted original sample from predicted noise also called
|
||||
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
|
||||
# 5. compute variance: "sigma_t(η)" -> see formula (16)
|
||||
# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
|
||||
# variance = self.scheduler._get_variance(timestep, prev_timestep)
|
||||
variance = get_variance(model, timestep) #, prev_timestep)
|
||||
std_dev_t = eta * variance ** (0.5)
|
||||
# Take care of asymetric reverse process (asyrp)
|
||||
model_output_direction = model_output
|
||||
# 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
# pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output_direction
|
||||
pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * model_output_direction
|
||||
# 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
|
||||
# 8. Add noice if eta > 0
|
||||
if eta > 0:
|
||||
if variance_noise is None:
|
||||
variance_noise = torch.randn(model_output.shape, device=model.device, dtype =torch.float16)
|
||||
sigma_z = eta * variance ** (0.5) * variance_noise
|
||||
prev_sample = prev_sample + sigma_z
|
||||
|
||||
return prev_sample
|
||||
|
||||
def inversion_reverse_process(model,
|
||||
xT,
|
||||
etas = 0,
|
||||
prompts = "",
|
||||
cfg_scales = None,
|
||||
prog_bar = False,
|
||||
zs = None,
|
||||
controller=None,
|
||||
asyrp = False):
|
||||
|
||||
batch_size = len(prompts)
|
||||
|
||||
cfg_scales_tensor = torch.Tensor(cfg_scales).view(-1,1,1,1).to(model.device, dtype=torch.float16)
|
||||
|
||||
text_embeddings = encode_text(model, prompts)
|
||||
uncond_embedding = encode_text(model, [""] * batch_size)
|
||||
|
||||
if etas is None: etas = 0
|
||||
if type(etas) in [int, float]: etas = [etas]*model.scheduler.num_inference_steps
|
||||
assert len(etas) == model.scheduler.num_inference_steps
|
||||
timesteps = model.scheduler.timesteps.to(model.device)
|
||||
|
||||
xt = xT.expand(batch_size, -1, -1, -1)
|
||||
op = tqdm(timesteps[-zs.shape[0]:]) if prog_bar else timesteps[-zs.shape[0]:]
|
||||
|
||||
t_to_idx = {int(v):k for k,v in enumerate(timesteps[-zs.shape[0]:])}
|
||||
|
||||
for t in op:
|
||||
idx = t_to_idx[int(t)]
|
||||
## Unconditional embedding
|
||||
with torch.no_grad():
|
||||
uncond_out = model.unet.forward(xt, timestep = t,
|
||||
encoder_hidden_states = uncond_embedding)
|
||||
|
||||
## Conditional embedding
|
||||
if prompts:
|
||||
with torch.no_grad():
|
||||
cond_out = model.unet.forward(xt, timestep = t,
|
||||
encoder_hidden_states = text_embeddings)
|
||||
|
||||
|
||||
z = zs[idx] if not zs is None else None
|
||||
z = z.expand(batch_size, -1, -1, -1)
|
||||
if prompts:
|
||||
## classifier free guidance
|
||||
noise_pred = uncond_out.sample + cfg_scales_tensor * (cond_out.sample - uncond_out.sample)
|
||||
else:
|
||||
noise_pred = uncond_out.sample
|
||||
# 2. compute less noisy image and set x_t -> x_t-1
|
||||
xt = reverse_step(model, noise_pred, t, xt, eta = etas[idx], variance_noise = z)
|
||||
if controller is not None:
|
||||
xt = controller.step_callback(xt)
|
||||
return xt, zs
|
||||
@@ -0,0 +1,135 @@
|
||||
import gradio as gr
|
||||
import torch
|
||||
import numpy as np
|
||||
import requests
|
||||
import random
|
||||
from io import BytesIO
|
||||
from utils import *
|
||||
from constants import *
|
||||
from inversion_utils import *
|
||||
from modified_pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline
|
||||
from torch import autocast, inference_mode
|
||||
from diffusers import StableDiffusionPipeline
|
||||
from diffusers import DDIMScheduler
|
||||
from transformers import AutoProcessor, BlipForConditionalGeneration
|
||||
from share_btn import community_icon_html, loading_icon_html, share_js
|
||||
|
||||
from PIL import ImageFile
|
||||
import random
|
||||
|
||||
# load pipelines
|
||||
sd_model_id = "sd_model_v1-5"
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
sd_pipe = StableDiffusionPipeline.from_pretrained(sd_model_id).to(device)
|
||||
sd_pipe.scheduler = DDIMScheduler.from_config(sd_model_id, subfolder="scheduler")
|
||||
|
||||
sega_pipe = SemanticStableDiffusionPipeline.from_pretrained(sd_model_id).to(device)
|
||||
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
input_image = "./examples/00001.png" # @param
|
||||
|
||||
source_prompt = "human face" # @param
|
||||
target_prompt = "make up like a clown" # @param
|
||||
num_diffusion_steps = 100 # @param
|
||||
source_guidance_scale = 0 # @param
|
||||
reconstruct = True # @param
|
||||
skip_steps = 50 # @param
|
||||
target_guidance_scale = 10 # @param
|
||||
|
||||
# SEGA only params
|
||||
edit_concepts = ["star makeup", "heart makeup"] # @param
|
||||
edit_guidance_scales = [7, 15] # @param
|
||||
warmup_steps = [1, 1] # @param
|
||||
reverse_editing = [True, False] # @param
|
||||
thresholds = [0.95, 0.95] # @param
|
||||
|
||||
|
||||
def invert(x0: torch.FloatTensor, prompt_src: str = "", num_inference_steps=100, cfg_scale_src=3.5, eta=1):
|
||||
# inverts a real image according to Algorihm 1 in https://arxiv.org/pdf/2304.06140.pdf,
|
||||
# based on the code in https://github.com/inbarhub/DDPM_inversion
|
||||
|
||||
# returns wt, zs, wts:
|
||||
# wt - inverted latent
|
||||
# wts - intermediate inverted latents
|
||||
# zs - noise maps
|
||||
|
||||
sd_pipe.scheduler.set_timesteps(num_diffusion_steps)
|
||||
|
||||
# vae encode image
|
||||
with autocast("cuda"), inference_mode():
|
||||
w0 = (sd_pipe.vae.encode(x0).latent_dist.mode() * 0.18215).float()
|
||||
|
||||
# find Zs and wts - forward process
|
||||
wt, zs, wts = inversion_forward_process(sd_pipe, w0, etas=eta, prompt=prompt_src, cfg_scale=cfg_scale_src,
|
||||
prog_bar=True, num_inference_steps=num_diffusion_steps)
|
||||
return zs, wts
|
||||
|
||||
|
||||
def sample(zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
|
||||
# reverse process (via Zs and wT)
|
||||
w0, _ = inversion_reverse_process(sd_pipe, xT=wts[skip], etas=eta, prompts=[prompt_tar], cfg_scales=[cfg_scale_tar],
|
||||
prog_bar=True, zs=zs[skip:])
|
||||
|
||||
# vae decode image
|
||||
with autocast("cuda"), inference_mode():
|
||||
x0_dec = sd_pipe.vae.decode(1 / 0.18215 * w0).sample
|
||||
if x0_dec.dim() < 4:
|
||||
x0_dec = x0_dec[None, :, :, :]
|
||||
img = image_grid(x0_dec)
|
||||
return img
|
||||
|
||||
|
||||
def edit(wts, zs,
|
||||
tar_prompt="",
|
||||
steps=100,
|
||||
skip=36,
|
||||
tar_cfg_scale=8,
|
||||
edit_concept="",
|
||||
guidnace_scale=7,
|
||||
warmup=1,
|
||||
neg_guidance=False,
|
||||
threshold=0.95
|
||||
):
|
||||
# SEGA
|
||||
# parse concepts and neg guidance
|
||||
editing_args = dict(
|
||||
editing_prompt=edit_concept,
|
||||
reverse_editing_direction=neg_guidance,
|
||||
edit_warmup_steps=warmup,
|
||||
edit_guidance_scale=guidnace_scale,
|
||||
edit_threshold=threshold,
|
||||
edit_momentum_scale=0.5,
|
||||
edit_mom_beta=0.6,
|
||||
eta=1,
|
||||
)
|
||||
latnets = wts[skip].expand(1, -1, -1, -1)
|
||||
sega_out = sega_pipe(prompt=tar_prompt, latents=latnets, guidance_scale=tar_cfg_scale,
|
||||
num_images_per_prompt=1,
|
||||
num_inference_steps=steps,
|
||||
use_ddpm=True, wts=wts, zs=zs[skip:], **editing_args)
|
||||
return sega_out.images[0]
|
||||
|
||||
with open('prompt.txt') as file:
|
||||
lines = file.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
|
||||
for i in range(10000):
|
||||
try:
|
||||
input_image = 'origin_face' + str(i + 1) + '.png'
|
||||
randn = random.randint(0, len(lines)-1)
|
||||
target_prompt = lines[randn]
|
||||
x0 = load_512(input_image, device=device)
|
||||
# noise maps and latents
|
||||
zs, wts = invert(x0=x0, prompt_src=source_prompt, num_inference_steps=num_diffusion_steps,
|
||||
cfg_scale_src=source_guidance_scale)
|
||||
if reconstruct:
|
||||
ddpm_out_img = sample(zs, wts, prompt_tar=target_prompt, skip=skip_steps, cfg_scale_tar=target_guidance_scale)
|
||||
ddpm_out_img.save(f'makeup/edit_{i+1}.png')
|
||||
except:
|
||||
continue
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,754 @@
|
||||
|
||||
import inspect
|
||||
import warnings
|
||||
from itertools import repeat
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import torch
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
||||
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
||||
# from . import SemanticStableDiffusionPipelineOutput
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class SemanticStableDiffusionPipeline(DiffusionPipeline):
|
||||
r"""
|
||||
Pipeline for text-to-image generation with latent editing.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
|
||||
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
|
||||
|
||||
This model builds on the implementation of ['StableDiffusionPipeline']
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
|
||||
text_encoder ([`CLIPTextModel`]):
|
||||
Frozen text-encoder. Stable Diffusion uses the text portion of
|
||||
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
|
||||
the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
|
||||
tokenizer (`CLIPTokenizer`):
|
||||
Tokenizer of class
|
||||
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
|
||||
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
|
||||
scheduler ([`SchedulerMixin`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of
|
||||
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
|
||||
safety_checker ([`Q16SafetyChecker`]):
|
||||
Classification module that estimates whether generated images could be considered offensive or harmful.
|
||||
Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.
|
||||
feature_extractor ([`CLIPImageProcessor`]):
|
||||
Model that extracts features from generated images to be used as inputs for the `safety_checker`.
|
||||
"""
|
||||
|
||||
_optional_components = ["safety_checker", "feature_extractor"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
safety_checker: StableDiffusionSafetyChecker,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
requires_safety_checker: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if safety_checker is None and requires_safety_checker:
|
||||
logger.warning(
|
||||
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
|
||||
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
|
||||
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
|
||||
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
|
||||
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
|
||||
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
|
||||
)
|
||||
|
||||
if safety_checker is not None and feature_extractor is None:
|
||||
raise ValueError(
|
||||
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
|
||||
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
|
||||
)
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
safety_checker=safety_checker,
|
||||
feature_extractor=feature_extractor,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
self.register_to_config(requires_safety_checker=requires_safety_checker)
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
|
||||
def run_safety_checker(self, image, device, dtype):
|
||||
if self.safety_checker is None:
|
||||
has_nsfw_concept = None
|
||||
else:
|
||||
if torch.is_tensor(image):
|
||||
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
|
||||
else:
|
||||
feature_extractor_input = self.image_processor.numpy_to_pil(image)
|
||||
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
|
||||
image, has_nsfw_concept = self.safety_checker(
|
||||
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
|
||||
)
|
||||
return image, has_nsfw_concept
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
|
||||
def decode_latents(self, latents):
|
||||
warnings.warn(
|
||||
"The decode_latents method is deprecated and will be removed in a future version. Please"
|
||||
" use VaeImageProcessor instead",
|
||||
FutureWarning,
|
||||
)
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
height,
|
||||
width,
|
||||
callback_steps,
|
||||
negative_prompt=None,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
):
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
if (callback_steps is None) or (
|
||||
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
|
||||
f" {type(callback_steps)}."
|
||||
)
|
||||
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif prompt is None and prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
||||
)
|
||||
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
|
||||
if negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
|
||||
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
||||
)
|
||||
|
||||
if prompt_embeds is not None and negative_prompt_embeds is not None:
|
||||
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
||||
raise ValueError(
|
||||
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
||||
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
||||
f" {negative_prompt_embeds.shape}."
|
||||
)
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
|
||||
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
|
||||
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
num_images_per_prompt: int = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: int = 1,
|
||||
editing_prompt: Optional[Union[str, List[str]]] = None,
|
||||
editing_prompt_embeddings: Optional[torch.Tensor] = None,
|
||||
reverse_editing_direction: Optional[Union[bool, List[bool]]] = False,
|
||||
edit_guidance_scale: Optional[Union[float, List[float]]] = 5,
|
||||
edit_warmup_steps: Optional[Union[int, List[int]]] = 10,
|
||||
edit_cooldown_steps: Optional[Union[int, List[int]]] = None,
|
||||
edit_threshold: Optional[Union[float, List[float]]] = 0.9,
|
||||
edit_momentum_scale: Optional[float] = 0.1,
|
||||
edit_mom_beta: Optional[float] = 0.4,
|
||||
edit_weights: Optional[List[float]] = None,
|
||||
sem_guidance: Optional[List[torch.Tensor]] = None,
|
||||
|
||||
# DDPM additions
|
||||
use_ddpm: bool = False,
|
||||
wts: Optional[List[torch.Tensor]] = None,
|
||||
zs: Optional[List[torch.Tensor]] = None
|
||||
):
|
||||
r"""
|
||||
Function invoked when calling the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`):
|
||||
The prompt or prompts to guide the image generation.
|
||||
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The width in pixels of the generated image.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
guidance_scale (`float`, *optional*, defaults to 7.5):
|
||||
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
||||
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
||||
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
||||
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
||||
usually at the expense of lower image quality.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
|
||||
if `guidance_scale` is less than `1`).
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
eta (`float`, *optional*, defaults to 0.0):
|
||||
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
||||
[`schedulers.DDIMScheduler`], will be ignored for others.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
||||
to make generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor will ge generated by sampling using the supplied random `generator`.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generate image. Choose between
|
||||
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
||||
plain tuple.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that will be called every `callback_steps` steps during inference. The function will be
|
||||
called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
||||
callback_steps (`int`, *optional*, defaults to 1):
|
||||
The frequency at which the `callback` function will be called. If not specified, the callback will be
|
||||
called at every step.
|
||||
editing_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to use for Semantic guidance. Semantic guidance is disabled by setting
|
||||
`editing_prompt = None`. Guidance direction of prompt should be specified via
|
||||
`reverse_editing_direction`.
|
||||
editing_prompt_embeddings (`torch.Tensor>`, *optional*):
|
||||
Pre-computed embeddings to use for semantic guidance. Guidance direction of embedding should be
|
||||
specified via `reverse_editing_direction`.
|
||||
reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`):
|
||||
Whether the corresponding prompt in `editing_prompt` should be increased or decreased.
|
||||
edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5):
|
||||
Guidance scale for semantic guidance. If provided as list values should correspond to `editing_prompt`.
|
||||
`edit_guidance_scale` is defined as `s_e` of equation 6 of [SEGA
|
||||
Paper](https://arxiv.org/pdf/2301.12247.pdf).
|
||||
edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10):
|
||||
Number of diffusion steps (for each prompt) for which semantic guidance will not be applied. Momentum
|
||||
will still be calculated for those steps and applied once all warmup periods are over.
|
||||
`edit_warmup_steps` is defined as `delta` (δ) of [SEGA Paper](https://arxiv.org/pdf/2301.12247.pdf).
|
||||
edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`):
|
||||
Number of diffusion steps (for each prompt) after which semantic guidance will no longer be applied.
|
||||
edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9):
|
||||
Threshold of semantic guidance.
|
||||
edit_momentum_scale (`float`, *optional*, defaults to 0.1):
|
||||
Scale of the momentum to be added to the semantic guidance at each diffusion step. If set to 0.0
|
||||
momentum will be disabled. Momentum is already built up during warmup, i.e. for diffusion steps smaller
|
||||
than `sld_warmup_steps`. Momentum will only be added to latent guidance once all warmup periods are
|
||||
finished. `edit_momentum_scale` is defined as `s_m` of equation 7 of [SEGA
|
||||
Paper](https://arxiv.org/pdf/2301.12247.pdf).
|
||||
edit_mom_beta (`float`, *optional*, defaults to 0.4):
|
||||
Defines how semantic guidance momentum builds up. `edit_mom_beta` indicates how much of the previous
|
||||
momentum will be kept. Momentum is already built up during warmup, i.e. for diffusion steps smaller
|
||||
than `edit_warmup_steps`. `edit_mom_beta` is defined as `beta_m` (β) of equation 8 of [SEGA
|
||||
Paper](https://arxiv.org/pdf/2301.12247.pdf).
|
||||
edit_weights (`List[float]`, *optional*, defaults to `None`):
|
||||
Indicates how much each individual concept should influence the overall guidance. If no weights are
|
||||
provided all concepts are applied equally. `edit_mom_beta` is defined as `g_i` of equation 9 of [SEGA
|
||||
Paper](https://arxiv.org/pdf/2301.12247.pdf).
|
||||
sem_guidance (`List[torch.Tensor]`, *optional*):
|
||||
List of pre-generated guidance vectors to be applied at generation. Length of the list has to
|
||||
correspond to `num_inference_steps`.
|
||||
|
||||
Returns:
|
||||
[`~pipelines.semantic_stable_diffusion.SemanticStableDiffusionPipelineOutput`] or `tuple`:
|
||||
[`~pipelines.semantic_stable_diffusion.SemanticStableDiffusionPipelineOutput`] if `return_dict` is True,
|
||||
otherwise a `tuple. When returning a tuple, the first element is a list with the generated images, and the
|
||||
second element is a list of `bool`s denoting whether the corresponding generated image likely represents
|
||||
"not-safe-for-work" (nsfw) content, according to the `safety_checker`.
|
||||
"""
|
||||
# 0. Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(prompt, height, width, callback_steps)
|
||||
|
||||
# 2. Define call parameters
|
||||
batch_size = 1 if isinstance(prompt, str) else len(prompt)
|
||||
|
||||
if editing_prompt:
|
||||
enable_edit_guidance = True
|
||||
if isinstance(editing_prompt, str):
|
||||
editing_prompt = [editing_prompt]
|
||||
enabled_editing_prompts = len(editing_prompt)
|
||||
elif editing_prompt_embeddings is not None:
|
||||
enable_edit_guidance = True
|
||||
enabled_editing_prompts = editing_prompt_embeddings.shape[0]
|
||||
else:
|
||||
enabled_editing_prompts = 0
|
||||
enable_edit_guidance = False
|
||||
|
||||
# get prompt text embeddings
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=self.tokenizer.model_max_length,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
|
||||
if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
|
||||
removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])
|
||||
logger.warning(
|
||||
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
||||
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
|
||||
)
|
||||
text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]
|
||||
text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
bs_embed, seq_len, _ = text_embeddings.shape
|
||||
text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
|
||||
text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
|
||||
|
||||
if enable_edit_guidance:
|
||||
# get safety text embeddings
|
||||
if editing_prompt_embeddings is None:
|
||||
edit_concepts_input = self.tokenizer(
|
||||
[x for item in editing_prompt for x in repeat(item, batch_size)],
|
||||
padding="max_length",
|
||||
max_length=self.tokenizer.model_max_length,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
edit_concepts_input_ids = edit_concepts_input.input_ids
|
||||
|
||||
if edit_concepts_input_ids.shape[-1] > self.tokenizer.model_max_length:
|
||||
removed_text = self.tokenizer.batch_decode(
|
||||
edit_concepts_input_ids[:, self.tokenizer.model_max_length :]
|
||||
)
|
||||
logger.warning(
|
||||
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
||||
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
|
||||
)
|
||||
edit_concepts_input_ids = edit_concepts_input_ids[:, : self.tokenizer.model_max_length]
|
||||
edit_concepts = self.text_encoder(edit_concepts_input_ids.to(self.device))[0]
|
||||
else:
|
||||
edit_concepts = editing_prompt_embeddings.to(self.device).repeat(batch_size, 1, 1)
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
bs_embed_edit, seq_len_edit, _ = edit_concepts.shape
|
||||
edit_concepts = edit_concepts.repeat(1, num_images_per_prompt, 1)
|
||||
edit_concepts = edit_concepts.view(bs_embed_edit * num_images_per_prompt, seq_len_edit, -1)
|
||||
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
# get unconditional embeddings for classifier free guidance
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
uncond_tokens: List[str]
|
||||
if negative_prompt is None:
|
||||
uncond_tokens = [""]
|
||||
elif type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif isinstance(negative_prompt, str):
|
||||
uncond_tokens = [negative_prompt]
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
else:
|
||||
uncond_tokens = negative_prompt
|
||||
|
||||
max_length = text_input_ids.shape[-1]
|
||||
uncond_input = self.tokenizer(
|
||||
uncond_tokens,
|
||||
padding="max_length",
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
|
||||
|
||||
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
||||
seq_len = uncond_embeddings.shape[1]
|
||||
uncond_embeddings = uncond_embeddings.repeat(batch_size, num_images_per_prompt, 1)
|
||||
uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
if enable_edit_guidance:
|
||||
text_embeddings = torch.cat([uncond_embeddings, text_embeddings, edit_concepts])
|
||||
else:
|
||||
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
|
||||
# get the initial random noise unless the user supplied it
|
||||
|
||||
# 4. Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=self.device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
if use_ddpm:
|
||||
t_to_idx = {int(v):k for k,v in enumerate(timesteps[-zs.shape[0]:])}
|
||||
timesteps = timesteps[-zs.shape[0]:]
|
||||
|
||||
# 5. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
text_embeddings.dtype,
|
||||
self.device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 6. Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# Initialize edit_momentum to None
|
||||
edit_momentum = None
|
||||
|
||||
self.uncond_estimates = None
|
||||
self.text_estimates = None
|
||||
self.edit_estimates = None
|
||||
self.sem_guidance = None
|
||||
|
||||
for i, t in enumerate(self.progress_bar(timesteps)):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = (
|
||||
torch.cat([latents] * (2 + enabled_editing_prompts)) if do_classifier_free_guidance else latents
|
||||
)
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_out = noise_pred.chunk(2 + enabled_editing_prompts) # [b,4, 64, 64]
|
||||
noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1]
|
||||
noise_pred_edit_concepts = noise_pred_out[2:]
|
||||
|
||||
# default text guidance
|
||||
noise_guidance = guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
# noise_guidance = (noise_pred_text - noise_pred_edit_concepts[0])
|
||||
|
||||
if self.uncond_estimates is None:
|
||||
self.uncond_estimates = torch.zeros((num_inference_steps + 1, *noise_pred_uncond.shape))
|
||||
self.uncond_estimates[i] = noise_pred_uncond.detach().cpu()
|
||||
|
||||
if self.text_estimates is None:
|
||||
self.text_estimates = torch.zeros((num_inference_steps + 1, *noise_pred_text.shape))
|
||||
self.text_estimates[i] = noise_pred_text.detach().cpu()
|
||||
|
||||
if self.edit_estimates is None and enable_edit_guidance:
|
||||
self.edit_estimates = torch.zeros(
|
||||
(num_inference_steps + 1, len(noise_pred_edit_concepts), *noise_pred_edit_concepts[0].shape)
|
||||
)
|
||||
|
||||
if self.sem_guidance is None:
|
||||
self.sem_guidance = torch.zeros((num_inference_steps + 1, *noise_pred_text.shape))
|
||||
|
||||
if edit_momentum is None:
|
||||
edit_momentum = torch.zeros_like(noise_guidance)
|
||||
|
||||
if enable_edit_guidance:
|
||||
concept_weights = torch.zeros(
|
||||
(len(noise_pred_edit_concepts), noise_guidance.shape[0]),
|
||||
device=self.device,
|
||||
dtype=noise_guidance.dtype,
|
||||
)
|
||||
noise_guidance_edit = torch.zeros(
|
||||
(len(noise_pred_edit_concepts), *noise_guidance.shape),
|
||||
device=self.device,
|
||||
dtype=noise_guidance.dtype,
|
||||
)
|
||||
# noise_guidance_edit = torch.zeros_like(noise_guidance)
|
||||
warmup_inds = []
|
||||
for c, noise_pred_edit_concept in enumerate(noise_pred_edit_concepts):
|
||||
self.edit_estimates[i, c] = noise_pred_edit_concept
|
||||
if isinstance(edit_guidance_scale, list):
|
||||
edit_guidance_scale_c = edit_guidance_scale[c]
|
||||
else:
|
||||
edit_guidance_scale_c = edit_guidance_scale
|
||||
|
||||
if isinstance(edit_threshold, list):
|
||||
edit_threshold_c = edit_threshold[c]
|
||||
else:
|
||||
edit_threshold_c = edit_threshold
|
||||
if isinstance(reverse_editing_direction, list):
|
||||
reverse_editing_direction_c = reverse_editing_direction[c]
|
||||
else:
|
||||
reverse_editing_direction_c = reverse_editing_direction
|
||||
if edit_weights:
|
||||
edit_weight_c = edit_weights[c]
|
||||
else:
|
||||
edit_weight_c = 1.0
|
||||
if isinstance(edit_warmup_steps, list):
|
||||
edit_warmup_steps_c = edit_warmup_steps[c]
|
||||
else:
|
||||
edit_warmup_steps_c = edit_warmup_steps
|
||||
|
||||
if isinstance(edit_cooldown_steps, list):
|
||||
edit_cooldown_steps_c = edit_cooldown_steps[c]
|
||||
elif edit_cooldown_steps is None:
|
||||
edit_cooldown_steps_c = i + 1
|
||||
else:
|
||||
edit_cooldown_steps_c = edit_cooldown_steps
|
||||
if i >= edit_warmup_steps_c:
|
||||
warmup_inds.append(c)
|
||||
if i >= edit_cooldown_steps_c:
|
||||
noise_guidance_edit[c, :, :, :, :] = torch.zeros_like(noise_pred_edit_concept)
|
||||
continue
|
||||
|
||||
noise_guidance_edit_tmp = noise_pred_edit_concept - noise_pred_uncond
|
||||
# tmp_weights = (noise_pred_text - noise_pred_edit_concept).sum(dim=(1, 2, 3))
|
||||
tmp_weights = (noise_guidance - noise_pred_edit_concept).sum(dim=(1, 2, 3))
|
||||
|
||||
tmp_weights = torch.full_like(tmp_weights, edit_weight_c) # * (1 / enabled_editing_prompts)
|
||||
if reverse_editing_direction_c:
|
||||
noise_guidance_edit_tmp = noise_guidance_edit_tmp * -1
|
||||
concept_weights[c, :] = tmp_weights
|
||||
|
||||
noise_guidance_edit_tmp = noise_guidance_edit_tmp * edit_guidance_scale_c
|
||||
|
||||
# torch.quantile function expects float32
|
||||
if noise_guidance_edit_tmp.dtype == torch.float32:
|
||||
tmp = torch.quantile(
|
||||
torch.abs(noise_guidance_edit_tmp).flatten(start_dim=2),
|
||||
edit_threshold_c,
|
||||
dim=2,
|
||||
keepdim=False,
|
||||
)
|
||||
else:
|
||||
tmp = torch.quantile(
|
||||
torch.abs(noise_guidance_edit_tmp).flatten(start_dim=2).to(torch.float32),
|
||||
edit_threshold_c,
|
||||
dim=2,
|
||||
keepdim=False,
|
||||
).to(noise_guidance_edit_tmp.dtype)
|
||||
|
||||
noise_guidance_edit_tmp = torch.where(
|
||||
torch.abs(noise_guidance_edit_tmp) >= tmp[:, :, None, None],
|
||||
noise_guidance_edit_tmp,
|
||||
torch.zeros_like(noise_guidance_edit_tmp),
|
||||
)
|
||||
noise_guidance_edit[c, :, :, :, :] = noise_guidance_edit_tmp
|
||||
|
||||
# noise_guidance_edit = noise_guidance_edit + noise_guidance_edit_tmp
|
||||
|
||||
warmup_inds = torch.tensor(warmup_inds).to(self.device)
|
||||
if len(noise_pred_edit_concepts) > warmup_inds.shape[0] > 0:
|
||||
concept_weights = concept_weights.to("cpu") # Offload to cpu
|
||||
noise_guidance_edit = noise_guidance_edit.to("cpu")
|
||||
|
||||
concept_weights_tmp = torch.index_select(concept_weights.to(self.device), 0, warmup_inds)
|
||||
concept_weights_tmp = torch.where(
|
||||
concept_weights_tmp < 0, torch.zeros_like(concept_weights_tmp), concept_weights_tmp
|
||||
)
|
||||
concept_weights_tmp = concept_weights_tmp / concept_weights_tmp.sum(dim=0)
|
||||
# concept_weights_tmp = torch.nan_to_num(concept_weights_tmp)
|
||||
|
||||
noise_guidance_edit_tmp = torch.index_select(
|
||||
noise_guidance_edit.to(self.device), 0, warmup_inds
|
||||
)
|
||||
noise_guidance_edit_tmp = torch.einsum(
|
||||
"cb,cbijk->bijk", concept_weights_tmp, noise_guidance_edit_tmp
|
||||
)
|
||||
noise_guidance_edit_tmp = noise_guidance_edit_tmp
|
||||
noise_guidance = noise_guidance + noise_guidance_edit_tmp
|
||||
|
||||
self.sem_guidance[i] = noise_guidance_edit_tmp.detach().cpu()
|
||||
|
||||
del noise_guidance_edit_tmp
|
||||
del concept_weights_tmp
|
||||
concept_weights = concept_weights.to(self.device)
|
||||
noise_guidance_edit = noise_guidance_edit.to(self.device)
|
||||
|
||||
concept_weights = torch.where(
|
||||
concept_weights < 0, torch.zeros_like(concept_weights), concept_weights
|
||||
)
|
||||
|
||||
concept_weights = torch.nan_to_num(concept_weights)
|
||||
|
||||
noise_guidance_edit = torch.einsum("cb,cbijk->bijk", concept_weights, noise_guidance_edit)
|
||||
|
||||
noise_guidance_edit = noise_guidance_edit + edit_momentum_scale * edit_momentum
|
||||
|
||||
edit_momentum = edit_mom_beta * edit_momentum + (1 - edit_mom_beta) * noise_guidance_edit
|
||||
|
||||
if warmup_inds.shape[0] == len(noise_pred_edit_concepts):
|
||||
noise_guidance = noise_guidance + noise_guidance_edit
|
||||
self.sem_guidance[i] = noise_guidance_edit.detach().cpu()
|
||||
|
||||
if sem_guidance is not None:
|
||||
edit_guidance = sem_guidance[i].to(self.device)
|
||||
noise_guidance = noise_guidance + edit_guidance
|
||||
|
||||
noise_pred = noise_pred_uncond + noise_guidance
|
||||
## ddpm ###########################################################
|
||||
if use_ddpm:
|
||||
|
||||
idx = t_to_idx[int(t)]
|
||||
z = zs[idx] if not zs is None else None
|
||||
|
||||
# 1. get previous step value (=t-1)
|
||||
prev_timestep = t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
|
||||
# 2. compute alphas, betas
|
||||
alpha_prod_t = self.scheduler.alphas_cumprod[t]
|
||||
alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
|
||||
# 3. compute predicted original sample from predicted noise also called
|
||||
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
|
||||
|
||||
|
||||
# 5. compute variance: "sigma_t(η)" -> see formula (16)
|
||||
# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
|
||||
# variance = self.scheduler._get_variance(timestep, prev_timestep)
|
||||
# variance = get_variance(model, t) #, prev_timestep)
|
||||
beta_prod_t_prev = 1 - alpha_prod_t_prev
|
||||
variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
|
||||
|
||||
|
||||
|
||||
std_dev_t = eta * variance ** (0.5)
|
||||
# Take care of asymetric reverse process (asyrp)
|
||||
noise_pred_direction = noise_pred
|
||||
|
||||
# 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
# pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output_direction
|
||||
pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * noise_pred_direction
|
||||
|
||||
# 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
|
||||
prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
|
||||
# 8. Add noice if eta > 0
|
||||
if eta > 0:
|
||||
if z is None:
|
||||
z = torch.randn(noise_pred.shape, device=self.device)
|
||||
sigma_z = eta * variance ** (0.5) * z
|
||||
latents = prev_sample + sigma_z
|
||||
|
||||
## ddpm ##########################################################
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
else: #if not use_ddpm:
|
||||
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
|
||||
|
||||
# call the callback, if provided
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
callback(i, t, latents)
|
||||
|
||||
|
||||
# # 8. Post-processing
|
||||
# image = self.decode_latents(latents)
|
||||
|
||||
# # 9. Run safety checker
|
||||
# image, has_nsfw_concept = self.run_safety_checker(image, self.device, text_embeddings.dtype)
|
||||
|
||||
# # 10. Convert to PIL
|
||||
# if output_type == "pil":
|
||||
# image = self.numpy_to_pil(image)
|
||||
|
||||
# if not return_dict:
|
||||
# return (image, has_nsfw_concept)
|
||||
|
||||
# return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
||||
|
||||
# 8. Post-processing
|
||||
if not output_type == "latent":
|
||||
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
||||
image, has_nsfw_concept = self.run_safety_checker(image, self.device, text_embeddings.dtype)
|
||||
else:
|
||||
image = latents
|
||||
has_nsfw_concept = None
|
||||
|
||||
if has_nsfw_concept is None:
|
||||
do_denormalize = [True] * image.shape[0]
|
||||
else:
|
||||
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
||||
|
||||
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
||||
|
||||
if not return_dict:
|
||||
return (image, has_nsfw_concept)
|
||||
|
||||
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
||||
@@ -0,0 +1,292 @@
|
||||
Make it qilin makeup expressing mythical mercies
|
||||
Make it jackalope makeup combining whimsy and wanderlust
|
||||
Make it thunderbird makeup harnessing stormy supremacy
|
||||
Make it nuckelavee makeup spreading haunting maritime malice
|
||||
Make it cockatrice makeup radiating curious calamity
|
||||
Make it manticore makeup manifesting monstrous majesty
|
||||
Make it basilisk makeup hypnotizing with hidden handsomeness
|
||||
Make it minotaur makeup representing rugged resilience
|
||||
Make it chimera makeup captivating with extraordinary eclecticism
|
||||
Make it centaur makeup celebrating untamed vitality
|
||||
Make it yeti makeup symbolizing solemn solitude
|
||||
Make it kraken makeup oozing ancient aqueous authority
|
||||
Make it hydra makeup asserting immense indifference
|
||||
Make it kitsune makeup weaving supernatural surprise
|
||||
Make it cerberus makeup guarding against trespassing terrors
|
||||
Make it gryphon makeup amalgamating aerie alacrity
|
||||
Make it wyvern makeup whispering windswept wonderment
|
||||
Make it pegasus makeup prancing with persistent passion
|
||||
Make it phoenix makeup perpetuating primal persistency
|
||||
Make it ceres makeup commanding natural nobility
|
||||
Make it hecate makeup harnessing haunting lunar luminosity
|
||||
Make it neptune makeup navigating numinous mysteries
|
||||
Make it aphrodite makeup radiating romantic rapture
|
||||
Make it thanatos makeup transporting ephemeral allure
|
||||
Make it dionysus makeup indulging Dionysian delights
|
||||
Make it apollo makeup enlightening oracular omniscience
|
||||
Make it poseidon makeup dominating dashing dominance
|
||||
Make it hades makeup haunting the underworld's obscure obsessions
|
||||
Make it athena makeup accentuating adaptive acumen
|
||||
Make it zeus makeup bearing bolts of benevolent bravado
|
||||
Make it ra makeup blazing ferocious ethereal fire
|
||||
Make it amaterasu makeup illuminating innate inspiration
|
||||
Make it freyja makeup captivating with celestial sensuality
|
||||
Make it loki makeup twisting tales with mischievous mirth
|
||||
Make it odin makeup cultivating arcane acumen and acuity
|
||||
Make it persephone makeup tempting with transient temptations
|
||||
Make it morrigan makeup coursing ferocity with flecks of fear
|
||||
Make it cauchemar makeup spreading nocturnal nightmares
|
||||
Make it nain rouge makeup rousing Dread among the daring
|
||||
Make it jack o'lantern makeup gleaming with grinning ghosts
|
||||
Make it yuki-onna makeup whispering winter's wildest wonders
|
||||
Make it krampus makeup punishing with primordial panic
|
||||
Make it la llorona makeup luring with lamenting lures
|
||||
Make it veles makeup communicating communion with courts cryptic
|
||||
Make it melusine makeup shimmering with serpentine splendor
|
||||
Make it aswang makeup transforming traveling terrors
|
||||
Make it nue makeup taking tangled temerity forms
|
||||
Make it tengu makeup uttering otherworldly omens
|
||||
Make it cherry soul makeup blooming bittersweet beauty
|
||||
Make it azuki makeup radiating good fortune's fortunate finds
|
||||
Make it tsuchigumo makeup scheming shrewd silken shrouds
|
||||
Make it ubume makeup haunting history's harshest horrors
|
||||
Make it kappa makeup keeping waters mysterious motives masked
|
||||
Make it raiju makeup rumbling rambunctious nature's nuance
|
||||
Make it bake-kujira makeup emanating ethereal evening emanations
|
||||
Make it ishutski makeup invoking imaginings of icicle islands
|
||||
Make it futakuchi-onna makeup multiplying misdirecting mobilities
|
||||
Make it noppera-bo makeup tricking travels in tenebrous times
|
||||
Make it rokurokubi makeup lengthening lurid legends legacies
|
||||
Make it namahage makeup menacing midnight mountain mysteries
|
||||
Make it obake makeup morphing many marvelous manifestations
|
||||
Make it nogitsune makeup navigating nature's nightly nuances
|
||||
Make it itsumade makeup imbuing intimations of impermanence
|
||||
Make it mujina makeup manifesting mercurial metamorphoses
|
||||
Make it orochi makeup oozing ostentatious octuple opulence
|
||||
Make it yamabiko makeup resonating resounding rocky regions
|
||||
Make it mono no ke makeup concealing countenance with cryptic calm
|
||||
Make it iwanko makeup wandering wistful woodland ways
|
||||
Make it merlion makeup mixing marine mystique
|
||||
Make it qilin makeup expressing mythical mysticism
|
||||
Make it jackalope makeup conjuring whimsical woodland magic
|
||||
Make it thunderbird makeup manifesting majestic avian aura
|
||||
Make it cockatrice makeup radiating legendary lethal loveliness
|
||||
Make it manticore makeup menacing with monstrous magnificence
|
||||
Make it basilisk makeup hypnotizing with hypnotic handsomeness
|
||||
Make it minotaur makeup representing mythical muscular might
|
||||
Make it chimera makeup captivating with chimerical charisma
|
||||
Make it centaur makeup celebrating classical crossbred charm
|
||||
Make it yeti makeup yodeling elusive alpine allure
|
||||
Make it kraken makeup oozing ominous oceanic omnipotence
|
||||
Make it hydra makeup haunting with heads of horrifying handsomeness
|
||||
Make it kitsune makeup masking with supernatural slyness
|
||||
Make it cerberus makeup guarding with ghastly canine grandeur
|
||||
Make it griffin makeup galvanizing with golden grandeur
|
||||
Make it wyvern makeup whispering windswept wonder
|
||||
Make it pegasus makeup prancing with perennial passion
|
||||
Make it phoenix makeup perpetuating primeval powers
|
||||
Make it ceres makeup commanding cosmic carnal control
|
||||
Make it hecate makeup harnessing hauntedlunar luminosity
|
||||
Make it neptune makeup navigating nepenthean mystique
|
||||
Make it aphrodite makeup radiating romantic rapture
|
||||
Make it thanatos makeup transporting ephemeral eternal allure
|
||||
Make it dionysus makeup indulging divine decadence
|
||||
Make it apollo makeup enlightening oracular allure
|
||||
Make it poseidon makeup dominating dashing aquatic style
|
||||
Make it hades makeup haunting the underworld's enigmatic embrace
|
||||
Make it athena makeup accentuating agile aesthetic acumen
|
||||
Make it zeus makeup bearing bold bolts of beauty
|
||||
Make it ra makeup blazing ferocious ethereal allure
|
||||
Make it amaterasu makeup illuminating celestial comeliness
|
||||
Make it freyja makeup captivating with celestial sensuality
|
||||
Make it loki makeup twisting wiles with a wicked wink
|
||||
Make it odin makeup patining mystical wisdom and war
|
||||
Make it persephone makeup tempting with seasonal allure
|
||||
Make it morrigan makeup couraging ferocity with flecksof fear
|
||||
Make it fairy makeup bringing bubbly bliss
|
||||
Make it angel makeup emanating ethereal elegance
|
||||
Make it knight makeup exuding exemplary excellence
|
||||
Make it pirate makeup exuding eccentric exuberance
|
||||
Make it dinosaur makeup roaring with rugged resilience
|
||||
Make it monkey makeup swinging with silly spirits
|
||||
Make it peacock makeup displaying dazzling decorations
|
||||
Make it raccoon makeup masked with mischievous mystery
|
||||
Make it deer makeup gleaming with gentle grace
|
||||
Make it owl makeup wise with watchful eyes
|
||||
Make it llama makeup laughing with likable levity
|
||||
Make it bee makeup buzzing with busy beauty
|
||||
Make it shark makeup swimming with savage strength
|
||||
Make it octopus makeup mystifying with mobile motion
|
||||
Make it koala makeup cozy with calm care
|
||||
Make it ladybug makeup lucky with loveable looks
|
||||
Make it flamingo makeup pretty with pink poise
|
||||
Make it horse makeup handsome with heroic hustle
|
||||
Make it tiger makeup tough with tenacious talent
|
||||
Make it elephant makeup elegant with enduring endurance
|
||||
Make it penguin makeup playful with preppy polish
|
||||
Make it dolphin makeup dancing with delightful dexterity
|
||||
Make it bear makeup brawny with burly brawn
|
||||
Make it camel makeup companions with considerate care
|
||||
Make it kangaroo makeup jumping with joyous joviality
|
||||
Make it mouse makeup merry with mindful mannerisms
|
||||
Make it turtle makeup tranquil with trustworthy truth
|
||||
Make it snail makeup steady with surefooted slowness
|
||||
Make it goat makeup gleeful with gusto and glee
|
||||
Make it sheep makeup following with faithful flock
|
||||
Make it donkey makeup dauntless with durable diligence
|
||||
Make it moose makeup majestic with mighty magnitude
|
||||
Make it buffalo makeup bold with benevolent brawn
|
||||
Make it hamster makeup happy with hurried hustle
|
||||
Make it parrot makeup personable with pretty plumes
|
||||
Make it deer makeup gracious with gentle grace
|
||||
Make it otter makeup playful with a love of water
|
||||
Make it chipmunk makeup charming with cheerfully busy habits
|
||||
Make it bunny makeup bouncy with a brisk and bold spirit
|
||||
Make it panda makeup precious with peerless peacefulness
|
||||
Make it fox makeup foxy with discreet daring
|
||||
Make it raccoon makeup rascal with reserved resourcefulness
|
||||
Make it squirrel makeup spirited with swift and sprightly styles
|
||||
Make it woodpecker makeup zany with zippy zealousness
|
||||
Make it skunk makeup sly with a striking social scheme
|
||||
Make it badger makeup brave with bearing and backbone
|
||||
Make it armadillo makeup armored with amiable armor
|
||||
Make it opossum makeup opportunistic with occasional optimism
|
||||
Make it clown makeup spotlighting funny features
|
||||
Make it zombie makeup emphasizing ghoulish details
|
||||
Make it unicorn makeup accenting magical aura
|
||||
Make it cartoon makeup bringing whimsical drawings to life
|
||||
Make it vampire makeup enhancing ominous presence
|
||||
Make it fairy makeup highlighting ethereal glow
|
||||
Make it cat makeup emphasizing feline facial lines
|
||||
Make it princess makeup framing regal radiance
|
||||
Make it mime makeup articulating silent stories
|
||||
Make it bat makeup hovering mysterious shadows
|
||||
Make it pig makeup oinking humor out front
|
||||
Make it devil makeup flaming fiendish functions
|
||||
Make it skeleton makeup gleaming grim grin
|
||||
Make it mermaid makeup glimmering undersea ambience
|
||||
Make it witch makeup hexing bewitching beauty
|
||||
Make it superhero makeup spotlighting courageous character
|
||||
Make it harlequin makeup jestering jovial jokes
|
||||
Make it butterfly makeup fluttering fanciful features
|
||||
Make it panda makeup bearing bearish black and white
|
||||
Make it flower makeup blossoming botanical beauty
|
||||
Make it pumpkin makeup gourdifying ghoulish glee
|
||||
Make it robot makeup powering mechanical mischief
|
||||
Make it cow makeup mooing mirthful moments
|
||||
Make it fox makeup cunning crafty complexion
|
||||
Make it ninja makeup concealing covert craft
|
||||
Make it peacock makeup feathering flashy face
|
||||
Make it elf makeup evoking ethereal enchantment
|
||||
Make it wolf makeup prowling predatory persona
|
||||
Make it star makeup beaming bright visage
|
||||
Make it scarecrow makeup stuffing strawfilled fright
|
||||
Make it angel makeup basking beatific beauty
|
||||
Make it zombie bride makeup rising reanimated radiance
|
||||
Make it bumblebee makeup buzzing bold bugliness
|
||||
Make it dalmatian makeup spotting canine comeliness
|
||||
Make it phoenix makeup igniting fiery features
|
||||
Make it chameleon makeup adapting active appearance
|
||||
Make it tiger makeup stalking striped strength
|
||||
Make it snake makeup slithering sly visage
|
||||
Make it raccoon makeup banditing bashful brigandry
|
||||
Make it spider makeup weaving wicked wonder
|
||||
Make it geisha makeup flowing fanciful femininity
|
||||
Make it ballerina makeup dancing delicate disposition
|
||||
Make it giraffe makeup lengthening lofty looks
|
||||
Make it frog makeup hopping herpetological humor
|
||||
Make it merman makeup plunging amphibious aesthetics
|
||||
Make it pirate makeup plundering privateering panache
|
||||
Make it rainbow makeup painting prismatic perspective
|
||||
Make it pegasus makeup soaring celestial spirit
|
||||
Make it knight makeup marching heroic honor
|
||||
Make it butterfly makeup fluttering colorful composition
|
||||
Make it dragon makeup exhaling exotic elegance
|
||||
Make it witch makeup casting captivating charms
|
||||
Make it reindeer makeup prancing picturesque personality
|
||||
Make it disco makeup boogying brassy bombast
|
||||
Make it teddy makeup bear hugging homey honesty
|
||||
Make it churro makeup sweetening spirited savoriness
|
||||
Make it rainbow makeup tinting pristine panoply
|
||||
Make it mermaid makeup swirling oceanic opulence
|
||||
Make it unicorn makeup radiating rainbow resplendence
|
||||
Make it fairy makeup conveying fantastical facets
|
||||
Make it zombie makeup decomposing dreadful decay
|
||||
Make it devil makeup depicting damnable dimensions
|
||||
Make it angel makeup expressing ethereal ends
|
||||
Make it vampire makeup vivifying vicious visage
|
||||
Make it witch makeup weaving wizardly worlds
|
||||
Make it gypsy makeup guiding gyrovago grandeur
|
||||
Make it ninja makeup narrating nuanced obfuscation
|
||||
Make it geisha makeup granting gracious gentility
|
||||
Make it pirate makeup pilfering plundered prosperity
|
||||
Make it knight makeup chronicling chivalric chronicles
|
||||
Make it princess makeup privileging pomp and pageantry
|
||||
Make it wizard makeup weaving wondrous wonders
|
||||
Make it zombie makeup memorializing macabre mortality
|
||||
Make it phoenix makeup portending perpetual progression
|
||||
Make it mermaid makeup chronicling coral chronologies
|
||||
Make it dragon makeup depicting dazzling diversions
|
||||
Make it tiger makeup tracking tenacious triumphs
|
||||
Make it unicorn makeup ushering utopian visions
|
||||
Make it phoenix makeup projecting prophetic pronouncements
|
||||
Make it elf makeup emitting ethereal effulgence
|
||||
Make it wizard makeup wielding wizened wizardry
|
||||
Make it fairy makeup facilitating fanciful forays
|
||||
Make it werewolf makeup warding wild wilderness
|
||||
Make it angel makeup announcing altruistic arrivals
|
||||
Make it witch makeup working wicked wonders
|
||||
Make it pirate makeup pilfering precious plunder
|
||||
Make it princess makeup portraying pure principles
|
||||
Make it geisha makeup granting gracious gestures
|
||||
Make it genie makeup granting grandiose gifts
|
||||
Make it unicorn makeup upholding utopian unity
|
||||
Make it elf makeup ensuring evanescent enchantment
|
||||
Make it mermaid makeup maintaining mystical mysteries
|
||||
Make it phoenix makeup promising perpetual perfection
|
||||
Make it werewolf makeup howling haunting hymns
|
||||
Make it vampire makeup venerating vehement victory
|
||||
Make it witch makeup weaving wonderful worlds
|
||||
Make it princess makeup portraying peaceful purposes
|
||||
Make it angel makeup advocating altruistic action
|
||||
Make it unicorn makeup upholding utopian undertakings
|
||||
Make it phoenix makeup projecting Promised prosperity
|
||||
Make it mermaid makeup maintaining mesmeric melodies
|
||||
Make it elf makeup ensuring effervescent energies
|
||||
Make it wizard makeup wielding wondrous wonders
|
||||
Make it genie makeup granting glorious gifts
|
||||
Make it werewolf makeup howling hallowed horizons
|
||||
Make it vampire makeup venerating vehement vigors
|
||||
Make it witch makeup weaving whimsical wonders
|
||||
Make it mermaid makeup maintaining mystical miracles
|
||||
Make it phoenix makeup promising perpetual progress
|
||||
Make it unicorn makeup upholding utopian unity
|
||||
Make it elf makeup ensuring ethereal enchantments
|
||||
Make it wizard makeup wielding wizened wisdom
|
||||
Make it genie makeup granting grand gestures
|
||||
Make it werewolf makeup howling haunting hymns
|
||||
Make it vampire makeup venerating vigorous victories
|
||||
Make it witch makeup weaving whimsical worlds
|
||||
Make it mermaid makeup maintaining mystical marvels
|
||||
Make it phoenix makeup portending perpetual prosperity
|
||||
Make it unicorn makeup upholding utopian understanding
|
||||
Make it elf makeup ensuring evanescent elation
|
||||
Make it wizard makeup wielding wizardly works
|
||||
Make it genie makeup granting glorious gifts
|
||||
Make it werewolf makeup howling hallowed harmonies
|
||||
Make it vampire makeup venerating vehement virtues
|
||||
Make it witch makeup warping wondrous wonders
|
||||
Make it mermaid makeup maintaining mesmeric mysteries
|
||||
Make it phoenix makeup promising perpetual perfection
|
||||
Make it unicorn makeup upholding utopian unity
|
||||
Make it elf makeup ensuring effervescent endeavors
|
||||
Make it wizard makeup wielding wizened wisdom
|
||||
Make it jorogumo makeup spinning sinister silk sorcery
|
||||
Make it inugami makeup masking merciless mystical motives
|
||||
Make it nurikabe makeup emerging earthen elegance
|
||||
Make it tenome makeup fighting fierce phantom forces
|
||||
Make it tanuki makeup taking trickster talents to task
|
||||
Make it Yamawaro makeup yodeling yields from young yokai realms
|
||||
Make it jikininki makeup displaying death's dreaded degradation
|
||||
Make it nopperabo makeup navigating nightly natural nuances
|
||||
Make it onikuma makeup merging monstrous mystique with mountain majesty
|
||||
@@ -0,0 +1,110 @@
|
||||
community_icon_html = """<svg id="share-btn-share-icon" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32">
|
||||
<path d="M20.6081 3C21.7684 3 22.8053 3.49196 23.5284 4.38415C23.9756 4.93678 24.4428 5.82749 24.4808 7.16133C24.9674 7.01707 25.4353 6.93643 25.8725 6.93643C26.9833 6.93643 27.9865 7.37587 28.696 8.17411C29.6075 9.19872 30.0124 10.4579 29.8361 11.7177C29.7523 12.3177 29.5581 12.8555 29.2678 13.3534C29.8798 13.8646 30.3306 14.5763 30.5485 15.4322C30.719 16.1032 30.8939 17.5006 29.9808 18.9403C30.0389 19.0342 30.0934 19.1319 30.1442 19.2318C30.6932 20.3074 30.7283 21.5229 30.2439 22.6548C29.5093 24.3704 27.6841 25.7219 24.1397 27.1727C21.9347 28.0753 19.9174 28.6523 19.8994 28.6575C16.9842 29.4379 14.3477 29.8345 12.0653 29.8345C7.87017 29.8345 4.8668 28.508 3.13831 25.8921C0.356375 21.6797 0.754104 17.8269 4.35369 14.1131C6.34591 12.058 7.67023 9.02782 7.94613 8.36275C8.50224 6.39343 9.97271 4.20438 12.4172 4.20438H12.4179C12.6236 4.20438 12.8314 4.2214 13.0364 4.25468C14.107 4.42854 15.0428 5.06476 15.7115 6.02205C16.4331 5.09583 17.134 4.359 17.7682 3.94323C18.7242 3.31737 19.6794 3 20.6081 3ZM20.6081 5.95917C20.2427 5.95917 19.7963 6.1197 19.3039 6.44225C17.7754 7.44319 14.8258 12.6772 13.7458 14.7131C13.3839 15.3952 12.7655 15.6837 12.2086 15.6837C11.1036 15.6837 10.2408 14.5497 12.1076 13.1085C14.9146 10.9402 13.9299 7.39584 12.5898 7.1776C12.5311 7.16799 12.4731 7.16355 12.4172 7.16355C11.1989 7.16355 10.6615 9.33114 10.6615 9.33114C10.6615 9.33114 9.0863 13.4148 6.38031 16.206C3.67434 18.998 3.5346 21.2388 5.50675 24.2246C6.85185 26.2606 9.42666 26.8753 12.0653 26.8753C14.8021 26.8753 17.6077 26.2139 19.1799 25.793C19.2574 25.7723 28.8193 22.984 27.6081 20.6107C27.4046 20.212 27.0693 20.0522 26.6471 20.0522C24.9416 20.0522 21.8393 22.6726 20.5057 22.6726C20.2076 22.6726 19.9976 22.5416 19.9116 22.222C19.3433 20.1173 28.552 19.2325 27.7758 16.1839C27.639 15.6445 27.2677 15.4256 26.746 15.4263C24.4923 15.4263 19.4358 19.5181 18.3759 19.5181C18.2949 19.5181 18.2368 19.4937 18.2053 19.4419C17.6743 18.557 17.9653 17.9394 21.7082 15.6009C25.4511 13.2617 28.0783 11.8545 26.5841 10.1752C26.4121 9.98141 26.1684 9.8956 25.8725 9.8956C23.6001 9.89634 18.2311 14.9403 18.2311 14.9403C18.2311 14.9403 16.7821 16.496 15.9057 16.496C15.7043 16.496 15.533 16.4139 15.4169 16.2112C14.7956 15.1296 21.1879 10.1286 21.5484 8.06535C21.7928 6.66715 21.3771 5.95917 20.6081 5.95917Z" fill="#FF9D00"></path>
|
||||
<path d="M5.50686 24.2246C3.53472 21.2387 3.67446 18.9979 6.38043 16.206C9.08641 13.4147 10.6615 9.33111 10.6615 9.33111C10.6615 9.33111 11.2499 6.95933 12.59 7.17757C13.93 7.39581 14.9139 10.9401 12.1069 13.1084C9.29997 15.276 12.6659 16.7489 13.7459 14.713C14.8258 12.6772 17.7747 7.44316 19.304 6.44221C20.8326 5.44128 21.9089 6.00204 21.5484 8.06532C21.188 10.1286 14.795 15.1295 15.4171 16.2118C16.0391 17.2934 18.2312 14.9402 18.2312 14.9402C18.2312 14.9402 25.0907 8.49588 26.5842 10.1752C28.0776 11.8545 25.4512 13.2616 21.7082 15.6008C17.9646 17.9393 17.6744 18.557 18.2054 19.4418C18.7372 20.3266 26.9998 13.1351 27.7759 16.1838C28.5513 19.2324 19.3434 20.1173 19.9117 22.2219C20.48 24.3274 26.3979 18.2382 27.6082 20.6107C28.8193 22.9839 19.2574 25.7722 19.18 25.7929C16.0914 26.62 8.24723 28.3726 5.50686 24.2246Z" fill="#FFD21E"></path>
|
||||
</svg>"""
|
||||
|
||||
loading_icon_html = """<svg id="share-btn-loading-icon" style="display:none;" class="animate-spin"
|
||||
style="color: #ffffff;
|
||||
"
|
||||
xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" fill="none" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 24 24"><circle style="opacity: 0.25;" cx="12" cy="12" r="10" stroke="white" stroke-width="4"></circle><path style="opacity: 0.75;" fill="white" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg>"""
|
||||
|
||||
share_js = r"""async () => {
|
||||
async function uploadFile(file){
|
||||
const UPLOAD_URL = 'https://huggingface.co/uploads';
|
||||
const response = await fetch(UPLOAD_URL, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': file.type,
|
||||
'X-Requested-With': 'XMLHttpRequest',
|
||||
},
|
||||
body: file,
|
||||
});
|
||||
const url = await response.text();
|
||||
return url;
|
||||
}
|
||||
|
||||
function getButtonText(componentId) {
|
||||
const buttonEl = gradioEl.querySelector(`${componentId} button`);
|
||||
return buttonEl ? buttonEl.textContent : '';
|
||||
}
|
||||
|
||||
const gradioEl = document.querySelector('body > gradio-app');
|
||||
const imgEls = [gradioEl.querySelector('#input_image img'), gradioEl.querySelector('#output_image img')];
|
||||
const concepts = [
|
||||
{ value: getButtonText('#box1'), parent: gradioEl.querySelector('#box1 span[data-testid="block-info"]') },
|
||||
{ value: getButtonText('#box2'), parent: gradioEl.querySelector('#box2 span[data-testid="block-info"]') },
|
||||
{ value: getButtonText('#box3'), parent: gradioEl.querySelector('#box3 span[data-testid="block-info"]') }
|
||||
];
|
||||
|
||||
const promptTxt = gradioEl.querySelector('#target_prompt input').value;
|
||||
const shareBtnEl = gradioEl.querySelector('#share-btn');
|
||||
const shareIconEl = gradioEl.querySelector('#share-btn-share-icon');
|
||||
const loadingIconEl = gradioEl.querySelector('#share-btn-loading-icon');
|
||||
if(!imgEls[1]){
|
||||
return;
|
||||
};
|
||||
shareBtnEl.style.pointerEvents = 'none';
|
||||
shareIconEl.style.display = 'none';
|
||||
loadingIconEl.style.removeProperty('display');
|
||||
|
||||
async function processImage(imgEl, imgId) {
|
||||
const res = await fetch(imgEl.src);
|
||||
const blob = await res.blob();
|
||||
const fileType = blob.type.includes('png') ? 'png' : 'jpg';
|
||||
const fileName = `diffuse-the-rest-${imgId}.${fileType}`;
|
||||
return new File([blob], fileName, { type: blob.type });
|
||||
}
|
||||
|
||||
const files = await Promise.all(imgEls.map((imgEl, index) => processImage(imgEl, Date.now() + index % 200)));
|
||||
const urls = await Promise.all(files.map((file) => uploadFile(file)));
|
||||
|
||||
const labels = ['Source image', 'Target image'];
|
||||
const htmlImgs = urls.map((url, index) => `<div>${labels[index]}: <img src='${url}' width='400' height='400' style="margin: 0"></div>`);
|
||||
|
||||
let descriptionMd = `<div style='display: flex; flex-wrap: wrap; column-gap: 0.75rem;'>${htmlImgs.join(`\n`)}</div>`;
|
||||
|
||||
if (promptTxt) {
|
||||
descriptionMd += `<span style="font-size:1.2em">Target image prompt: <b>${promptTxt}</b></span><br>`;
|
||||
} else {
|
||||
descriptionMd += `<span style="font-size:1.2em">Target image prompt: <b>""</b></span><br>`;
|
||||
}
|
||||
|
||||
const conceptHeaders = [];
|
||||
const conceptDescriptions = [];
|
||||
const conceptTableRows = [];
|
||||
concepts.forEach((concept, index) => {
|
||||
if (concept.value) {
|
||||
const label = concept.parent.textContent.includes('Negative') ? `remove concept` : `add concept`;
|
||||
conceptHeaders.push(`<th>${label}</th>`);
|
||||
conceptDescriptions.push(`${label}: ${concept.value}`);
|
||||
conceptTableRows.push(`<td>${concept.value}</td>`);
|
||||
}
|
||||
});
|
||||
|
||||
let title = 'Editing';
|
||||
if (promptTxt) {
|
||||
title += ` "${promptTxt}"`;
|
||||
}
|
||||
if (conceptDescriptions.length > 0) {
|
||||
title += ` to ${conceptDescriptions.join(', ')}`;
|
||||
descriptionMd += `<table style="font-size: 1.1em">
|
||||
<tr>
|
||||
${conceptHeaders.join('\n')}
|
||||
</tr>
|
||||
<tr>
|
||||
${conceptTableRows.join('\n')}
|
||||
</tr>
|
||||
</table>`;
|
||||
}
|
||||
|
||||
const params = new URLSearchParams({
|
||||
title: title,
|
||||
description: descriptionMd,
|
||||
preview: true,
|
||||
});
|
||||
const paramsStr = params.toString();
|
||||
window.open(`https://huggingface.co/spaces/editing-images/ledits/discussions/new?${paramsStr}`, '_blank');
|
||||
shareBtnEl.style.removeProperty('pointer-events');
|
||||
shareIconEl.style.removeProperty('display');
|
||||
loadingIconEl.style.display = 'none';
|
||||
}"""
|
||||
@@ -0,0 +1,114 @@
|
||||
import PIL
|
||||
from PIL import Image, ImageDraw ,ImageFont
|
||||
from matplotlib import pyplot as plt
|
||||
import torchvision.transforms as T
|
||||
import os
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
def show_torch_img(img):
|
||||
img = to_np_image(img)
|
||||
plt.imshow(img)
|
||||
plt.axis("off")
|
||||
|
||||
def to_np_image(all_images):
|
||||
all_images = (all_images.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu().numpy()[0]
|
||||
return all_images
|
||||
|
||||
def tensor_to_pil(tensor_imgs):
|
||||
if type(tensor_imgs) == list:
|
||||
tensor_imgs = torch.cat(tensor_imgs)
|
||||
tensor_imgs = (tensor_imgs / 2 + 0.5).clamp(0, 1)
|
||||
to_pil = T.ToPILImage()
|
||||
pil_imgs = [to_pil(img) for img in tensor_imgs]
|
||||
return pil_imgs
|
||||
|
||||
def pil_to_tensor(pil_imgs):
|
||||
to_torch = T.ToTensor()
|
||||
if type(pil_imgs) == PIL.Image.Image:
|
||||
tensor_imgs = to_torch(pil_imgs).unsqueeze(0)*2-1
|
||||
elif type(pil_imgs) == list:
|
||||
tensor_imgs = torch.cat([to_torch(pil_imgs).unsqueeze(0)*2-1 for img in pil_imgs]).to(device)
|
||||
else:
|
||||
raise Exception("Input need to be PIL.Image or list of PIL.Image")
|
||||
return tensor_imgs
|
||||
|
||||
|
||||
## TODO implement this
|
||||
# n = 10
|
||||
# num_rows = 4
|
||||
# num_col = n // num_rows
|
||||
# num_col = num_col + 1 if n % num_rows else num_col
|
||||
# num_col
|
||||
def add_margin(pil_img, top = 0, right = 0, bottom = 0,
|
||||
left = 0, color = (255,255,255)):
|
||||
width, height = pil_img.size
|
||||
new_width = width + right + left
|
||||
new_height = height + top + bottom
|
||||
result = Image.new(pil_img.mode, (new_width, new_height), color)
|
||||
|
||||
result.paste(pil_img, (left, top))
|
||||
return result
|
||||
|
||||
def image_grid(imgs, rows = 1, cols = None,
|
||||
size = None,
|
||||
titles = None, text_pos = (0, 0)):
|
||||
if type(imgs) == list and type(imgs[0]) == torch.Tensor:
|
||||
imgs = torch.cat(imgs)
|
||||
if type(imgs) == torch.Tensor:
|
||||
imgs = tensor_to_pil(imgs)
|
||||
|
||||
if not size is None:
|
||||
imgs = [img.resize((size,size)) for img in imgs]
|
||||
if cols is None:
|
||||
cols = len(imgs)
|
||||
assert len(imgs) >= rows*cols
|
||||
|
||||
top=20
|
||||
w, h = imgs[0].size
|
||||
delta = 0
|
||||
if len(imgs)> 1 and not imgs[1].size[1] == h:
|
||||
delta = top
|
||||
h = imgs[1].size[1]
|
||||
if not titles is None:
|
||||
font = ImageFont.truetype("/usr/share/fonts/truetype/freefont/FreeMono.ttf",
|
||||
size = 20, encoding="unic")
|
||||
h = top + h
|
||||
grid = Image.new('RGB', size=(cols*w, rows*h+delta))
|
||||
for i, img in enumerate(imgs):
|
||||
|
||||
if not titles is None:
|
||||
img = add_margin(img, top = top, bottom = 0,left=0)
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.text(text_pos, titles[i],(0,0,0),
|
||||
font = font)
|
||||
if not delta == 0 and i > 0:
|
||||
grid.paste(img, box=(i%cols*w, i//cols*h+delta))
|
||||
else:
|
||||
grid.paste(img, box=(i%cols*w, i//cols*h))
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
"""
|
||||
input_folder - dataset folder
|
||||
"""
|
||||
def load_dataset(input_folder):
|
||||
# full_file_names = glob.glob(input_folder)
|
||||
# class_names = [x[0] for x in os.walk(input_folder)]
|
||||
class_names = next(os.walk(input_folder))[1]
|
||||
class_names[:] = [d for d in class_names if not d[0] == '.']
|
||||
file_names=[]
|
||||
for class_name in class_names:
|
||||
cur_path = os.path.join(input_folder, class_name)
|
||||
filenames = next(os.walk(cur_path), (None, None, []))[2]
|
||||
filenames = [f for f in filenames if not f[0] == '.']
|
||||
file_names.append(filenames)
|
||||
return class_names, file_names
|
||||
|
||||
|
||||
def dataset_from_yaml(yaml_location):
|
||||
with open(yaml_location, 'r') as stream:
|
||||
data_loaded = yaml.safe_load(stream)
|
||||
|
||||
return data_loaded
|
||||
@@ -0,0 +1,211 @@
|
||||
import numpy as np
|
||||
import tqdm
|
||||
from PIL import Image
|
||||
from datasets import load_dataset
|
||||
from ..spiga.inference.config import ModelConfig
|
||||
from ..spiga.inference.framework import SPIGAFramework
|
||||
from ..facelib import FaceDetector
|
||||
import cv2
|
||||
import os
|
||||
|
||||
processor = SPIGAFramework(ModelConfig("300wpublic"))
|
||||
|
||||
def center_crop(image, size):
|
||||
width, height = image.size
|
||||
left = (width - size) // 2
|
||||
top = (height - size) // 2
|
||||
right = left + size
|
||||
bottom = top + size
|
||||
cropped_image = image.crop((left, top, right, bottom))
|
||||
return cropped_image
|
||||
|
||||
def resize(image, size):
|
||||
width, height = image.size
|
||||
if width > height:
|
||||
# 按宽度进行比例调整
|
||||
new_width = size
|
||||
new_height = int(height * (size / width))
|
||||
else:
|
||||
# 按高度进行比例调整
|
||||
new_height = size
|
||||
new_width = int(width * (size / height))
|
||||
resized_image = image.resize((new_width, new_height))
|
||||
return resized_image
|
||||
|
||||
def preprocess(example, name, path):
|
||||
image = resize(example, 512)
|
||||
# 调用中心剪裁函数
|
||||
cropped_image = center_crop(image, 512)
|
||||
# 保存剪裁后的图像
|
||||
cropped_image.save(path+name)
|
||||
return cropped_image
|
||||
|
||||
# We obtain the bbox from the existing landmarks in the dataset.
|
||||
# We could use `dlib`, but this should be faster.
|
||||
# Note that the `landmarks` are stored as strings.
|
||||
|
||||
def get_landmarks(image, detector):
|
||||
image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
|
||||
faces, boxes, scores, landmarks = detector.detect_align(image) # 一定要用align啊
|
||||
boxes = boxes.cpu().numpy()
|
||||
box_ls = []
|
||||
for box in boxes:
|
||||
x, y, x1, y1 = box
|
||||
box = x, y, x1 - x, y1 - y
|
||||
box_ls.append(box)
|
||||
if len(box_ls) == 0:
|
||||
return []
|
||||
else:
|
||||
features = processor.inference(image, box_ls)
|
||||
landmarks = np.array(features['landmarks'])
|
||||
return landmarks
|
||||
|
||||
|
||||
def parse_landmarks(landmarks):
|
||||
ldm = []
|
||||
for landmark in landmarks:
|
||||
ldm.append([(float(x), float(y)) for x, y in landmark])
|
||||
return ldm
|
||||
|
||||
|
||||
def bbox_from_landmarks(landmarks_):
|
||||
landmarks = parse_landmarks(landmarks_)
|
||||
bbox = []
|
||||
for ldm in landmarks:
|
||||
landmarks_x, landmarks_y = zip(*ldm)
|
||||
|
||||
x_min, x_max = min(landmarks_x), max(landmarks_x)
|
||||
y_min, y_max = min(landmarks_y), max(landmarks_y)
|
||||
width = x_max - x_min
|
||||
height = y_max - y_min
|
||||
|
||||
# Give it a little room; I think it works anyway
|
||||
x_min -= 5
|
||||
y_min -= 5
|
||||
width += 10
|
||||
height += 10
|
||||
bbox.append((x_min, y_min, width, height))
|
||||
return bbox
|
||||
|
||||
|
||||
def spiga_process(example, detector):
|
||||
ldms = get_landmarks(example, detector)
|
||||
|
||||
if len(ldms) == 0:
|
||||
return False
|
||||
|
||||
else:
|
||||
image = example
|
||||
image = np.array(image)
|
||||
# BGR
|
||||
image = image[:, :, ::-1]
|
||||
bbox = bbox_from_landmarks(ldms)
|
||||
features = processor.inference(image, [*bbox])
|
||||
landmarks = features["landmarks"]
|
||||
spigas = landmarks
|
||||
return spigas
|
||||
|
||||
|
||||
# For some reason this map doesn't work with num_proc > 1 :(
|
||||
# TODO: run inference on GPU
|
||||
|
||||
|
||||
# ## "Segmentation"
|
||||
|
||||
# We use bezier paths to draw contours and areas.
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
from matplotlib.path import Path
|
||||
import PIL
|
||||
|
||||
|
||||
def get_patch(landmarks, color='lime', closed=False):
|
||||
contour = landmarks
|
||||
ops = [Path.MOVETO] + [Path.LINETO] * (len(contour) - 1)
|
||||
facecolor = (0, 0, 0, 0) # Transparent fill color, if open
|
||||
if closed:
|
||||
contour.append(contour[0])
|
||||
ops.append(Path.CLOSEPOLY)
|
||||
facecolor = color
|
||||
path = Path(contour, ops)
|
||||
return patches.PathPatch(path, facecolor=facecolor, edgecolor=color, lw=4)
|
||||
|
||||
|
||||
# Draw to a buffer.
|
||||
|
||||
def conditioning_from_landmarks(landmarks_, size=512):
|
||||
# Precisely control output image size
|
||||
dpi = 72
|
||||
fig, ax = plt.subplots(1, figsize=[size / dpi, size / dpi], tight_layout={'pad': 0})
|
||||
fig.set_dpi(dpi)
|
||||
|
||||
black = np.zeros((size, size, 3))
|
||||
ax.imshow(black)
|
||||
|
||||
for landmarks in landmarks_:
|
||||
face_patch = get_patch(landmarks[0:17])
|
||||
l_eyebrow = get_patch(landmarks[17:22], color='yellow')
|
||||
r_eyebrow = get_patch(landmarks[22:27], color='yellow')
|
||||
nose_v = get_patch(landmarks[27:31], color='orange')
|
||||
nose_h = get_patch(landmarks[31:36], color='orange')
|
||||
l_eye = get_patch(landmarks[36:42], color='magenta', closed=True)
|
||||
r_eye = get_patch(landmarks[42:48], color='magenta', closed=True)
|
||||
outer_lips = get_patch(landmarks[48:60], color='cyan', closed=True)
|
||||
inner_lips = get_patch(landmarks[60:68], color='blue', closed=True)
|
||||
|
||||
ax.add_patch(face_patch)
|
||||
ax.add_patch(l_eyebrow)
|
||||
ax.add_patch(r_eyebrow)
|
||||
ax.add_patch(nose_v)
|
||||
ax.add_patch(nose_h)
|
||||
ax.add_patch(l_eye)
|
||||
ax.add_patch(r_eye)
|
||||
ax.add_patch(outer_lips)
|
||||
ax.add_patch(inner_lips)
|
||||
|
||||
plt.axis('off')
|
||||
|
||||
fig.canvas.draw()
|
||||
buffer, (width, height) = fig.canvas.print_to_buffer()
|
||||
assert width == height
|
||||
assert width == size
|
||||
|
||||
buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
|
||||
buffer = buffer[:, :, 0:3]
|
||||
plt.close(fig)
|
||||
return PIL.Image.fromarray(buffer)
|
||||
|
||||
|
||||
def spiga_segmentation(spiga):
|
||||
landmarks = spiga
|
||||
spiga_seg = conditioning_from_landmarks(landmarks)
|
||||
return spiga_seg
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# ## Obtain SPIGA features
|
||||
processor = SPIGAFramework(ModelConfig("300wpublic"))
|
||||
detector = FaceDetector(weight_path="mobilenet0.25_Final.pth")
|
||||
|
||||
id_folder = "id"
|
||||
pose_folder = "pose"
|
||||
|
||||
if not os.path.exists(pose_folder):
|
||||
os.makedirs(pose_folder)
|
||||
|
||||
pbar = tqdm.tqdm(os.listdir(id_folder))
|
||||
for name in pbar:
|
||||
face = Image.open(id_folder+name).convert("RGB").resize((512, 512))
|
||||
face.save(id_folder+name)
|
||||
spigas = spiga_process(face, detector)
|
||||
if spigas == False:
|
||||
height = 512
|
||||
width = 512
|
||||
channels = 3
|
||||
black_image = np.zeros((height, width, channels), dtype=np.uint8)
|
||||
black_image_cv2 = cv2.cvtColor(black_image, cv2.COLOR_RGB2BGR)
|
||||
continue
|
||||
else:
|
||||
spigas_faces = spiga_segmentation(spigas)
|
||||
spigas_faces.save(pose_folder + name)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,687 @@
|
||||
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
from torchvision import transforms
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
class SSRAttnProcessor(nn.Module):
|
||||
r"""
|
||||
Attention processor for SSR-Adapater.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
# self.to_q_SSR = nn.Linear(hidden_size, hidden_size, bias=False)
|
||||
self.to_k_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
self.to_v_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
# query = self.to_q_SSR(hidden_states)
|
||||
query = attn.to_q(hidden_states)
|
||||
query = attn.head_to_batch_dim(query)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
_hidden_states = encoder_hidden_states
|
||||
_key = self.to_k_SSR(_hidden_states)
|
||||
_value = self.to_v_SSR(_hidden_states)
|
||||
_key = attn.head_to_batch_dim(_key)
|
||||
_value = attn.head_to_batch_dim(_value)
|
||||
_attention_probs = attn.get_attention_scores(query, _key, None)
|
||||
_hidden_states = torch.bmm(_attention_probs, _value)
|
||||
_hidden_states = attn.batch_to_head_dim(_hidden_states)
|
||||
hidden_states = self.scale * _hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SSRAttnProcessor2_0(torch.nn.Module):
|
||||
r"""
|
||||
Attention processor for SSR-Adapater for PyTorch 2.0.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0):
|
||||
super().__init__()
|
||||
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
# self.to_q_SSR = nn.Linear(hidden_size, hidden_size, bias=False)
|
||||
self.to_k_SSR = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
self.to_v_SSR = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
# query = self.to_q_SSR(hidden_states)
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
# split hidden states
|
||||
_hidden_states = encoder_hidden_states
|
||||
|
||||
_key = self.to_k_SSR(_hidden_states)
|
||||
_value = self.to_v_SSR(_hidden_states)
|
||||
inner_dim = _key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
_key = _key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
_value = _value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
_hidden_states = F.scaled_dot_product_attention(
|
||||
query, _key, _value, attn_mask=None, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
_hidden_states = _hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
_hidden_states = _hidden_states.to(query.dtype)
|
||||
|
||||
hidden_states = self.scale * _hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class AttnProcessor2_0(torch.nn.Module):
|
||||
r"""
|
||||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=None,
|
||||
cross_attention_dim=None,
|
||||
):
|
||||
super().__init__()
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
class AttnProcessor(nn.Module):
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=None,
|
||||
cross_attention_dim=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ConvAttnProcessor:
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
):
|
||||
## map to 2D
|
||||
if len(hidden_states.shape) == 4:
|
||||
shape = hidden_states.shape
|
||||
hidden_states = torch.reshape(hidden_states, (shape[0], shape[1], shape[2] * shape[3]))
|
||||
hidden_states = hidden_states.permute(0, 2, 1)
|
||||
if encoder_hidden_states is not None:
|
||||
if len(encoder_hidden_states.shape) == 4:
|
||||
kv_shape = encoder_hidden_states.shape
|
||||
encoder_hidden_states = torch.reshape(
|
||||
encoder_hidden_states, (kv_shape[0], kv_shape[1], kv_shape[2] * kv_shape[3])
|
||||
)
|
||||
encoder_hidden_states = encoder_hidden_states.permute(0, 2, 1)
|
||||
|
||||
# the same to standard attn
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
# map back to 4D
|
||||
if len(hidden_states.shape) == 3:
|
||||
hidden_states = hidden_states.permute(0, 2, 1)
|
||||
hidden_states = torch.reshape(hidden_states, (shape[0], shape[1], shape[2], shape[3]))
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SSRAttnProcessor_text(nn.Module):
|
||||
r"""
|
||||
Attention processor for SSR-Adapater.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1):
|
||||
super().__init__()
|
||||
self.text_context_len = 77
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.to_k_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
self.to_v_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
query = attn.head_to_batch_dim(query)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
# split hidden states
|
||||
encoder_hidden_states, _hidden_states = encoder_hidden_states[:, :self.text_context_len,
|
||||
:], encoder_hidden_states[:, self.text_context_len:, :]
|
||||
encoder_hidden_states = encoder_hidden_states[:, :, :768]
|
||||
# for text
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# for image
|
||||
_key = self.to_k_SSR(_hidden_states)
|
||||
_value = self.to_v_SSR(_hidden_states)
|
||||
_key = attn.head_to_batch_dim(_key)
|
||||
_value = attn.head_to_batch_dim(_value)
|
||||
_attention_probs = attn.get_attention_scores(query, _key, None)
|
||||
_hidden_states = torch.bmm(_attention_probs, _value)
|
||||
_hidden_states = attn.batch_to_head_dim(_hidden_states)
|
||||
hidden_states = self.scale * _hidden_states + hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SSRAttnProcessor2_0_text(torch.nn.Module):
|
||||
r"""
|
||||
Attention processor for SSR-Adapater for PyTorch 2.0.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0):
|
||||
super().__init__()
|
||||
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
self.text_context_len = 77
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.to_k_SSR = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
self.to_v_SSR = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
# split hidden states
|
||||
encoder_hidden_states, _hidden_states = encoder_hidden_states[:, :self.text_context_len,
|
||||
:], encoder_hidden_states[:, self.text_context_len:, :]
|
||||
|
||||
encoder_hidden_states = encoder_hidden_states[:, :, :768]
|
||||
# for text
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p = 0.0, is_causal = False
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# for image
|
||||
_key = self.to_k_SSR(_hidden_states)
|
||||
_value = self.to_v_SSR(_hidden_states)
|
||||
inner_dim = _key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
_key = _key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
_value = _value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
_hidden_states = F.scaled_dot_product_attention(
|
||||
query, _key, _value, attn_mask=None, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
_hidden_states = _hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
_hidden_states = _hidden_states.to(query.dtype)
|
||||
|
||||
hidden_states = self.scale * _hidden_states + hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SSRAttnProcessor_visual(nn.Module):
|
||||
r"""
|
||||
Attention processor for attn visualization.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1, attnstore=None, place_in_unet=None):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.to_k_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
self.to_v_SSR = nn.Linear(cross_attention_dim, hidden_size, bias=False)
|
||||
self.attnstore = attnstore
|
||||
self.place_in_unet = place_in_unet
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
# query = self.to_q_SSR(hidden_states)
|
||||
query = attn.to_q(hidden_states)
|
||||
query = attn.head_to_batch_dim(query)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
_hidden_states = encoder_hidden_states
|
||||
_key = self.to_k_SSR(_hidden_states)
|
||||
_value = self.to_v_SSR(_hidden_states)
|
||||
_key = attn.head_to_batch_dim(_key)
|
||||
_value = attn.head_to_batch_dim(_value)
|
||||
_attention_probs = attn.get_attention_scores(query, _key, None)
|
||||
|
||||
# store attention maps
|
||||
is_cross = encoder_hidden_states is not None
|
||||
self.attnstore(_attention_probs, is_cross, self.place_in_unet)
|
||||
|
||||
_hidden_states = torch.bmm(_attention_probs, _value)
|
||||
_hidden_states = attn.batch_to_head_dim(_hidden_states)
|
||||
hidden_states = self.scale * _hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,113 @@
|
||||
from typing import List
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from transformers import CLIPImageProcessor
|
||||
from transformers import CLIPVisionModel as OriginalCLIPVisionModel
|
||||
from ._clip import CLIPVisionModel
|
||||
from PIL import Image
|
||||
import torch.nn.functional as F
|
||||
import torch.nn as nn
|
||||
import os
|
||||
|
||||
def is_torch2_available():
|
||||
return hasattr(F, "scaled_dot_product_attention")
|
||||
if is_torch2_available():
|
||||
from .attention_processor import SSRAttnProcessor2_0 as SSRAttnProcessor, AttnProcessor2_0 as AttnProcessor
|
||||
else:
|
||||
from .attention_processor import SSRAttnProcessor, AttnProcessor
|
||||
from .resampler import Resampler
|
||||
|
||||
class detail_encoder(torch.nn.Module):
|
||||
"""from SSR-encoder"""
|
||||
def __init__(self, unet, image_encoder_path, device="cuda", dtype=torch.float32):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
# load image encoder
|
||||
clip_encoder = OriginalCLIPVisionModel.from_pretrained(image_encoder_path)
|
||||
self.image_encoder = CLIPVisionModel(clip_encoder.config)
|
||||
state_dict = clip_encoder.state_dict()
|
||||
self.image_encoder.load_state_dict(state_dict, strict=False)
|
||||
self.image_encoder.to(self.device, self.dtype)
|
||||
del clip_encoder
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
|
||||
# load SSR layers
|
||||
attn_procs = {}
|
||||
for name in unet.attn_processors.keys():
|
||||
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
|
||||
if name.startswith("mid_block"):
|
||||
hidden_size = unet.config.block_out_channels[-1]
|
||||
elif name.startswith("up_blocks"):
|
||||
block_id = int(name[len("up_blocks.")])
|
||||
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
|
||||
elif name.startswith("down_blocks"):
|
||||
block_id = int(name[len("down_blocks.")])
|
||||
hidden_size = unet.config.block_out_channels[block_id]
|
||||
if cross_attention_dim is None:
|
||||
attn_procs[name] = AttnProcessor()
|
||||
else:
|
||||
attn_procs[name] = SSRAttnProcessor(hidden_size=hidden_size, cross_attention_dim=1024, scale=1).to(self.device, dtype=self.dtype)
|
||||
unet.set_attn_processor(attn_procs)
|
||||
adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
|
||||
self.SSR_layers = adapter_modules
|
||||
self.SSR_layers.to(self.device, dtype=self.dtype)
|
||||
self.resampler = self.init_proj()
|
||||
|
||||
def init_proj(self):
|
||||
resampler = Resampler().to(self.device, dtype=self.dtype)
|
||||
return resampler
|
||||
|
||||
def forward(self, img):
|
||||
image_embeds = self.image_encoder(img, output_hidden_states=True)['hidden_states'][2::2]
|
||||
image_embeds = torch.cat(image_embeds, dim=1)
|
||||
image_embeds = self.resampler(image_embeds)
|
||||
return image_embeds
|
||||
|
||||
@torch.inference_mode()
|
||||
def get_image_embeds(self, pil_image):
|
||||
if isinstance(pil_image, Image.Image):
|
||||
pil_image = [pil_image]
|
||||
clip_image = []
|
||||
for pil in pil_image:
|
||||
tensor_image = self.clip_image_processor(images=pil, return_tensors="pt").pixel_values.to(self.device, dtype=self.dtype)
|
||||
clip_image.append(tensor_image)
|
||||
clip_image = torch.cat(clip_image, dim=0)
|
||||
|
||||
# cond
|
||||
clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True)['hidden_states'][2::2] # 1 257*12 1024
|
||||
clip_image_embeds = torch.cat(clip_image_embeds, dim=1)
|
||||
uncond_clip_image_embeds = self.image_encoder(torch.zeros_like(clip_image), output_hidden_states=True)['hidden_states'][2::2]
|
||||
uncond_clip_image_embeds = torch.cat(uncond_clip_image_embeds, dim=1)
|
||||
clip_image_embeds = self.resampler(clip_image_embeds)
|
||||
uncond_clip_image_embeds = self.resampler(uncond_clip_image_embeds)
|
||||
return clip_image_embeds, uncond_clip_image_embeds
|
||||
|
||||
def generate(
|
||||
self,
|
||||
id_image,
|
||||
makeup_image,
|
||||
seed=None,
|
||||
guidance_scale=2,
|
||||
num_inference_steps=30,
|
||||
pipe=None,
|
||||
**kwargs,
|
||||
):
|
||||
image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(makeup_image)
|
||||
|
||||
prompt_embeds = image_prompt_embeds
|
||||
negative_prompt_embeds = uncond_image_prompt_embeds
|
||||
|
||||
generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
|
||||
image = pipe(
|
||||
image=id_image,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
guidance_scale=guidance_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
generator=generator,
|
||||
**kwargs,
|
||||
).images[0]
|
||||
|
||||
return image
|
||||
@@ -0,0 +1,112 @@
|
||||
import torch
|
||||
import math
|
||||
import torch.nn.functional as F
|
||||
from torch import nn, einsum
|
||||
from inspect import isfunction
|
||||
|
||||
|
||||
def exists(val):
|
||||
return val is not None
|
||||
|
||||
def uniq(arr):
|
||||
return{el: True for el in arr}.keys()
|
||||
|
||||
|
||||
def default(val, d):
|
||||
if exists(val):
|
||||
return val
|
||||
return d() if isfunction(d) else d
|
||||
|
||||
|
||||
def max_neg_value(t):
|
||||
return -torch.finfo(t.dtype).max
|
||||
|
||||
|
||||
def init_(tensor):
|
||||
dim = tensor.shape[-1]
|
||||
std = 1 / math.sqrt(dim)
|
||||
tensor.uniform_(-std, std)
|
||||
return tensor
|
||||
|
||||
|
||||
# feedforward
|
||||
class GEGLU(nn.Module):
|
||||
def __init__(self, dim_in, dim_out):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(dim_in, dim_out * 2)
|
||||
|
||||
def forward(self, x):
|
||||
x, gate = self.proj(x).chunk(2, dim=-1)
|
||||
return x * F.gelu(gate)
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, dim, dim_out=None, mult=4, glu=True, dropout=0.):
|
||||
super().__init__()
|
||||
inner_dim = int(dim * mult)
|
||||
dim_out = default(dim_out, dim)
|
||||
project_in = nn.Sequential(
|
||||
nn.Linear(dim, inner_dim),
|
||||
nn.GELU()
|
||||
) if not glu else GEGLU(dim, inner_dim)
|
||||
|
||||
self.net = nn.Sequential(
|
||||
project_in,
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(inner_dim, dim_out)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, query_dim, heads=8, dim_head=64, dropout=0.):
|
||||
super().__init__()
|
||||
inner_dim = dim_head * heads
|
||||
self.scale = dim_head ** -0.5
|
||||
self.heads = heads
|
||||
|
||||
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
|
||||
self.to_k = nn.Linear(query_dim, inner_dim, bias=False)
|
||||
self.to_v = nn.Linear(query_dim, inner_dim, bias=False)
|
||||
|
||||
self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) )
|
||||
|
||||
def forward(self, x):
|
||||
q = self.to_q(x) # B*N*(H*C)
|
||||
k = self.to_k(x) # B*N*(H*C)
|
||||
v = self.to_v(x) # B*N*(H*C)
|
||||
|
||||
B, N, HC = q.shape
|
||||
H = self.heads
|
||||
C = HC // H
|
||||
|
||||
q = q.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C
|
||||
k = k.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C
|
||||
v = v.view(B,N,H,C).permute(0,2,1,3).reshape(B*H,N,C) # (B*H)*N*C
|
||||
|
||||
sim = torch.einsum('b i c, b j c -> b i j', q, k) * self.scale # (B*H)*N*N
|
||||
attn = sim.softmax(dim=-1) # (B*H)*N*N
|
||||
|
||||
out = torch.einsum('b i j, b j c -> b i c', attn, v) # (B*H)*N*C
|
||||
out = out.view(B,H,N,C).permute(0,2,1,3).reshape(B,N,(H*C)) # B*N*(H*C)
|
||||
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
|
||||
class Resampler(nn.Module):
|
||||
def __init__(self, query_dim=1024, n_heads=8, d_head=64):
|
||||
super().__init__()
|
||||
|
||||
self.attn = SelfAttention(query_dim=query_dim, heads=n_heads, dim_head=d_head)
|
||||
self.ff = FeedForward(query_dim, glu=True)
|
||||
|
||||
self.norm1 = nn.LayerNorm(query_dim)
|
||||
self.norm2 = nn.LayerNorm(query_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.attn(self.norm1(x))
|
||||
x = x + self.ff(self.norm2(x))
|
||||
return x
|
||||
@@ -0,0 +1,62 @@
|
||||
from .models.model import ShuffleneTiny, ShuffleneFull
|
||||
from ..utils import download_weight
|
||||
import torch
|
||||
import os
|
||||
|
||||
class AgeGenderEstimator:
|
||||
|
||||
def __init__(self, name='full', weight_path=None, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
"""
|
||||
Age and gender Detector
|
||||
:param name: name of backbone (full or tiny)
|
||||
:param device: model run in cpu or gpu (cuda, cpu)
|
||||
:param weight_path: path of network weight
|
||||
|
||||
Notice: image size must be 112x112
|
||||
but cun run with 224x224
|
||||
|
||||
Method detect:
|
||||
:param faces: 4d tensor of face for example size(1, 3, 112, 112)
|
||||
:returns genders list and ages list
|
||||
"""
|
||||
if name == 'tiny':
|
||||
model = ShuffleneTiny()
|
||||
elif name == 'full':
|
||||
model = ShuffleneFull()
|
||||
else:
|
||||
exit('from AgeGender Detector: model dose not support just(tiny, full)')
|
||||
|
||||
# download the default weigth
|
||||
if weight_path is None:
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'ShufflenetFull.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/ShufflenetFull.pth')
|
||||
if os.path.isfile(weight_path) == False:
|
||||
print('from AgeGenderEstimator: download defualt weight started')
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://drive.google.com/uc?export=download&id=1rnOZo46RjGZYrUb6Wup6sSOP37ol5I9E', file_name=file_name)
|
||||
os.rename(file_name, weight_path)
|
||||
|
||||
model.load_state_dict(torch.load(weight_path, map_location=device))
|
||||
model.to(device).eval()
|
||||
self.model = model
|
||||
self.device = device
|
||||
print('from AgeGenderEstimator: weights loaded')
|
||||
|
||||
def detect(self, faces):
|
||||
faces = faces.permute(0, 3, 1, 2)
|
||||
faces = faces.float().div(255).to(self.device)
|
||||
|
||||
mu = torch.as_tensor([0.485, 0.456, 0.406], dtype=faces.dtype, device=faces.device)
|
||||
std = torch.as_tensor([0.229, 0.224, 0.225], dtype=faces.dtype, device=faces.device)
|
||||
faces[:].sub_(mu[:, None, None]).div_(std[:, None, None])
|
||||
|
||||
outputs = self.model(faces)
|
||||
genders = []
|
||||
ages = []
|
||||
for out in outputs:
|
||||
gender = torch.argmax(out[:2])
|
||||
gender = 'Male' if gender == 0 else 'Female'
|
||||
genders.append(gender)
|
||||
ages.append(int(out[-1]))
|
||||
|
||||
return genders, ages
|
||||
@@ -0,0 +1,37 @@
|
||||
# from facelib import FaceDetector
|
||||
# from facelib import AgeGenderEstimator
|
||||
# from facelib import special_draw
|
||||
|
||||
from ..Retinaface.Retinaface import FaceDetector
|
||||
from ..AgeGender.Detector import AgeGenderEstimator
|
||||
from ..InsightFace.models.utils import special_draw
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
class WebcamAgeGenderEstimator:
|
||||
def __init__(self, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
print('loading ...')
|
||||
self.face_detector = FaceDetector(device=device)
|
||||
self.age_gender_detector = AgeGenderEstimator(device=device)
|
||||
|
||||
|
||||
def run(self, camera_index=0):
|
||||
cap = cv2.VideoCapture(camera_index)
|
||||
cap.set(3, 1280)
|
||||
cap.set(4, 720)
|
||||
print('type q for exit')
|
||||
while cap.isOpened():
|
||||
ret , frame = cap.read()
|
||||
if ret == False:
|
||||
raise Exception('the camera not recognized: change camera_index param to ' + str(0 if camera_index == 1 else 1))
|
||||
faces, boxes, scores, landmarks = self.face_detector.detect_align(frame)
|
||||
if len(faces.shape) > 1:
|
||||
genders, ages = self.age_gender_detector.detect(faces)
|
||||
for i, b in enumerate(boxes):
|
||||
special_draw(frame, b, landmarks[i], name=genders[i]+' '+str(ages[i]))
|
||||
|
||||
cv2.imshow('frame', frame)
|
||||
if cv2.waitKey(1) == ord('q'):
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,144 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.models as models
|
||||
import time
|
||||
import sys
|
||||
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
sz = 112
|
||||
|
||||
|
||||
class ShuffleneTiny(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super(ShuffleneTiny, self).__init__()
|
||||
self.model = models.shufflenet_v2_x0_5()
|
||||
self.model.fc = nn.Sequential(
|
||||
nn.BatchNorm1d(1024),
|
||||
nn.Linear(1024, 3))
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
|
||||
class ShuffleneFull(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super(ShuffleneFull, self).__init__()
|
||||
self.model = models.shufflenet_v2_x1_0()
|
||||
self.model.fc = nn.Sequential(
|
||||
nn.BatchNorm1d(1024),
|
||||
nn.Linear(1024, 3))
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
|
||||
class TrainModel:
|
||||
|
||||
def __init__(self, model, train_dl, valid_dl, optimizer, certrion, scheduler, num_epochs):
|
||||
|
||||
self.num_epochs = num_epochs
|
||||
self.model = model
|
||||
self.scheduler = scheduler
|
||||
self.train_dl = train_dl
|
||||
self.valid_dl = valid_dl
|
||||
self.optimizer = optimizer
|
||||
self.certrion = certrion
|
||||
|
||||
self.loss_history = []
|
||||
self.best_acc_valid = 0.0
|
||||
self.best_wieght = None
|
||||
|
||||
self.training()
|
||||
|
||||
def training(self):
|
||||
|
||||
valid_acc = 0
|
||||
for epoch in range(self.num_epochs):
|
||||
|
||||
print('Epoch %2d/%2d' % (epoch + 1, self.num_epochs))
|
||||
print('-' * 15)
|
||||
|
||||
t0 = time.time()
|
||||
train_acc = self.train_model()
|
||||
valid_acc = self.valid_model()
|
||||
if self.scheduler:
|
||||
self.scheduler.step()
|
||||
|
||||
time_elapsed = time.time() - t0
|
||||
print(' Training complete in: %.0fm %.0fs' % (time_elapsed // 60, time_elapsed % 60))
|
||||
print('| val_acc_gender | val_l1_loss | acc_gender | l1_loss |')
|
||||
print('| %.3f | %.3f | %.3f | %.3f \n' % (valid_acc[0], valid_acc[1], train_acc[0], train_acc[1]))
|
||||
|
||||
if valid_acc[0] > self.best_acc_valid:
|
||||
self.best_acc_valid = valid_acc[1]
|
||||
self.best_wieght = self.model.state_dict().copy()
|
||||
return
|
||||
|
||||
def train_model(self):
|
||||
|
||||
self.model.train()
|
||||
N = len(self.train_dl.dataset)
|
||||
step = N // self.train_dl.batch_size
|
||||
|
||||
avg_loss = 0.0
|
||||
acc_gender = 0.0
|
||||
loss_age = 0.0
|
||||
|
||||
for i, (x, y) in enumerate(self.train_dl):
|
||||
x, y = x.cuda(), y.cuda()
|
||||
# forward
|
||||
pred_8 = self.model(x)
|
||||
# loss
|
||||
loss = self.certrion(pred_8, y)
|
||||
# backward
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
# statistics of model training
|
||||
avg_loss = (avg_loss * i + loss) / (i + 1)
|
||||
acc_gender += accuracy_gender(pred_8, y)
|
||||
loss_age += l1loss_age(pred_8, y)
|
||||
|
||||
self.loss_history.append(avg_loss)
|
||||
|
||||
# report statistics
|
||||
sys.stdout.flush()
|
||||
sys.stdout.write("\r Train_Step: %d/%d | runing_loss: %.4f" % (i + 1, step, avg_loss))
|
||||
|
||||
sys.stdout.flush()
|
||||
return torch.tensor([acc_gender, loss_age]) / N
|
||||
|
||||
def valid_model(self):
|
||||
print()
|
||||
self.model.eval()
|
||||
N = len(self.valid_dl.dataset)
|
||||
step = N // self.valid_dl.batch_size
|
||||
acc_gender = 0.0
|
||||
loss_age = 0.0
|
||||
|
||||
with torch.no_grad():
|
||||
for i, (x, y) in enumerate(self.valid_dl):
|
||||
x, y = x.cuda(), y.cuda()
|
||||
|
||||
score = self.model(x)
|
||||
acc_gender += accuracy_gender(score, y)
|
||||
loss_age += l1loss_age(score, y)
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stdout.write("\r Vaild_Step: %d/%d" % (i, step))
|
||||
|
||||
sys.stdout.flush()
|
||||
return torch.tensor([acc_gender, loss_age]) / N
|
||||
|
||||
|
||||
def accuracy_gender(input, targs):
|
||||
pred = torch.argmax(input[:, :2], dim=1)
|
||||
y = targs[:, 0]
|
||||
return torch.sum(pred == y)
|
||||
|
||||
|
||||
def l1loss_age(input, targs):
|
||||
return F.l1_loss(input[:, -1], targs[:, -1]).mean()
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
Sajjad Ayoubi: Age Gender Detection
|
||||
I use UTKFace DataSet
|
||||
from: https://susanqq.github.io/UTKFace/
|
||||
download it and put it on FaceSet dir
|
||||
and I create a annotation file data.npy
|
||||
which there is in weights folder
|
||||
"""
|
||||
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
import torchvision.transforms as transforms
|
||||
from .model import ShuffleneFull, TrainModel
|
||||
|
||||
|
||||
class MultitaskDataset(Dataset):
|
||||
|
||||
def __init__(self, data, tfms, root='FaceSet/'):
|
||||
self.root = root
|
||||
self.tfms = tfms
|
||||
self.ages = data[:, 3]
|
||||
self.races = data[:, 2]
|
||||
self.genders = data[:, 1]
|
||||
self.imgs = data[:, 0]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs)
|
||||
|
||||
def __getitem__(self, i):
|
||||
return self.tfms(Image.open(self.root + self.imgs[i])), torch.tensor(
|
||||
[self.genders[i], self.races[i], self.ages[i]]).float()
|
||||
|
||||
def __repr__(self):
|
||||
return f'{type(self).__name__} of len {len(self)}'
|
||||
|
||||
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
sz = 112
|
||||
bs = 256
|
||||
|
||||
tf = {'train': transforms.Compose([
|
||||
transforms.RandomRotation(degrees=0.2),
|
||||
transforms.RandomHorizontalFlip(p=.5),
|
||||
transforms.RandomGrayscale(p=.2),
|
||||
transforms.Resize((sz, sz)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std)]),
|
||||
'test': transforms.Compose([
|
||||
transforms.Resize((sz, sz)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std)])}
|
||||
|
||||
data = np.load('data.npy', allow_pickle=True)
|
||||
train_data = data[data[:, -1] == 1]
|
||||
valid_data = data[data[:, -1] == 0]
|
||||
|
||||
valid_ds = MultitaskDataset(data=valid_data, tfms=tf['test'])
|
||||
train_ds = MultitaskDataset(data=train_data, tfms=tf['train'])
|
||||
|
||||
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True, num_workers=8)
|
||||
valid_dl = DataLoader(valid_ds, batch_size=bs, shuffle=True, num_workers=4)
|
||||
|
||||
|
||||
def multitask_loss(input, target):
|
||||
input_gender = input[:, :2]
|
||||
input_age = input[:, -1]
|
||||
|
||||
loss_gender = F.cross_entropy(input_gender, target[:, 0].long())
|
||||
loss_age = F.l1_loss(input_age, target[:, 2])
|
||||
|
||||
return loss_gender / (.16) + loss_age * 2
|
||||
|
||||
|
||||
model = ShuffleneFull().cuda()
|
||||
optimizer = optim.Adam(params=model.parameters(), lr=0.001, weight_decay=0.001)
|
||||
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
|
||||
history = TrainModel(model, train_dl, valid_dl, optimizer, multitask_loss, scheduler, 5)
|
||||
print(history)
|
||||
Binary file not shown.
@@ -0,0 +1,59 @@
|
||||
import torch
|
||||
import os
|
||||
import numpy as np
|
||||
from ..utils import download_weight
|
||||
from .models.densenet import densenet121
|
||||
from .models.resnet import resnet34
|
||||
|
||||
labels = np.array(['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral'])
|
||||
|
||||
class EmotionDetector:
|
||||
|
||||
def __init__(self, name='densnet121', weight_path=None, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
"""
|
||||
Residual Masking Emotion Detector from a list of labels
|
||||
:param name: name of backbone of networks (resnet34, densenet121)
|
||||
:param device: model run in cpu or gpu (cuda, cpu)
|
||||
:param weight_path: path of network weight
|
||||
|
||||
Notice: image size must be 224x224
|
||||
|
||||
Method detect_emotion:
|
||||
:param faces: 4d tensor of face for example size(1, 3, 224, 224)
|
||||
:returns emotions list and probability of emotions
|
||||
"""
|
||||
|
||||
self.device = device
|
||||
self.model = None
|
||||
if name == 'resnet34':
|
||||
self.model = resnet34()
|
||||
elif name == 'densnet121':
|
||||
self.model = densenet121()
|
||||
else:
|
||||
exit('from EmotionDetector: Network does not support!! \n just(resnet34, densnet121)')
|
||||
|
||||
# download the default weigth
|
||||
if weight_path is None:
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'densnet121.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/densnet121.pth')
|
||||
if os.path.isfile(weight_path) == False:
|
||||
print('from EmotionDetector: download defualt weight started')
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://drive.google.com/uc?export=download&id=1G3VsfgiQb16VyFnOwEVDgm2g8-9qN0-9', file_name=file_name)
|
||||
os.rename(file_name, weight_path)
|
||||
|
||||
self.model.load_state_dict(torch.load(weight_path, map_location=device))
|
||||
self.model.to(device).eval()
|
||||
print('from EmotionDetector: weights loaded')
|
||||
|
||||
|
||||
def detect_emotion(self, faces):
|
||||
if len(faces) > 0:
|
||||
faces = faces.permute(0, 3, 1, 2)
|
||||
faces = faces.float().div(255).to(self.device)
|
||||
emotions = self.model(faces)
|
||||
prob = torch.softmax(emotions, dim=1)
|
||||
emo_prob, emo_idx = torch.max(prob, dim=1)
|
||||
return labels[emo_idx.tolist()], emo_prob.tolist()
|
||||
else:
|
||||
return 0, 0
|
||||
@@ -0,0 +1,38 @@
|
||||
# from facelib import FaceDetector
|
||||
# from facelib import EmotionDetector
|
||||
# from facelib import special_draw
|
||||
|
||||
from ..Retinaface.Retinaface import FaceDetector
|
||||
from .FaceExpression import EmotionDetector
|
||||
from ..InsightFace.models.utils import special_draw
|
||||
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
class WebcamEmotionDetector:
|
||||
def __init__(self, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
print('loading ...')
|
||||
self.face_detector = FaceDetector(face_size=(224, 224), device=device)
|
||||
self.emotion_detector = EmotionDetector(device=device)
|
||||
|
||||
|
||||
def run(self, camera_index=0):
|
||||
cap = cv2.VideoCapture(camera_index)
|
||||
cap.set(3, 1280)
|
||||
cap.set(4, 720)
|
||||
print('type q for exit')
|
||||
while cap.isOpened():
|
||||
ret , frame = cap.read()
|
||||
if ret == False:
|
||||
raise Exception('the camera not recognized: change camera_index param to ' + str(0 if camera_index == 1 else 1))
|
||||
faces, boxes, scores, landmarks = self.face_detector.detect_align(frame)
|
||||
if len(faces.shape) > 1:
|
||||
emotions, emo_probs = self.emotion_detector.detect_emotion(faces)
|
||||
for i, b in enumerate(boxes):
|
||||
special_draw(frame, b, landmarks[i], name=emotions[i], score=emo_probs[i])
|
||||
|
||||
cv2.imshow('frame', frame)
|
||||
if cv2.waitKey(1) == ord('q'):
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,246 @@
|
||||
import re
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint as cp
|
||||
from collections import OrderedDict
|
||||
try:
|
||||
from torch.hub import load_state_dict_from_url
|
||||
except ImportError:
|
||||
from torch.utils.model_zoo import load_url as load_state_dict_from_url
|
||||
|
||||
|
||||
__all__ = ['DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161']
|
||||
|
||||
model_urls = {
|
||||
'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth',
|
||||
'densenet169': 'https://download.pytorch.org/models/densenet169-b2777c0a.pth',
|
||||
'densenet201': 'https://download.pytorch.org/models/densenet201-c1103571.pth',
|
||||
'densenet161': 'https://download.pytorch.org/models/densenet161-8d451a50.pth',
|
||||
}
|
||||
|
||||
|
||||
def _bn_function_factory(norm, relu, conv):
|
||||
def bn_function(*inputs):
|
||||
concated_features = torch.cat(inputs, 1)
|
||||
bottleneck_output = conv(relu(norm(concated_features)))
|
||||
return bottleneck_output
|
||||
|
||||
return bn_function
|
||||
|
||||
|
||||
class _DenseLayer(nn.Sequential):
|
||||
def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, memory_efficient=False):
|
||||
super(_DenseLayer, self).__init__()
|
||||
self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
|
||||
self.add_module('relu1', nn.ReLU(inplace=True)),
|
||||
self.add_module('conv1', nn.Conv2d(num_input_features, bn_size *
|
||||
growth_rate, kernel_size=1, stride=1,
|
||||
bias=False)),
|
||||
self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
|
||||
self.add_module('relu2', nn.ReLU(inplace=True)),
|
||||
self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,
|
||||
kernel_size=3, stride=1, padding=1,
|
||||
bias=False)),
|
||||
self.drop_rate = drop_rate
|
||||
self.memory_efficient = memory_efficient
|
||||
|
||||
def forward(self, *prev_features):
|
||||
bn_function = _bn_function_factory(self.norm1, self.relu1, self.conv1)
|
||||
if self.memory_efficient and any(prev_feature.requires_grad for prev_feature in prev_features):
|
||||
bottleneck_output = cp.checkpoint(bn_function, *prev_features)
|
||||
else:
|
||||
bottleneck_output = bn_function(*prev_features)
|
||||
new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))
|
||||
if self.drop_rate > 0:
|
||||
new_features = F.dropout(new_features, p=self.drop_rate,
|
||||
training=self.training)
|
||||
return new_features
|
||||
|
||||
|
||||
class _DenseBlock(nn.Module):
|
||||
def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False):
|
||||
super(_DenseBlock, self).__init__()
|
||||
for i in range(num_layers):
|
||||
layer = _DenseLayer(
|
||||
num_input_features + i * growth_rate,
|
||||
growth_rate=growth_rate,
|
||||
bn_size=bn_size,
|
||||
drop_rate=drop_rate,
|
||||
memory_efficient=memory_efficient,
|
||||
)
|
||||
self.add_module('denselayer%d' % (i + 1), layer)
|
||||
|
||||
def forward(self, init_features):
|
||||
features = [init_features]
|
||||
for name, layer in self.named_children():
|
||||
new_features = layer(*features)
|
||||
features.append(new_features)
|
||||
return torch.cat(features, 1)
|
||||
|
||||
|
||||
class _Transition(nn.Sequential):
|
||||
def __init__(self, num_input_features, num_output_features):
|
||||
super(_Transition, self).__init__()
|
||||
self.add_module('norm', nn.BatchNorm2d(num_input_features))
|
||||
self.add_module('relu', nn.ReLU(inplace=True))
|
||||
self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,
|
||||
kernel_size=1, stride=1, bias=False))
|
||||
self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
|
||||
|
||||
|
||||
class DenseNet(nn.Module):
|
||||
r"""Densenet-BC model class, based on
|
||||
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
|
||||
|
||||
Args:
|
||||
growth_rate (int) - how many filters to add each layer (`k` in paper)
|
||||
block_config (list of 4 ints) - how many layers in each pooling block
|
||||
num_init_features (int) - the number of filters to learn in the first convolution layer
|
||||
bn_size (int) - multiplicative factor for number of bottle neck layers
|
||||
(i.e. bn_size * k features in the bottleneck layer)
|
||||
drop_rate (float) - dropout rate after each dense layer
|
||||
num_classes (int) - number of classification classes
|
||||
memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
|
||||
but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
|
||||
"""
|
||||
|
||||
def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), #in_channels=3, num_classes=7,
|
||||
num_init_features=64, bn_size=4, drop_rate=0, memory_efficient=False):
|
||||
|
||||
super(DenseNet, self).__init__()
|
||||
|
||||
# First convolution
|
||||
self.features = nn.Sequential(OrderedDict([
|
||||
('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2,
|
||||
padding=3, bias=False)),
|
||||
('norm0', nn.BatchNorm2d(num_init_features)),
|
||||
('relu0', nn.ReLU(inplace=True)),
|
||||
('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
|
||||
]))
|
||||
|
||||
# Each denseblock
|
||||
num_features = num_init_features
|
||||
for i, num_layers in enumerate(block_config):
|
||||
block = _DenseBlock(
|
||||
num_layers=num_layers,
|
||||
num_input_features=num_features,
|
||||
bn_size=bn_size,
|
||||
growth_rate=growth_rate,
|
||||
drop_rate=drop_rate,
|
||||
memory_efficient=memory_efficient
|
||||
)
|
||||
self.features.add_module('denseblock%d' % (i + 1), block)
|
||||
num_features = num_features + num_layers * growth_rate
|
||||
if i != len(block_config) - 1:
|
||||
trans = _Transition(num_input_features=num_features,
|
||||
num_output_features=num_features // 2)
|
||||
self.features.add_module('transition%d' % (i + 1), trans)
|
||||
num_features = num_features // 2
|
||||
|
||||
# Final batch norm
|
||||
self.features.add_module('norm5', nn.BatchNorm2d(num_features))
|
||||
|
||||
# Linear layer
|
||||
# NOTE: strictly set to 1000 to load pretrained model
|
||||
# self.classifier = nn.Linear(num_features, num_classes)
|
||||
self.classifier = nn.Linear(num_features, 1000)
|
||||
|
||||
# Official init from torch repo.
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
features = self.features(x)
|
||||
out = F.relu(features, inplace=True)
|
||||
out = F.adaptive_avg_pool2d(out, (1, 1))
|
||||
out = torch.flatten(out, 1)
|
||||
out = self.classifier(out)
|
||||
return out
|
||||
|
||||
|
||||
def _load_state_dict(model, model_url, progress):
|
||||
# '.'s are no longer allowed in module names, but previous _DenseLayer
|
||||
# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
|
||||
# They are also in the checkpoints in model_urls. This pattern is used
|
||||
# to find such keys.
|
||||
pattern = re.compile(
|
||||
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
|
||||
|
||||
state_dict = load_state_dict_from_url(model_url, progress=progress)
|
||||
for key in list(state_dict.keys()):
|
||||
res = pattern.match(key)
|
||||
if res:
|
||||
new_key = res.group(1) + res.group(2)
|
||||
state_dict[new_key] = state_dict[key]
|
||||
del state_dict[key]
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
|
||||
def _densenet(arch, growth_rate, block_config, num_init_features, pretrained, progress, **kwargs):
|
||||
model = DenseNet(growth_rate, block_config, num_init_features, **kwargs)
|
||||
if pretrained:
|
||||
_load_state_dict(model, model_urls[arch], progress)
|
||||
|
||||
model.classifier = nn.Linear(1024, 7)
|
||||
return model
|
||||
|
||||
|
||||
def densenet121(pretrained=False, progress=True, **kwargs):
|
||||
r"""Densenet-121 model from `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
|
||||
but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
|
||||
"""
|
||||
return _densenet('densenet121', 32, (6, 12, 24, 16), 64, pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def densenet161(pretrained=False, progress=True, **kwargs):
|
||||
r"""Densenet-161 model from
|
||||
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
|
||||
but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
|
||||
"""
|
||||
return _densenet('densenet161', 48, (6, 12, 36, 24), 96, pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def densenet169(pretrained=False, progress=True, **kwargs):
|
||||
r"""Densenet-169 model from
|
||||
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
|
||||
but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
|
||||
"""
|
||||
return _densenet('densenet169', 32, (6, 12, 32, 32), 64, pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def densenet201(pretrained=False, progress=True, **kwargs):
|
||||
r"""Densenet-201 model from
|
||||
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
|
||||
but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_
|
||||
"""
|
||||
return _densenet('densenet201', 32, (6, 12, 48, 32), 64, pretrained, progress,
|
||||
**kwargs)
|
||||
@@ -0,0 +1,362 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
try:
|
||||
from torch.hub import load_state_dict_from_url
|
||||
except ImportError:
|
||||
from torch.utils.model_zoo import load_url as load_state_dict_from_url
|
||||
|
||||
|
||||
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||
'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
|
||||
'wide_resnet50_2', 'wide_resnet101_2']
|
||||
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
|
||||
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
|
||||
'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
|
||||
'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
|
||||
}
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
__constants__ = ['downsample']
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
||||
base_width=64, dilation=1, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
if groups != 1 or base_width != 64:
|
||||
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
|
||||
if dilation > 1:
|
||||
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
__constants__ = ['downsample']
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
||||
base_width=64, dilation=1, norm_layer=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
width = int(planes * (base_width / 64.)) * groups
|
||||
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv1x1(inplanes, width)
|
||||
self.bn1 = norm_layer(width)
|
||||
self.conv2 = conv3x3(width, width, stride, groups, dilation)
|
||||
self.bn2 = norm_layer(width)
|
||||
self.conv3 = conv1x1(width, planes * self.expansion)
|
||||
self.bn3 = norm_layer(planes * self.expansion)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, in_channels=3, num_classes=7, zero_init_residual=False,
|
||||
groups=1, width_per_group=64, replace_stride_with_dilation=None,
|
||||
norm_layer=None):
|
||||
super(ResNet, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.dilation = 1
|
||||
if replace_stride_with_dilation is None:
|
||||
# each element in the tuple indicates if we should replace
|
||||
# the 2x2 stride with a dilated convolution instead
|
||||
replace_stride_with_dilation = [False, False, False]
|
||||
if len(replace_stride_with_dilation) != 3:
|
||||
raise ValueError("replace_stride_with_dilation should be None "
|
||||
"or a 3-element tuple, got {}".format(replace_stride_with_dilation))
|
||||
self.groups = groups
|
||||
self.base_width = width_per_group
|
||||
|
||||
# NOTE: strictly set the in_channels = 3 to load the pretrained model
|
||||
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.bn1 = norm_layer(self.inplanes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
|
||||
dilate=replace_stride_with_dilation[0])
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
|
||||
dilate=replace_stride_with_dilation[1])
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
|
||||
dilate=replace_stride_with_dilation[2])
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
|
||||
# NOTE: strictly set the num_classes = 1000 to load the pretrained model
|
||||
self.fc = nn.Linear(512 * block.expansion, 1000)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck):
|
||||
nn.init.constant_(m.bn3.weight, 0)
|
||||
elif isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
previous_dilation = self.dilation
|
||||
if dilate:
|
||||
self.dilation *= stride
|
||||
stride = 1
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
|
||||
self.base_width, previous_dilation, norm_layer))
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, groups=self.groups,
|
||||
base_width=self.base_width, dilation=self.dilation,
|
||||
norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
|
||||
model = ResNet(block, layers, **kwargs)
|
||||
if pretrained:
|
||||
state_dict = load_state_dict_from_url(model_urls[arch],progress=progress)
|
||||
model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet18(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-18 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
model = _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
# model.fc = nn.Linear(512, kwargs['num_classes'])
|
||||
model.fc = nn.Linear(512, 7)
|
||||
return model
|
||||
|
||||
|
||||
def resnet34(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-34 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
model = _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,**kwargs)
|
||||
model.fc = nn.Linear(512, 7)
|
||||
return model
|
||||
|
||||
|
||||
def resnet50(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-50 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
model = _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
model.fc = nn.Linear(2048, kwargs['num_classes'])
|
||||
return model
|
||||
|
||||
|
||||
def resnet101(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-101 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
model = _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
model.fc = nn.Linear(2048, kwargs['num_classes'])
|
||||
return model
|
||||
|
||||
|
||||
def resnet152(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-152 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
model = _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
model.fc = nn.Linear(2048, kwargs['num_classes'])
|
||||
return model
|
||||
|
||||
|
||||
def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNeXt-50 32x4d model from
|
||||
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['groups'] = 32
|
||||
kwargs['width_per_group'] = 4
|
||||
return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNeXt-101 32x8d model from
|
||||
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['groups'] = 32
|
||||
kwargs['width_per_group'] = 8
|
||||
return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
|
||||
r"""Wide ResNet-50-2 model from
|
||||
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
|
||||
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
|
||||
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['width_per_group'] = 64 * 2
|
||||
return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
|
||||
r"""Wide ResNet-101-2 model from
|
||||
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
|
||||
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
|
||||
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['width_per_group'] = 64 * 2
|
||||
return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
from ..Retinaface.Retinaface import FaceDetector
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def add_from_webcam(person_name='unknow', camera_index=0):
|
||||
print('loading ...')
|
||||
# create facebank folder if is not exists
|
||||
|
||||
save_path = Path(os.path.dirname(os.path.realpath(__file__)), 'models/data/facebank')
|
||||
if not save_path.exists():
|
||||
save_path.mkdir()
|
||||
|
||||
# create a new folder with (name) for new person
|
||||
save_path = save_path/person_name
|
||||
if not save_path.exists():
|
||||
save_path.mkdir()
|
||||
|
||||
print('for exit: use q keyword')
|
||||
# init camera
|
||||
cap = cv2.VideoCapture(camera_index)
|
||||
cap.set(3, 1280)
|
||||
cap.set(4, 720)
|
||||
# init detector
|
||||
detector = FaceDetector()
|
||||
count = 4
|
||||
print('type q for exit')
|
||||
while cap.isOpened():
|
||||
ret , frame = cap.read()
|
||||
if ret == False:
|
||||
raise Exception('the camera not recognized: change camera_index param to ' + str(0 if camera_index == 1 else 1))
|
||||
|
||||
frame = cv2.putText(
|
||||
frame, f'Press t to take {count} pictures', (10, 50),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 2, (0,100,0), 3, cv2.LINE_AA)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord('t'):
|
||||
count -= 1
|
||||
faces = detector.detect_align(frame)[0].cpu().numpy()
|
||||
if len(faces.shape) > 1:
|
||||
cv2.imwrite(f'{save_path}/{person_name}_{count}.jpg', faces[0])
|
||||
if count <= 0:
|
||||
break
|
||||
else:
|
||||
print('there is not face in this frame')
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cv2.imshow("add new person from camera", frame)
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
print('images saved in this path: ', save_path)
|
||||
|
||||
|
||||
|
||||
def add_from_folder(folder_path='./', person_name='unknow'):
|
||||
|
||||
print('only a face in each image and all image from the same person')
|
||||
dir_path = Path(folder_path)
|
||||
if not dir_path.is_dir():
|
||||
exit('dir does not exists !!')
|
||||
|
||||
# create facebank folder if is not exists
|
||||
save_path = Path(os.path.dirname(os.path.realpath(__file__)), 'models/data/facebank')
|
||||
if not save_path.exists():
|
||||
save_path.mkdir()
|
||||
|
||||
|
||||
save_path = Path(f'{save_path}/{person_name}')
|
||||
if not save_path.exists():
|
||||
save_path.mkdir()
|
||||
print('loading ...')
|
||||
# init detector
|
||||
detector = FaceDetector()
|
||||
|
||||
counter = 0
|
||||
for img_path in dir_path.iterdir():
|
||||
img = cv2.imread(str(img_path))
|
||||
if img is None:
|
||||
raise Exception('this image has a problem: ', img_path)
|
||||
face = detector.detect_align(img)[0].cpu().numpy()
|
||||
if len(face.shape) > 1:
|
||||
save_name = f'{save_path}/{dir_path.name}_{counter}.jpg'
|
||||
cv2.imwrite(save_name, face[0])
|
||||
counter += 1
|
||||
else:
|
||||
print(img_path, 'in this image did not detect any face')
|
||||
print('images saved in this path: ', save_path)
|
||||
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
from .model import Backbone, Arcface, MobileFaceNet, l2_norm
|
||||
from .evaluatation import evaluate
|
||||
import torch
|
||||
from torch import optim
|
||||
import numpy as np
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
from ...utils import download_weight
|
||||
from .utils import get_time, gen_plot, separate_bn_paras
|
||||
from .utils import faces_preprocessing
|
||||
from PIL import Image
|
||||
from torchvision import transforms as trans
|
||||
import math
|
||||
from matplotlib import pyplot as plt
|
||||
plt.switch_backend('agg')
|
||||
|
||||
|
||||
class FaceRecognizer:
|
||||
|
||||
def __init__(self, conf, verbose=True):
|
||||
self.device = conf.device
|
||||
self.threshold = conf.threshold
|
||||
|
||||
try:
|
||||
if conf.use_mobilfacenet:
|
||||
self.model = MobileFaceNet(conf.embedding_size).to(conf.device)
|
||||
# download the default weigth
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'mobilenet.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/mobilenet.pth')
|
||||
if os.path.isfile(weight_path) == False:
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://drive.google.com/uc?export=download&id=1W9nM7LE6zUKQ4tncL6OnBn-aXNyiRPNH',
|
||||
file_name=file_name,
|
||||
verbose=verbose
|
||||
)
|
||||
os.rename(file_name, weight_path)
|
||||
|
||||
self.model.load_state_dict(torch.load(weight_path, map_location=conf.device))
|
||||
else:
|
||||
self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device)
|
||||
# download the default weigth
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'ir_se50.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/ir_se50.pth')
|
||||
if os.path.isfile(weight_path) == False:
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://www.dropbox.com/s/a570ucg6a5z22zf/ir_se50.pth?dl=1',
|
||||
file_name=file_name,
|
||||
verbose=verbose
|
||||
)
|
||||
os.rename(file_name, weight_path)
|
||||
|
||||
self.model.load_state_dict(torch.load(weight_path, map_location=conf.device))
|
||||
except IOError as e:
|
||||
exit(f'from FaceRecognizer Exit: the weight does not exist,'
|
||||
f' \n download and putting up in "{conf.work_path}" folder \n {e}')
|
||||
|
||||
self.model.eval()
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
|
||||
def evaluate(self, conf, carray, issame, nrof_folds=5, tta=False):
|
||||
self.model.eval()
|
||||
idx = 0
|
||||
embeddings = np.zeros([len(carray), conf.embedding_size])
|
||||
with torch.no_grad():
|
||||
while idx + conf.batch_size <= len(carray):
|
||||
batch = torch.tensor(carray[idx:idx + conf.batch_size])
|
||||
if tta:
|
||||
fliped = batch.flip(-1) # I do not test in this case
|
||||
emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device))
|
||||
embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch)
|
||||
else:
|
||||
embeddings[idx:idx + conf.batch_size] = self.model(batch.to(conf.device)).cpu()
|
||||
idx += conf.batch_size
|
||||
if idx < len(carray):
|
||||
batch = torch.tensor(carray[idx:])
|
||||
if tta:
|
||||
fliped = batch.flip(-1)
|
||||
emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device))
|
||||
embeddings[idx:] = l2_norm(emb_batch)
|
||||
else:
|
||||
embeddings[idx:] = self.model(batch.to(conf.device)).cpu()
|
||||
tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds)
|
||||
buf = gen_plot(fpr, tpr)
|
||||
roc_curve = Image.open(buf)
|
||||
roc_curve_tensor = trans.ToTensor()(roc_curve)
|
||||
return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor
|
||||
|
||||
|
||||
def train(self, conf, epochs):
|
||||
self.model.train()
|
||||
running_loss = 0.
|
||||
for e in range(epochs):
|
||||
print('epoch {} started'.format(e))
|
||||
if e == self.milestones[0]:
|
||||
self.schedule_lr()
|
||||
if e == self.milestones[1]:
|
||||
self.schedule_lr()
|
||||
if e == self.milestones[2]:
|
||||
self.schedule_lr()
|
||||
for imgs, labels in tqdm(iter(self.loader)):
|
||||
imgs = imgs.to(conf.device)
|
||||
labels = labels.to(conf.device)
|
||||
self.optimizer.zero_grad()
|
||||
embeddings = self.model(imgs)
|
||||
thetas = self.head(embeddings, labels)
|
||||
loss = conf.ce_loss(thetas, labels)
|
||||
loss.backward()
|
||||
running_loss += loss.item()
|
||||
self.optimizer.step()
|
||||
|
||||
if self.step % self.board_loss_every == 0 and self.step != 0:
|
||||
loss_board = running_loss / self.board_loss_every
|
||||
running_loss = 0.
|
||||
|
||||
if self.step % self.evaluate_every == 0 and self.step != 0:
|
||||
accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.agedb_30,
|
||||
self.agedb_30_issame)
|
||||
accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.lfw, self.lfw_issame)
|
||||
accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.cfp_fp, self.cfp_fp_issame)
|
||||
self.model.train()
|
||||
if self.step % self.save_every == 0 and self.step != 0:
|
||||
self.save_state(conf, accuracy)
|
||||
|
||||
self.step += 1
|
||||
|
||||
self.save_state(conf, accuracy, to_save_folder=True, extra='final')
|
||||
|
||||
def schedule_lr(self):
|
||||
for params in self.optimizer.param_groups:
|
||||
params['lr'] /= 10
|
||||
print(self.optimizer)
|
||||
|
||||
def infer(self, faces, target_embs, tta=False):
|
||||
"""
|
||||
faces : list of PIL Image
|
||||
target_embs : [n, 512] computed embeddings of faces in facebank
|
||||
names : recorded names of faces in facebank
|
||||
tta : test time augmentation (hfilp, that's all)
|
||||
"""
|
||||
faces = faces_preprocessing(faces, self.device)
|
||||
if tta:
|
||||
faces_emb = self.model(faces)
|
||||
hflip_emb = self.model(faces.flip(-1)) # image horizontal flip
|
||||
embs = l2_norm((faces_emb + hflip_emb)/2) # take mean
|
||||
else:
|
||||
embs = self.model(faces)
|
||||
|
||||
diff = embs.unsqueeze(-1) - target_embs.transpose(1, 0).unsqueeze(0)
|
||||
dist = torch.sum(torch.pow(diff, 2), dim=1)
|
||||
minimum, min_idx = torch.min(dist, dim=1)
|
||||
min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1
|
||||
return min_idx, minimum
|
||||
|
||||
def feature_extractor(self, faces, tta=False):
|
||||
"""
|
||||
faces : stack of torch.tensors
|
||||
tta : test time augmentation (hfilp, that's all)
|
||||
"""
|
||||
faces = faces_preprocessing(faces, self.device)
|
||||
if tta:
|
||||
faces_emb = self.model(faces)
|
||||
hflip_emb = self.model(faces.flip(-1)) # horizontal flip
|
||||
embs = l2_norm((faces_emb + hflip_emb)/2) # take mean
|
||||
else:
|
||||
embs = self.model(faces)
|
||||
return embs
|
||||
@@ -0,0 +1,46 @@
|
||||
from easydict import EasyDict as edict
|
||||
from pathlib import Path
|
||||
import os
|
||||
import torch
|
||||
from torch.nn import CrossEntropyLoss
|
||||
|
||||
|
||||
def get_config(inference=True):
|
||||
conf = edict()
|
||||
conf.data_path = Path(os.path.dirname(os.path.realpath(__file__)))
|
||||
conf.work_path = Path('weights/')
|
||||
conf.model_path = conf.work_path / 'models'
|
||||
conf.log_path = conf.work_path / 'log'
|
||||
conf.save_path = conf.work_path
|
||||
conf.input_size = [112, 112]
|
||||
conf.embedding_size = 512
|
||||
conf.use_mobilfacenet = True
|
||||
conf.net_depth = 50
|
||||
conf.drop_ratio = 0.6
|
||||
conf.net_mode = 'ir_se' # or 'ir'
|
||||
conf.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
conf.data_mode = 'emore'
|
||||
conf.vgg_folder = conf.data_path / 'faces_vgg_112x112'
|
||||
conf.ms1m_folder = conf.data_path / 'faces_ms1m_112x112'
|
||||
conf.emore_folder = conf.data_path / 'faces_emore'
|
||||
conf.batch_size = 100 # irse net depth 50
|
||||
# conf.batch_size = 200 # mobilefacenet
|
||||
if inference:
|
||||
# --------------------Inference Config ------------------------
|
||||
conf.facebank_path = conf.data_path / 'facebank'
|
||||
conf.threshold = 1.5
|
||||
conf.face_limit = 10
|
||||
# when inference, at maximum detect 10 faces in one image, my laptop is slow
|
||||
else:
|
||||
# --------------------Training Config ------------------------
|
||||
conf.log_path = conf.work_path / 'log'
|
||||
conf.save_path = conf.work_path / 'save'
|
||||
# conf.weight_decay = 5e-4
|
||||
conf.lr = 1e-3
|
||||
conf.momentum = 0.9
|
||||
conf.pin_memory = True
|
||||
# conf.num_workers = 4 # when batchsize is 200
|
||||
conf.num_workers = 3
|
||||
conf.ce_loss = CrossEntropyLoss()
|
||||
|
||||
return conf
|
||||
@@ -0,0 +1,168 @@
|
||||
""""Helper for evaluation on the Labeled Faces in the Wild dataset"""
|
||||
|
||||
# MIT License
|
||||
#
|
||||
# Copyright (c) 2016 David Sandberg
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to the following conditions:
|
||||
#
|
||||
# The above copyright notice and this permission notice shall be included in all
|
||||
# copies or substantial portions of the Software.
|
||||
#
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
# SOFTWARE.
|
||||
|
||||
import numpy as np
|
||||
from sklearn.model_selection import KFold
|
||||
from sklearn.decomposition import PCA
|
||||
import sklearn
|
||||
from scipy import interpolate
|
||||
|
||||
|
||||
def calculate_roc(thresholds, embeddings1, embeddings2, actual_issame, nrof_folds=10, pca=0):
|
||||
assert (embeddings1.shape[0] == embeddings2.shape[0])
|
||||
assert (embeddings1.shape[1] == embeddings2.shape[1])
|
||||
nrof_pairs = min(len(actual_issame), embeddings1.shape[0])
|
||||
nrof_thresholds = len(thresholds)
|
||||
k_fold = KFold(n_splits=nrof_folds, shuffle=False)
|
||||
|
||||
tprs = np.zeros((nrof_folds, nrof_thresholds))
|
||||
fprs = np.zeros((nrof_folds, nrof_thresholds))
|
||||
accuracy = np.zeros((nrof_folds))
|
||||
best_thresholds = np.zeros((nrof_folds))
|
||||
indices = np.arange(nrof_pairs)
|
||||
# print('pca', pca)
|
||||
|
||||
if pca == 0:
|
||||
diff = np.subtract(embeddings1, embeddings2)
|
||||
dist = np.sum(np.square(diff), 1)
|
||||
|
||||
for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):
|
||||
# print('train_set', train_set)
|
||||
# print('test_set', test_set)
|
||||
if pca > 0:
|
||||
print('doing pca on', fold_idx)
|
||||
embed1_train = embeddings1[train_set]
|
||||
embed2_train = embeddings2[train_set]
|
||||
_embed_train = np.concatenate((embed1_train, embed2_train), axis=0)
|
||||
# print(_embed_train.shape)
|
||||
pca_model = PCA(n_components=pca)
|
||||
pca_model.fit(_embed_train)
|
||||
embed1 = pca_model.transform(embeddings1)
|
||||
embed2 = pca_model.transform(embeddings2)
|
||||
embed1 = sklearn.preprocessing.normalize(embed1)
|
||||
embed2 = sklearn.preprocessing.normalize(embed2)
|
||||
# print(embed1.shape, embed2.shape)
|
||||
diff = np.subtract(embed1, embed2)
|
||||
dist = np.sum(np.square(diff), 1)
|
||||
|
||||
# Find the best threshold for the fold
|
||||
acc_train = np.zeros((nrof_thresholds))
|
||||
for threshold_idx, threshold in enumerate(thresholds):
|
||||
_, _, acc_train[threshold_idx] = calculate_accuracy(threshold, dist[train_set], actual_issame[train_set])
|
||||
best_threshold_index = np.argmax(acc_train)
|
||||
# print('best_threshold_index', best_threshold_index, acc_train[best_threshold_index])
|
||||
best_thresholds[fold_idx] = thresholds[best_threshold_index]
|
||||
for threshold_idx, threshold in enumerate(thresholds):
|
||||
tprs[fold_idx, threshold_idx], fprs[fold_idx, threshold_idx], _ = calculate_accuracy(threshold,
|
||||
dist[test_set],
|
||||
actual_issame[
|
||||
test_set])
|
||||
_, _, accuracy[fold_idx] = calculate_accuracy(thresholds[best_threshold_index], dist[test_set],
|
||||
actual_issame[test_set])
|
||||
|
||||
tpr = np.mean(tprs, 0)
|
||||
fpr = np.mean(fprs, 0)
|
||||
return tpr, fpr, accuracy, best_thresholds
|
||||
|
||||
|
||||
def calculate_accuracy(threshold, dist, actual_issame):
|
||||
predict_issame = np.less(dist, threshold)
|
||||
tp = np.sum(np.logical_and(predict_issame, actual_issame))
|
||||
fp = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))
|
||||
tn = np.sum(np.logical_and(np.logical_not(predict_issame), np.logical_not(actual_issame)))
|
||||
fn = np.sum(np.logical_and(np.logical_not(predict_issame), actual_issame))
|
||||
|
||||
tpr = 0 if (tp + fn == 0) else float(tp) / float(tp + fn)
|
||||
fpr = 0 if (fp + tn == 0) else float(fp) / float(fp + tn)
|
||||
acc = float(tp + tn) / dist.size
|
||||
return tpr, fpr, acc
|
||||
|
||||
|
||||
def calculate_val(thresholds, embeddings1, embeddings2, actual_issame, far_target, nrof_folds=10):
|
||||
'''
|
||||
Copy from [insightface](https://github.com/deepinsight/insightface)
|
||||
:param thresholds:
|
||||
:param embeddings1:
|
||||
:param embeddings2:
|
||||
:param actual_issame:
|
||||
:param far_target:
|
||||
:param nrof_folds:
|
||||
:return:
|
||||
'''
|
||||
assert (embeddings1.shape[0] == embeddings2.shape[0])
|
||||
assert (embeddings1.shape[1] == embeddings2.shape[1])
|
||||
nrof_pairs = min(len(actual_issame), embeddings1.shape[0])
|
||||
nrof_thresholds = len(thresholds)
|
||||
k_fold = KFold(n_splits=nrof_folds, shuffle=False)
|
||||
|
||||
val = np.zeros(nrof_folds)
|
||||
far = np.zeros(nrof_folds)
|
||||
|
||||
diff = np.subtract(embeddings1, embeddings2)
|
||||
dist = np.sum(np.square(diff), 1)
|
||||
indices = np.arange(nrof_pairs)
|
||||
|
||||
for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):
|
||||
|
||||
# Find the threshold that gives FAR = far_target
|
||||
far_train = np.zeros(nrof_thresholds)
|
||||
for threshold_idx, threshold in enumerate(thresholds):
|
||||
_, far_train[threshold_idx] = calculate_val_far(threshold, dist[train_set], actual_issame[train_set])
|
||||
if np.max(far_train) >= far_target:
|
||||
f = interpolate.interp1d(far_train, thresholds, kind='slinear')
|
||||
threshold = f(far_target)
|
||||
else:
|
||||
threshold = 0.0
|
||||
|
||||
val[fold_idx], far[fold_idx] = calculate_val_far(threshold, dist[test_set], actual_issame[test_set])
|
||||
|
||||
val_mean = np.mean(val)
|
||||
far_mean = np.mean(far)
|
||||
val_std = np.std(val)
|
||||
return val_mean, val_std, far_mean
|
||||
|
||||
|
||||
def calculate_val_far(threshold, dist, actual_issame):
|
||||
predict_issame = np.less(dist, threshold)
|
||||
true_accept = np.sum(np.logical_and(predict_issame, actual_issame))
|
||||
false_accept = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))
|
||||
n_same = np.sum(actual_issame)
|
||||
n_diff = np.sum(np.logical_not(actual_issame))
|
||||
val = float(true_accept) / float(n_same)
|
||||
far = float(false_accept) / float(n_diff)
|
||||
return val, far
|
||||
|
||||
|
||||
def evaluate(embeddings, actual_issame, nrof_folds=10, pca=0):
|
||||
# Calculate evaluation metrics
|
||||
thresholds = np.arange(0, 4, 0.01)
|
||||
embeddings1 = embeddings[0::2]
|
||||
embeddings2 = embeddings[1::2]
|
||||
tpr, fpr, accuracy, best_thresholds = calculate_roc(thresholds, embeddings1, embeddings2,
|
||||
np.asarray(actual_issame), nrof_folds=nrof_folds, pca=pca)
|
||||
# thresholds = np.arange(0, 4, 0.001)
|
||||
# val, val_std, far = calculate_val(thresholds, embeddings1, embeddings2,
|
||||
# np.asarray(actual_issame), 1e-3, nrof_folds=nrof_folds)
|
||||
# return tpr, fpr, accuracy, best_thresholds, val, val_std, far
|
||||
return tpr, fpr, accuracy, best_thresholds
|
||||
@@ -0,0 +1,309 @@
|
||||
from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, \
|
||||
MaxPool2d, AdaptiveAvgPool2d, Sequential, Module, Parameter
|
||||
import torch
|
||||
from collections import namedtuple
|
||||
import math
|
||||
|
||||
|
||||
################################## Original Arcface Model ############################################################
|
||||
|
||||
class Flatten(Module):
|
||||
def forward(self, input):
|
||||
return input.view(input.size(0), -1)
|
||||
|
||||
|
||||
def l2_norm(input, axis=1):
|
||||
norm = torch.norm(input, 2, axis, True)
|
||||
output = torch.div(input, norm)
|
||||
return output
|
||||
|
||||
|
||||
class SEModule(Module):
|
||||
def __init__(self, channels, reduction):
|
||||
super(SEModule, self).__init__()
|
||||
self.avg_pool = AdaptiveAvgPool2d(1)
|
||||
self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False)
|
||||
self.relu = ReLU(inplace=True)
|
||||
self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False)
|
||||
self.sigmoid = Sigmoid()
|
||||
|
||||
def forward(self, x):
|
||||
module_input = x
|
||||
x = self.avg_pool(x)
|
||||
x = self.fc1(x)
|
||||
x = self.relu(x)
|
||||
x = self.fc2(x)
|
||||
x = self.sigmoid(x)
|
||||
return module_input * x
|
||||
|
||||
|
||||
class bottleneck_IR(Module):
|
||||
def __init__(self, in_channel, depth, stride):
|
||||
super(bottleneck_IR, self).__init__()
|
||||
if in_channel == depth:
|
||||
self.shortcut_layer = MaxPool2d(1, stride)
|
||||
else:
|
||||
self.shortcut_layer = Sequential(
|
||||
Conv2d(in_channel, depth, (1, 1), stride, bias=False), BatchNorm2d(depth))
|
||||
self.res_layer = Sequential(
|
||||
BatchNorm2d(in_channel),
|
||||
Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth),
|
||||
Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth))
|
||||
|
||||
def forward(self, x):
|
||||
shortcut = self.shortcut_layer(x)
|
||||
res = self.res_layer(x)
|
||||
return res + shortcut
|
||||
|
||||
|
||||
class bottleneck_IR_SE(Module):
|
||||
def __init__(self, in_channel, depth, stride):
|
||||
super(bottleneck_IR_SE, self).__init__()
|
||||
if in_channel == depth:
|
||||
self.shortcut_layer = MaxPool2d(1, stride)
|
||||
else:
|
||||
self.shortcut_layer = Sequential(Conv2d(in_channel, depth, (1, 1), stride, bias=False), BatchNorm2d(depth))
|
||||
self.res_layer = Sequential(
|
||||
BatchNorm2d(in_channel),
|
||||
Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
|
||||
PReLU(depth),
|
||||
Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
|
||||
BatchNorm2d(depth),
|
||||
SEModule(depth, 16)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
shortcut = self.shortcut_layer(x)
|
||||
res = self.res_layer(x)
|
||||
return res + shortcut
|
||||
|
||||
|
||||
class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
|
||||
'''A named tuple describing a ResNet block.'''
|
||||
|
||||
|
||||
def get_block(in_channel, depth, num_units, stride=2):
|
||||
return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
|
||||
|
||||
|
||||
def get_blocks(num_layers):
|
||||
if num_layers == 50:
|
||||
blocks = [
|
||||
get_block(in_channel=64, depth=64, num_units=3),
|
||||
get_block(in_channel=64, depth=128, num_units=4),
|
||||
get_block(in_channel=128, depth=256, num_units=14),
|
||||
get_block(in_channel=256, depth=512, num_units=3)
|
||||
]
|
||||
elif num_layers == 100:
|
||||
blocks = [
|
||||
get_block(in_channel=64, depth=64, num_units=3),
|
||||
get_block(in_channel=64, depth=128, num_units=13),
|
||||
get_block(in_channel=128, depth=256, num_units=30),
|
||||
get_block(in_channel=256, depth=512, num_units=3)
|
||||
]
|
||||
elif num_layers == 152:
|
||||
blocks = [
|
||||
get_block(in_channel=64, depth=64, num_units=3),
|
||||
get_block(in_channel=64, depth=128, num_units=8),
|
||||
get_block(in_channel=128, depth=256, num_units=36),
|
||||
get_block(in_channel=256, depth=512, num_units=3)
|
||||
]
|
||||
return blocks
|
||||
|
||||
|
||||
class Backbone(Module):
|
||||
def __init__(self, num_layers, drop_ratio, mode='ir'):
|
||||
super(Backbone, self).__init__()
|
||||
assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
|
||||
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
|
||||
blocks = get_blocks(num_layers)
|
||||
if mode == 'ir':
|
||||
unit_module = bottleneck_IR
|
||||
elif mode == 'ir_se':
|
||||
unit_module = bottleneck_IR_SE
|
||||
self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), BatchNorm2d(64), PReLU(64))
|
||||
self.output_layer = Sequential(BatchNorm2d(512),
|
||||
Dropout(drop_ratio),
|
||||
Flatten(),
|
||||
Linear(512 * 7 * 7, 512),
|
||||
BatchNorm1d(512))
|
||||
modules = []
|
||||
for block in blocks:
|
||||
for bottleneck in block:
|
||||
modules.append(unit_module(bottleneck.in_channel, bottleneck.depth, bottleneck.stride))
|
||||
self.body = Sequential(*modules)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.input_layer(x)
|
||||
x = self.body(x)
|
||||
x = self.output_layer(x)
|
||||
return l2_norm(x)
|
||||
|
||||
|
||||
################################## MobileFaceNet #############################################################
|
||||
|
||||
class Conv_block(Module):
|
||||
def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
|
||||
super(Conv_block, self).__init__()
|
||||
self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding,
|
||||
bias=False)
|
||||
self.bn = BatchNorm2d(out_c)
|
||||
self.prelu = PReLU(out_c)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
x = self.prelu(x)
|
||||
return x
|
||||
|
||||
|
||||
class Linear_block(Module):
|
||||
def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
|
||||
super(Linear_block, self).__init__()
|
||||
self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding,
|
||||
bias=False)
|
||||
self.bn = BatchNorm2d(out_c)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
return x
|
||||
|
||||
|
||||
class Depth_Wise(Module):
|
||||
def __init__(self, in_c, out_c, residual=False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1):
|
||||
super(Depth_Wise, self).__init__()
|
||||
self.conv = Conv_block(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
|
||||
self.conv_dw = Conv_block(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride)
|
||||
self.project = Linear_block(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
|
||||
self.residual = residual
|
||||
|
||||
def forward(self, x):
|
||||
if self.residual:
|
||||
short_cut = x
|
||||
x = self.conv(x)
|
||||
x = self.conv_dw(x)
|
||||
x = self.project(x)
|
||||
if self.residual:
|
||||
output = short_cut + x
|
||||
else:
|
||||
output = x
|
||||
return output
|
||||
|
||||
|
||||
class Residual(Module):
|
||||
def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):
|
||||
super(Residual, self).__init__()
|
||||
modules = []
|
||||
for _ in range(num_block):
|
||||
modules.append(
|
||||
Depth_Wise(c, c, residual=True, kernel=kernel, padding=padding, stride=stride, groups=groups))
|
||||
self.model = Sequential(*modules)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
|
||||
class MobileFaceNet(Module):
|
||||
def __init__(self, embedding_size):
|
||||
super(MobileFaceNet, self).__init__()
|
||||
self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1))
|
||||
self.conv2_dw = Conv_block(64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
|
||||
self.conv_23 = Depth_Wise(64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128)
|
||||
self.conv_3 = Residual(64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
self.conv_34 = Depth_Wise(64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256)
|
||||
self.conv_4 = Residual(128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
self.conv_45 = Depth_Wise(128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512)
|
||||
self.conv_5 = Residual(128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
self.conv_6_sep = Conv_block(128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0))
|
||||
self.conv_6_dw = Linear_block(512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0))
|
||||
self.conv_6_flatten = Flatten()
|
||||
self.linear = Linear(512, embedding_size, bias=False)
|
||||
self.bn = BatchNorm1d(embedding_size)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = self.conv2_dw(out)
|
||||
out = self.conv_23(out)
|
||||
out = self.conv_3(out)
|
||||
out = self.conv_34(out)
|
||||
out = self.conv_4(out)
|
||||
out = self.conv_45(out)
|
||||
out = self.conv_5(out)
|
||||
out = self.conv_6_sep(out)
|
||||
out = self.conv_6_dw(out)
|
||||
out = self.conv_6_flatten(out)
|
||||
out = self.linear(out)
|
||||
out = self.bn(out)
|
||||
return l2_norm(out)
|
||||
|
||||
|
||||
################################## Arcface head #############################################################
|
||||
|
||||
class Arcface(Module):
|
||||
# implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599
|
||||
def __init__(self, embedding_size=512, classnum=51332, s=64., m=0.5):
|
||||
super(Arcface, self).__init__()
|
||||
self.classnum = classnum
|
||||
self.kernel = Parameter(torch.Tensor(embedding_size, classnum))
|
||||
# initial kernel
|
||||
self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
|
||||
self.m = m # the margin value, default is 0.5
|
||||
self.s = s # scalar value default is 64, see normface https://arxiv.org/abs/1704.06369
|
||||
self.cos_m = math.cos(m)
|
||||
self.sin_m = math.sin(m)
|
||||
self.mm = self.sin_m * m # issue 1
|
||||
self.threshold = math.cos(math.pi - m)
|
||||
|
||||
def forward(self, embbedings, label):
|
||||
# weights norm
|
||||
nB = len(embbedings)
|
||||
kernel_norm = l2_norm(self.kernel, axis=0)
|
||||
# cos(theta+m)
|
||||
cos_theta = torch.mm(embbedings, kernel_norm)
|
||||
# output = torch.mm(embbedings,kernel_norm)
|
||||
cos_theta = cos_theta.clamp(-1, 1) # for numerical stability
|
||||
cos_theta_2 = torch.pow(cos_theta, 2)
|
||||
sin_theta_2 = 1 - cos_theta_2
|
||||
sin_theta = torch.sqrt(sin_theta_2)
|
||||
cos_theta_m = (cos_theta * self.cos_m - sin_theta * self.sin_m)
|
||||
# this condition controls the theta+m should in range [0, pi]
|
||||
# 0<=theta+m<=pi
|
||||
# -m<=theta<=pi-m
|
||||
cond_v = cos_theta - self.threshold
|
||||
cond_mask = cond_v <= 0
|
||||
keep_val = (cos_theta - self.mm) # when theta not in [0,pi], use cosface instead
|
||||
cos_theta_m[cond_mask] = keep_val[cond_mask]
|
||||
output = cos_theta * 1.0 # a little bit hacky way to prevent in_place operation on cos_theta
|
||||
idx_ = torch.arange(0, nB, dtype=torch.long)
|
||||
output[idx_, label] = cos_theta_m[idx_, label]
|
||||
output *= self.s # scale up in order to make softmax work, first introduced in normface
|
||||
return output
|
||||
|
||||
|
||||
################################## Cosface head #############################################################
|
||||
|
||||
class Am_softmax(Module):
|
||||
# implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599
|
||||
def __init__(self, embedding_size=512, classnum=51332):
|
||||
super(Am_softmax, self).__init__()
|
||||
self.classnum = classnum
|
||||
self.kernel = Parameter(torch.Tensor(embedding_size, classnum))
|
||||
# initial kernel
|
||||
self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
|
||||
self.m = 0.35 # additive margin recommended by the paper
|
||||
self.s = 30. # see normface https://arxiv.org/abs/1704.06369
|
||||
|
||||
def forward(self, embbedings, label):
|
||||
kernel_norm = l2_norm(self.kernel, axis=0)
|
||||
cos_theta = torch.mm(embbedings, kernel_norm)
|
||||
cos_theta = cos_theta.clamp(-1, 1) # for numerical stability
|
||||
phi = cos_theta - self.m
|
||||
label = label.view(-1, 1) # size=(B,1)
|
||||
index = cos_theta.data * 0.0 # size=(B,Classnum)
|
||||
index.scatter_(1, label.data.view(-1, 1), 1)
|
||||
index = index.byte()
|
||||
output = cos_theta * 1.0
|
||||
output[index] = phi[index] # only change the correct predicted output
|
||||
output *= self.s # scale up in order to make softmax work, first introduced in normface
|
||||
return output
|
||||
@@ -0,0 +1,31 @@
|
||||
from facelib import get_config
|
||||
from facelib import FaceRecognizer
|
||||
import argparse
|
||||
|
||||
# python train.py -net mobilefacenet -b 200 -w 4
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='for face verification')
|
||||
parser.add_argument("-e", "--epochs", help="training epochs", default=20, type=int)
|
||||
parser.add_argument("-net", "--net_mode", help="which network, [ir, ir_se, mobilefacenet]",default='ir_se', type=str)
|
||||
parser.add_argument("-depth", "--net_depth", help="how many layers [50,100,152]", default=50, type=int)
|
||||
parser.add_argument('-lr','--lr',help='learning rate',default=1e-3, type=float)
|
||||
parser.add_argument("-b", "--batch_size", help="batch_size", default=96, type=int)
|
||||
parser.add_argument("-w", "--num_workers", help="workers number", default=3, type=int)
|
||||
parser.add_argument("-d", "--data_mode", help="use which database, [vgg, ms1m, emore, concat]",default='emore', type=str)
|
||||
args = parser.parse_args()
|
||||
|
||||
conf = get_config()
|
||||
|
||||
if args.net_mode == 'mobilefacenet':
|
||||
conf.use_mobilfacenet = True
|
||||
else:
|
||||
conf.net_mode = args.net_mode
|
||||
conf.net_depth = args.net_depth
|
||||
|
||||
conf.lr = args.lr
|
||||
conf.batch_size = args.batch_size
|
||||
conf.num_workers = args.num_workers
|
||||
conf.data_mode = args.data_mode
|
||||
learner = FaceRecognizer(conf)
|
||||
learner.train(conf, args.epochs)
|
||||
@@ -0,0 +1,132 @@
|
||||
from datetime import datetime
|
||||
import numpy as np
|
||||
import io, cv2, os
|
||||
from .model import l2_norm
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
plt.switch_backend('agg')
|
||||
|
||||
|
||||
def faces_preprocessing(faces, device):
|
||||
faces = faces.permute(0, 3, 1, 2).float()
|
||||
faces = faces.div(255).to(device)
|
||||
mu = torch.as_tensor([.5, .5, .5], dtype=faces.dtype, device=device)
|
||||
faces[:].sub_(mu[:, None, None]).div_(mu[:, None, None])
|
||||
return faces
|
||||
|
||||
|
||||
def separate_bn_paras(modules):
|
||||
if not isinstance(modules, list):
|
||||
modules = [*modules.modules()]
|
||||
paras_only_bn = []
|
||||
paras_wo_bn = []
|
||||
for layer in modules:
|
||||
if 'model' in str(layer.__class__):
|
||||
continue
|
||||
if 'container' in str(layer.__class__):
|
||||
continue
|
||||
else:
|
||||
if 'batchnorm' in str(layer.__class__):
|
||||
paras_only_bn.extend([*layer.parameters()])
|
||||
else:
|
||||
paras_wo_bn.extend([*layer.parameters()])
|
||||
return paras_only_bn, paras_wo_bn
|
||||
|
||||
|
||||
def update_facebank(conf, model, detector, tta=True):
|
||||
if os.path.exists(conf.facebank_path) == False:
|
||||
raise Exception("you don't have facebank yet: create with add_from_webcam or add_from_folder function")
|
||||
model.eval()
|
||||
faces_embs = torch.empty(0).to(conf.device)
|
||||
names = np.array(['Unknown'])
|
||||
for path in conf.facebank_path.iterdir():
|
||||
if path.is_file():
|
||||
continue
|
||||
faces = []
|
||||
for img_path in path.iterdir():
|
||||
face = cv2.imread(str(img_path))
|
||||
if face.shape[:2] != (112, 112): # if img be not face
|
||||
face = detector.detect_align(face)[0]
|
||||
cv2.imwrite(img_path, face)
|
||||
else:
|
||||
face = torch.tensor(face).unsqueeze(0)
|
||||
faces.append(face)
|
||||
|
||||
faces = torch.cat(faces)
|
||||
if len(faces.shape) <= 3:
|
||||
continue
|
||||
|
||||
with torch.no_grad():
|
||||
faces = faces_preprocessing(faces, device=conf.device)
|
||||
if tta:
|
||||
face_emb = model(faces)
|
||||
hflip_emb = model(faces.flip(-1)) # image horizontal flip
|
||||
face_embs = l2_norm(face_emb + hflip_emb)
|
||||
else:
|
||||
face_embs = model(faces)
|
||||
|
||||
faces_embs = torch.cat((faces_embs, face_embs.mean(0, keepdim=True)))
|
||||
names = np.append(names, path.name)
|
||||
|
||||
torch.save(faces_embs, conf.facebank_path/'facebank.pth')
|
||||
np.save(conf.facebank_path/'names', names)
|
||||
print('from FaceRecognizer: facebank updated')
|
||||
return faces_embs, names
|
||||
|
||||
|
||||
def load_facebank(conf):
|
||||
if os.path.exists(conf.facebank_path) == False:
|
||||
raise Exception("you don't have facebank yet: create with add_from_webcam or add_from_folder function")
|
||||
|
||||
embs = torch.load(conf.facebank_path/'facebank.pth')
|
||||
names = np.load(conf.facebank_path/'names.npy')
|
||||
print('from FaceRecognizer: facebank loaded')
|
||||
return embs, names
|
||||
|
||||
|
||||
def get_time():
|
||||
return (str(datetime.now())[:-10]).replace(' ', '-').replace(':', '-')
|
||||
|
||||
|
||||
def gen_plot(fpr, tpr):
|
||||
"""Create a pyplot plot and save to buffer."""
|
||||
plt.figure()
|
||||
plt.xlabel("FPR", fontsize=14)
|
||||
plt.ylabel("TPR", fontsize=14)
|
||||
plt.title("ROC Curve", fontsize=14)
|
||||
plt.plot(fpr, tpr, linewidth=2)
|
||||
buf = io.BytesIO()
|
||||
plt.savefig(buf, format='jpeg')
|
||||
buf.seek(0)
|
||||
plt.close()
|
||||
return buf
|
||||
|
||||
|
||||
def draw(bbox, name, frame):
|
||||
frame = cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 6)
|
||||
frame = cv2.putText(frame, name, (bbox[0], bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 4, cv2.LINE_AA)
|
||||
return frame
|
||||
|
||||
|
||||
def special_draw(img, box, landmarsk, name, score=100):
|
||||
"""draw a bounding box on image"""
|
||||
color = (148, 133, 0)
|
||||
tl = round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line thickness
|
||||
c1, c2 = (box[0], box[1]), (box[2], box[3])
|
||||
# draw bounding box
|
||||
cv2.rectangle(img, c1, c2, color, thickness=tl)
|
||||
# draw landmark
|
||||
for land in landmarsk:
|
||||
cv2.circle(img, tuple(land.int().tolist()), 3, color, -1)
|
||||
# draw score
|
||||
score = 100-(score*100/1.4)
|
||||
score = 0 if score < 0 else score
|
||||
bar = (box[3] + 2) - (box[1] - 2)
|
||||
score_final = bar - (score*bar/100)
|
||||
cv2.rectangle(img, (box[2] + 1, box[1] - 2 + score_final), (box[2] + (tl+5), box[3] + 2), color, -1)
|
||||
# draw label
|
||||
tf = max(tl - 1, 1) # font thickness
|
||||
t_size = cv2.getTextSize(name, 0, fontScale=tl / 3, thickness=tf)[0]
|
||||
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
||||
cv2.rectangle(img, c1, c2, color, -1) # filled
|
||||
cv2.putText(img, name, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
@@ -0,0 +1,54 @@
|
||||
import cv2
|
||||
import torch
|
||||
from ..InsightFace.models.data.config import get_config
|
||||
from ..InsightFace.models.utils import update_facebank, load_facebank,special_draw
|
||||
from ..Retinaface.Retinaface import FaceDetector
|
||||
from ..InsightFace.models.Learner import FaceRecognizer
|
||||
class WebcamVerify:
|
||||
"""
|
||||
WebcamVerify: face verfication with cv2
|
||||
|
||||
if you add new person in to facebank
|
||||
pass update True
|
||||
"""
|
||||
def __init__(self, update=True, tta=True, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
print('loading ...')
|
||||
self.tta = tta
|
||||
self.detector = FaceDetector(device=device)
|
||||
self.conf = get_config()
|
||||
self.conf.device = device
|
||||
recognizer = FaceRecognizer(self.conf)
|
||||
recognizer.model.eval()
|
||||
self.recognizer = recognizer
|
||||
|
||||
if update:
|
||||
self.targets, self.names = update_facebank(self.conf, recognizer.model, self.detector, tta=self.tta)
|
||||
else:
|
||||
self.targets, self.names = load_facebank(self.conf)
|
||||
|
||||
|
||||
|
||||
def run(self, camera_index=0):
|
||||
if len(self.targets) < 1:
|
||||
raise Exception("you don't have any person in facebank: add new person with 'add_from_webcam' or 'add_from_folder' function")
|
||||
|
||||
cap = cv2.VideoCapture(camera_index)
|
||||
cap.set(3, 1280)
|
||||
cap.set(4, 720)
|
||||
# frame rate 6 due to my laptop is quite slow...
|
||||
print('type q for exit')
|
||||
while cap.isOpened():
|
||||
ret , frame = cap.read()
|
||||
if ret == False:
|
||||
raise Exception('the camera not recognized: change camera_index param to ' + str(0 if camera_index == 1 else 1))
|
||||
faces, boxes, scores, landmarks = self.detector.detect_align(frame)
|
||||
if len(faces.shape) > 1:
|
||||
results, score = self.recognizer.infer(faces, self.targets, tta=self.tta)
|
||||
for idx, bbox in enumerate(boxes):
|
||||
special_draw(frame, bbox, landmarks[idx], self.names[results[idx] + 1], score[idx])
|
||||
cv2.imshow('face Capture', frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,182 @@
|
||||
import os
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
from skimage import transform
|
||||
|
||||
from .utils.alignment import get_reference_facial_points, FaceWarpException, alignment
|
||||
from .utils.box_utils import decode, decode_landmark, prior_box, nms
|
||||
from .utils.config import cfg_mnet, cfg_re50
|
||||
from .models.retinaface import RetinaFace
|
||||
from ..utils import download_weight
|
||||
|
||||
|
||||
class FaceDetector:
|
||||
|
||||
def __init__(self, name='mobilenet', weight_path=None, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), confidence_threshold=0.99,
|
||||
top_k=5000, nms_threshold=0.4, keep_top_k=750, face_size=(112, 112), crop_size=(96, 112), verbose=True):
|
||||
"""
|
||||
RetinaFace Detector with 5points landmarks
|
||||
Args:
|
||||
name: name of backbone (resnet, mobilenet, slim, rfb)
|
||||
weight_path: path of network weight
|
||||
device: running device (cuda, cpu)
|
||||
face_size: final face size
|
||||
face_padding: padding for bounding boxes
|
||||
"""
|
||||
|
||||
model, cfg = None, None
|
||||
if name == 'mobilenet':
|
||||
cfg = cfg_mnet
|
||||
model = RetinaFace(cfg=cfg, phase='test')
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'mobilenet.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/mobilenet.pth')
|
||||
weight_path="F:/test/ComfyUI/models/stable_makeup/mobilenet0.25_Final.pth"
|
||||
if os.path.isfile(weight_path) == False:
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://drive.google.com/uc?export=download&id=15zP8BP-5IvWXWZoYTNdvUJUiBqZ1hxu1',
|
||||
file_name=file_name,
|
||||
verbose=verbose)
|
||||
os.rename(file_name, weight_path)
|
||||
elif name == 'resnet':
|
||||
cfg = cfg_re50
|
||||
model = RetinaFace(cfg=cfg, phase='test')
|
||||
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'resnet50.pth')
|
||||
weight_path = os.path.join(os.path.dirname(file_name), 'weights/resnet50.pth')
|
||||
if os.path.isfile(weight_path) == False:
|
||||
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
|
||||
download_weight(link='https://www.dropbox.com/s/8sxkgc9voel6ost/resnet50.pth?dl=1',
|
||||
file_name=file_name,
|
||||
verbose=verbose)
|
||||
os.rename(file_name, weight_path)
|
||||
else:
|
||||
exit('FaceDetector Exit: model name can be either mobilenet or resnet')
|
||||
|
||||
|
||||
# settings for model
|
||||
model.load_state_dict(torch.load(weight_path, map_location=device))
|
||||
model.to(device).eval()
|
||||
self.model = model
|
||||
self.device = device
|
||||
self.cfg = cfg
|
||||
# settings for face detection
|
||||
self.thresh = confidence_threshold
|
||||
self.top_k = top_k
|
||||
self.nms_thresh = nms_threshold
|
||||
self.keep_top_k = keep_top_k
|
||||
# settings for face alignment
|
||||
self.trans = transform.SimilarityTransform()
|
||||
self.out_size = face_size
|
||||
self.ref_pts = get_reference_facial_points(output_size=face_size, crop_size=crop_size)
|
||||
|
||||
|
||||
def preprocessor(self, img_raw):
|
||||
img = torch.tensor(img_raw, dtype=torch.float32).to(self.device)
|
||||
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]).to(self.device)
|
||||
img -= torch.tensor([104, 117, 123]).to(self.device)
|
||||
img = img.permute(2, 0, 1).unsqueeze(0)
|
||||
return img, scale
|
||||
|
||||
def detect_faces(self, img_raw):
|
||||
"""
|
||||
get a image from ndarray, detect faces in image
|
||||
Args:
|
||||
img_raw: original image from cv2(BGR) or PIL(RGB)
|
||||
Notes:
|
||||
coordinate is corresponding to original image
|
||||
and type of return image is corresponding to input(cv2, PIL)
|
||||
Returns:
|
||||
boxes:
|
||||
faces bounding box for each face
|
||||
scores:
|
||||
percentage of each face
|
||||
landmarks:
|
||||
faces landmarks for each face
|
||||
"""
|
||||
|
||||
img, scale = self.preprocessor(img_raw)
|
||||
# tic = time.time()
|
||||
with torch.no_grad():
|
||||
loc, conf, landmarks = self.model(img) # forward pass
|
||||
# print('net forward time: {:.4f}'.format(time.time() - tic))
|
||||
|
||||
priors = prior_box(self.cfg, image_size=img.shape[2:]).to(self.device)
|
||||
boxes = decode(loc.data.squeeze(0), priors, self.cfg['variance'])
|
||||
boxes = boxes * scale
|
||||
scores = conf.squeeze(0)[:, 1]
|
||||
landmarks = decode_landmark(landmarks.squeeze(0), priors, self.cfg['variance'])
|
||||
scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2]]).to(self.device)
|
||||
landmarks = landmarks * scale1
|
||||
|
||||
# ignore low scores
|
||||
index = torch.where(scores > self.thresh)[0]
|
||||
boxes = boxes[index]
|
||||
landmarks = landmarks[index]
|
||||
scores = scores[index]
|
||||
|
||||
# keep top-K before NMS
|
||||
order = scores.argsort(dim=0, descending=True)[:self.top_k]
|
||||
boxes = boxes[order]
|
||||
landmarks = landmarks[order]
|
||||
scores = scores[order]
|
||||
|
||||
# Do NMS
|
||||
keep = nms(boxes, scores, self.nms_thresh)
|
||||
boxes = torch.abs(boxes[keep, :])
|
||||
scores = scores[:, None][keep, :]
|
||||
landmarks = landmarks[keep, :].reshape(-1, 5, 2)
|
||||
|
||||
# # keep top-K faster NMS
|
||||
landmarks = landmarks[:self.keep_top_k, :]
|
||||
scores = scores[:self.keep_top_k, :]
|
||||
boxes = boxes[:self.keep_top_k, :]
|
||||
|
||||
return boxes, scores, landmarks
|
||||
|
||||
def detect_align(self, img):
|
||||
"""
|
||||
get a image from ndarray, detect faces in image,
|
||||
cropped face and align face
|
||||
Args:
|
||||
img: original image from cv2(BGR) or PIL(RGB)
|
||||
Notes:
|
||||
coordinate is corresponding to original image
|
||||
and type of return image is corresponding to input(cv2, PIL)
|
||||
|
||||
Returns:
|
||||
faces:
|
||||
a tensor(n, 112, 112, 3) of faces that aligned
|
||||
"""
|
||||
boxes, scores, landmarks = self.detect_faces(img)
|
||||
|
||||
warped = []
|
||||
for src_pts in landmarks:
|
||||
if max(src_pts.shape) < 3 or min(src_pts.shape) != 2:
|
||||
raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')
|
||||
|
||||
if src_pts.shape[0] == 2:
|
||||
src_pts = src_pts.T
|
||||
|
||||
if src_pts.shape != self.ref_pts.shape:
|
||||
raise FaceWarpException('facial_pts and reference_pts must have the same shape')
|
||||
|
||||
self.trans.estimate(src_pts.cpu().numpy(), self.ref_pts)
|
||||
face_img = cv2.warpAffine(img, self.trans.params[0:2, :], self.out_size)
|
||||
warped.append(face_img)
|
||||
|
||||
faces = torch.tensor(np.array(warped)).to(self.device)
|
||||
return faces, boxes, scores, landmarks
|
||||
|
||||
|
||||
def detect(self, img, align=False):
|
||||
boxes, scores, landmarks = self.detect_faces(img)
|
||||
box = boxes[0].int()
|
||||
face = img[box[1]: box[3], box[0]: box[2]]
|
||||
face = cv2.resize(face, self.out_size)
|
||||
if align:
|
||||
landmark = landmarks[0].int()[:3].numpy()
|
||||
return alignment(face, *landmark)
|
||||
else:
|
||||
return face
|
||||
@@ -0,0 +1,33 @@
|
||||
import cv2
|
||||
import torch
|
||||
#from facelib import special_draw
|
||||
#from facelib import FaceDetector
|
||||
from ..InsightFace.models.utils import special_draw
|
||||
from ..Retinaface.Retinaface import FaceDetector
|
||||
class WebcamFaceDetector:
|
||||
def __init__(self, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||||
print('loading ...')
|
||||
self.detector = FaceDetector(face_size=(224, 224), device=device)
|
||||
|
||||
|
||||
def run(self, camera_index=0):
|
||||
cap = cv2.VideoCapture(camera_index)
|
||||
cap.set(3, 1280)
|
||||
cap.set(4, 720)
|
||||
print('type q for exit')
|
||||
while cap.isOpened():
|
||||
ret , frame = cap.read()
|
||||
if ret == False:
|
||||
raise Exception('the camera not recognized: change camera_index param to ' + str(0 if camera_index == 1 else 1))
|
||||
# boxes, scores, landmarks = detector.detect_faces(frame)
|
||||
faces, boxes, scores, landmarks = self.detector.detect_align(frame)
|
||||
if len(faces.shape) > 1:
|
||||
for idx, bbox in enumerate(boxes):
|
||||
special_draw(frame, bbox, landmarks[idx], name='face', score=scores[idx])
|
||||
|
||||
cv2.imshow('frame', frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,137 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def conv_bn(inp, oup, stride=1, leaky=0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True)
|
||||
)
|
||||
|
||||
|
||||
def conv_bn_no_relu(inp, oup, stride):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
)
|
||||
|
||||
|
||||
def conv_bn1X1(inp, oup, stride, leaky=0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True)
|
||||
)
|
||||
|
||||
|
||||
def conv_dw(inp, oup, stride, leaky=0.1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
|
||||
nn.BatchNorm2d(inp),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True),
|
||||
|
||||
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True),
|
||||
)
|
||||
|
||||
|
||||
class SSH(nn.Module):
|
||||
def __init__(self, in_channel, out_channel):
|
||||
super(SSH, self).__init__()
|
||||
assert out_channel % 4 == 0
|
||||
leaky = 0
|
||||
if (out_channel <= 64):
|
||||
leaky = 0.1
|
||||
self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
|
||||
|
||||
self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
|
||||
self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
|
||||
|
||||
self.conv7X7_2 = conv_bn(out_channel // 4, out_channel // 4, stride=1, leaky=leaky)
|
||||
self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
|
||||
|
||||
def forward(self, input):
|
||||
conv3X3 = self.conv3X3(input)
|
||||
|
||||
conv5X5_1 = self.conv5X5_1(input)
|
||||
conv5X5 = self.conv5X5_2(conv5X5_1)
|
||||
|
||||
conv7X7_2 = self.conv7X7_2(conv5X5_1)
|
||||
conv7X7 = self.conv7x7_3(conv7X7_2)
|
||||
|
||||
out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class FPN(nn.Module):
|
||||
def __init__(self, in_channels_list, out_channels):
|
||||
super(FPN, self).__init__()
|
||||
leaky = 0
|
||||
if (out_channels <= 64):
|
||||
leaky = 0.1
|
||||
self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride=1, leaky=leaky)
|
||||
self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride=1, leaky=leaky)
|
||||
self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride=1, leaky=leaky)
|
||||
|
||||
self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
|
||||
self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
|
||||
|
||||
def forward(self, input):
|
||||
# names = list(input.keys())
|
||||
input = list(input.values())
|
||||
|
||||
output1 = self.output1(input[0])
|
||||
output2 = self.output2(input[1])
|
||||
output3 = self.output3(input[2])
|
||||
|
||||
up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest")
|
||||
output2 = output2 + up3
|
||||
output2 = self.merge2(output2)
|
||||
|
||||
up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest")
|
||||
output1 = output1 + up2
|
||||
output1 = self.merge1(output1)
|
||||
|
||||
out = [output1, output2, output3]
|
||||
return out
|
||||
|
||||
|
||||
class MobileNetV1(nn.Module):
|
||||
def __init__(self):
|
||||
super(MobileNetV1, self).__init__()
|
||||
self.stage1 = nn.Sequential(
|
||||
conv_bn(3, 8, 2, leaky=0.1), # 3
|
||||
conv_dw(8, 16, 1), # 7
|
||||
conv_dw(16, 32, 2), # 11
|
||||
conv_dw(32, 32, 1), # 19
|
||||
conv_dw(32, 64, 2), # 27
|
||||
conv_dw(64, 64, 1), # 43
|
||||
)
|
||||
self.stage2 = nn.Sequential(
|
||||
conv_dw(64, 128, 2), # 43 + 16 = 59
|
||||
conv_dw(128, 128, 1), # 59 + 32 = 91
|
||||
conv_dw(128, 128, 1), # 91 + 32 = 123
|
||||
conv_dw(128, 128, 1), # 123 + 32 = 155
|
||||
conv_dw(128, 128, 1), # 155 + 32 = 187
|
||||
conv_dw(128, 128, 1), # 187 + 32 = 219
|
||||
)
|
||||
self.stage3 = nn.Sequential(
|
||||
conv_dw(128, 256, 2), # 219 +3 2 = 241
|
||||
conv_dw(256, 256, 1), # 241 + 64 = 301
|
||||
)
|
||||
self.avg = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.fc = nn.Linear(256, 1000)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stage1(x)
|
||||
x = self.stage2(x)
|
||||
x = self.stage3(x)
|
||||
x = self.avg(x)
|
||||
# x = self.model(x)
|
||||
x = x.view(-1, 256)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
@@ -0,0 +1,118 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.models._utils as _utils
|
||||
|
||||
from ..models.net import FPN as FPN
|
||||
from ..models.net import MobileNetV1 as MobileNetV1
|
||||
from ..models.net import SSH as SSH
|
||||
|
||||
|
||||
class ClassHead(nn.Module):
|
||||
def __init__(self, inchannels=512, num_anchors=3):
|
||||
super(ClassHead, self).__init__()
|
||||
self.num_anchors = num_anchors
|
||||
self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 2)
|
||||
|
||||
|
||||
class BboxHead(nn.Module):
|
||||
def __init__(self, inchannels=512, num_anchors=3):
|
||||
super(BboxHead, self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 4)
|
||||
|
||||
|
||||
class LandmarkHead(nn.Module):
|
||||
def __init__(self, inchannels=512, num_anchors=3):
|
||||
super(LandmarkHead, self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 10)
|
||||
|
||||
|
||||
class RetinaFace(nn.Module):
|
||||
def __init__(self, cfg=None, phase='train'):
|
||||
"""
|
||||
:param cfg: Network related settings.
|
||||
:param phase: train or test.
|
||||
"""
|
||||
super(RetinaFace, self).__init__()
|
||||
self.phase = phase
|
||||
# backbone = MobileNetV1()
|
||||
if cfg['name'] == 'mobilenet0.25':
|
||||
backbone = MobileNetV1()
|
||||
elif cfg['name'] == 'Resnet50':
|
||||
import torchvision.models as models
|
||||
backbone = models.resnet50(pretrained=cfg['pretrain'])
|
||||
|
||||
self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])
|
||||
in_channels_stage2 = cfg['in_channel']
|
||||
in_channels_list = [
|
||||
in_channels_stage2 * 2,
|
||||
in_channels_stage2 * 4,
|
||||
in_channels_stage2 * 8,
|
||||
]
|
||||
out_channels = cfg['out_channel']
|
||||
self.fpn = FPN(in_channels_list, out_channels)
|
||||
self.ssh1 = SSH(out_channels, out_channels)
|
||||
self.ssh2 = SSH(out_channels, out_channels)
|
||||
self.ssh3 = SSH(out_channels, out_channels)
|
||||
|
||||
self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
|
||||
def _make_class_head(self, fpn_num=3, inchannels=64, anchor_num=2):
|
||||
classhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
classhead.append(ClassHead(inchannels, anchor_num))
|
||||
return classhead
|
||||
|
||||
def _make_bbox_head(self, fpn_num=3, inchannels=64, anchor_num=2):
|
||||
bboxhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
bboxhead.append(BboxHead(inchannels, anchor_num))
|
||||
return bboxhead
|
||||
|
||||
def _make_landmark_head(self, fpn_num=3, inchannels=64, anchor_num=2):
|
||||
landmarkhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
landmarkhead.append(LandmarkHead(inchannels, anchor_num))
|
||||
return landmarkhead
|
||||
|
||||
def forward(self, inputs):
|
||||
out = self.body(inputs)
|
||||
|
||||
# FPN
|
||||
fpn = self.fpn(out)
|
||||
|
||||
# SSH
|
||||
feature1 = self.ssh1(fpn[0])
|
||||
feature2 = self.ssh2(fpn[1])
|
||||
feature3 = self.ssh3(fpn[2])
|
||||
features = [feature1, feature2, feature3]
|
||||
|
||||
bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
|
||||
if self.phase == 'train':
|
||||
output = (bbox_regressions, classifications, ldm_regressions)
|
||||
else:
|
||||
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
||||
return output
|
||||
@@ -0,0 +1,197 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class BasicConv(nn.Module):
|
||||
|
||||
def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True):
|
||||
super(BasicConv, self).__init__()
|
||||
self.out_channels = out_planes
|
||||
if bn:
|
||||
self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=False)
|
||||
self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True)
|
||||
self.relu = nn.ReLU(inplace=True) if relu else None
|
||||
else:
|
||||
self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=True)
|
||||
self.bn = None
|
||||
self.relu = nn.ReLU(inplace=True) if relu else None
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
if self.bn is not None:
|
||||
x = self.bn(x)
|
||||
if self.relu is not None:
|
||||
x = self.relu(x)
|
||||
return x
|
||||
|
||||
|
||||
class BasicRFB(nn.Module):
|
||||
|
||||
def __init__(self, in_planes, out_planes, stride=1, scale=0.1, map_reduce=8, vision=1, groups=1):
|
||||
super(BasicRFB, self).__init__()
|
||||
self.scale = scale
|
||||
self.out_channels = out_planes
|
||||
inter_planes = in_planes // map_reduce
|
||||
|
||||
self.branch0 = nn.Sequential(
|
||||
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
|
||||
BasicConv(inter_planes, 2 * inter_planes, kernel_size=(3, 3), stride=stride, padding=(1, 1), groups=groups),
|
||||
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 1, dilation=vision + 1, relu=False, groups=groups)
|
||||
)
|
||||
self.branch1 = nn.Sequential(
|
||||
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
|
||||
BasicConv(inter_planes, 2 * inter_planes, kernel_size=(3, 3), stride=stride, padding=(1, 1), groups=groups),
|
||||
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 2, dilation=vision + 2, relu=False, groups=groups)
|
||||
)
|
||||
self.branch2 = nn.Sequential(
|
||||
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
|
||||
BasicConv(inter_planes, (inter_planes // 2) * 3, kernel_size=3, stride=1, padding=1, groups=groups),
|
||||
BasicConv((inter_planes // 2) * 3, 2 * inter_planes, kernel_size=3, stride=stride, padding=1, groups=groups),
|
||||
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 4, dilation=vision + 4, relu=False, groups=groups)
|
||||
)
|
||||
|
||||
self.ConvLinear = BasicConv(6 * inter_planes, out_planes, kernel_size=1, stride=1, relu=False)
|
||||
self.shortcut = BasicConv(in_planes, out_planes, kernel_size=1, stride=stride, relu=False)
|
||||
self.relu = nn.ReLU(inplace=False)
|
||||
|
||||
def forward(self, x):
|
||||
x0 = self.branch0(x)
|
||||
x1 = self.branch1(x)
|
||||
x2 = self.branch2(x)
|
||||
|
||||
out = torch.cat((x0, x1, x2), 1)
|
||||
out = self.ConvLinear(out)
|
||||
short = self.shortcut(x)
|
||||
out = out * self.scale + short
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
|
||||
def conv_bn(inp, oup, stride = 1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def depth_conv2d(inp, oup, kernel=1, stride=1, pad=0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, kernel_size = kernel, stride = stride, padding=pad, groups=inp),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(inp, oup, kernel_size=1)
|
||||
)
|
||||
|
||||
def conv_dw(inp, oup, stride):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
|
||||
nn.BatchNorm2d(inp),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
class RFB(nn.Module):
|
||||
def __init__(self, cfg = None, phase = 'train'):
|
||||
"""
|
||||
:param cfg: Network related settings.
|
||||
:param phase: train or test.
|
||||
"""
|
||||
super(RFB,self).__init__()
|
||||
self.phase = phase
|
||||
self.num_classes = 2
|
||||
|
||||
self.conv1 = conv_bn(3, 16, 2)
|
||||
self.conv2 = conv_dw(16, 32, 1)
|
||||
self.conv3 = conv_dw(32, 32, 2)
|
||||
self.conv4 = conv_dw(32, 32, 1)
|
||||
self.conv5 = conv_dw(32, 64, 2)
|
||||
self.conv6 = conv_dw(64, 64, 1)
|
||||
self.conv7 = conv_dw(64, 64, 1)
|
||||
self.conv8 = BasicRFB(64, 64, stride=1, scale=1.0)
|
||||
|
||||
self.conv9 = conv_dw(64, 128, 2)
|
||||
self.conv10 = conv_dw(128, 128, 1)
|
||||
self.conv11 = conv_dw(128, 128, 1)
|
||||
|
||||
self.conv12 = conv_dw(128, 256, 2)
|
||||
self.conv13 = conv_dw(256, 256, 1)
|
||||
|
||||
self.conv14 = nn.Sequential(
|
||||
nn.Conv2d(in_channels=256, out_channels=64, kernel_size=1),
|
||||
nn.ReLU(inplace=True),
|
||||
depth_conv2d(64, 256, kernel=3, stride=2, pad=1),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
self.loc, self.conf, self.landm = self.multibox(self.num_classes);
|
||||
|
||||
def multibox(self, num_classes):
|
||||
loc_layers = []
|
||||
conf_layers = []
|
||||
landm_layers = []
|
||||
loc_layers += [depth_conv2d(64, 3 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(64, 3 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(64, 3 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [depth_conv2d(128, 2 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(128, 2 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(128, 2 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [depth_conv2d(256, 2 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(256, 2 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(256, 2 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [nn.Conv2d(256, 3 * 4, kernel_size=3, padding=1)]
|
||||
conf_layers += [nn.Conv2d(256, 3 * num_classes, kernel_size=3, padding=1)]
|
||||
landm_layers += [nn.Conv2d(256, 3 * 10, kernel_size=3, padding=1)]
|
||||
return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers), nn.Sequential(*landm_layers)
|
||||
|
||||
|
||||
def forward(self,inputs):
|
||||
detections = list()
|
||||
loc = list()
|
||||
conf = list()
|
||||
landm = list()
|
||||
|
||||
x1 = self.conv1(inputs)
|
||||
x2 = self.conv2(x1)
|
||||
x3 = self.conv3(x2)
|
||||
x4 = self.conv4(x3)
|
||||
x5 = self.conv5(x4)
|
||||
x6 = self.conv6(x5)
|
||||
x7 = self.conv7(x6)
|
||||
x8 = self.conv8(x7)
|
||||
detections.append(x8)
|
||||
|
||||
x9 = self.conv9(x8)
|
||||
x10 = self.conv10(x9)
|
||||
x11 = self.conv11(x10)
|
||||
detections.append(x11)
|
||||
|
||||
x12 = self.conv12(x11)
|
||||
x13 = self.conv13(x12)
|
||||
detections.append(x13)
|
||||
|
||||
x14= self.conv14(x13)
|
||||
detections.append(x14)
|
||||
|
||||
for (x, l, c, lam) in zip(detections, self.loc, self.conf, self.landm):
|
||||
loc.append(l(x).permute(0, 2, 3, 1).contiguous())
|
||||
conf.append(c(x).permute(0, 2, 3, 1).contiguous())
|
||||
landm.append(lam(x).permute(0, 2, 3, 1).contiguous())
|
||||
|
||||
bbox_regressions = torch.cat([o.view(o.size(0), -1, 4) for o in loc], 1)
|
||||
classifications = torch.cat([o.view(o.size(0), -1, 2) for o in conf], 1)
|
||||
ldm_regressions = torch.cat([o.view(o.size(0), -1, 10) for o in landm], 1)
|
||||
|
||||
|
||||
|
||||
if self.phase == 'train':
|
||||
output = (bbox_regressions, classifications, ldm_regressions)
|
||||
else:
|
||||
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
||||
return output
|
||||
@@ -0,0 +1,129 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
def conv_bn(inp, oup, stride = 1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def depth_conv2d(inp, oup, kernel=1, stride=1, pad=0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, kernel_size = kernel, stride = stride, padding=pad, groups=inp),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(inp, oup, kernel_size=1)
|
||||
)
|
||||
|
||||
def conv_dw(inp, oup, stride):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
|
||||
nn.BatchNorm2d(inp),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
class Slim(nn.Module):
|
||||
def __init__(self, cfg = None, phase = 'train'):
|
||||
"""
|
||||
:param cfg: Network related settings.
|
||||
:param phase: train or test.
|
||||
"""
|
||||
super(Slim, self).__init__()
|
||||
self.phase = phase
|
||||
self.num_classes = 2
|
||||
|
||||
self.conv1 = conv_bn(3, 16, 2)
|
||||
self.conv2 = conv_dw(16, 32, 1)
|
||||
self.conv3 = conv_dw(32, 32, 2)
|
||||
self.conv4 = conv_dw(32, 32, 1)
|
||||
self.conv5 = conv_dw(32, 64, 2)
|
||||
self.conv6 = conv_dw(64, 64, 1)
|
||||
self.conv7 = conv_dw(64, 64, 1)
|
||||
self.conv8 = conv_dw(64, 64, 1)
|
||||
|
||||
self.conv9 = conv_dw(64, 128, 2)
|
||||
self.conv10 = conv_dw(128, 128, 1)
|
||||
self.conv11 = conv_dw(128, 128, 1)
|
||||
|
||||
self.conv12 = conv_dw(128, 256, 2)
|
||||
self.conv13 = conv_dw(256, 256, 1)
|
||||
|
||||
self.conv14 = nn.Sequential(
|
||||
nn.Conv2d(in_channels=256, out_channels=64, kernel_size=1),
|
||||
nn.ReLU(inplace=True),
|
||||
depth_conv2d(64, 256, kernel=3, stride=2, pad=1),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
self.loc, self.conf, self.landm = self.multibox(self.num_classes);
|
||||
|
||||
def multibox(self, num_classes):
|
||||
loc_layers = []
|
||||
conf_layers = []
|
||||
landm_layers = []
|
||||
loc_layers += [depth_conv2d(64, 3 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(64, 3 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(64, 3 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [depth_conv2d(128, 2 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(128, 2 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(128, 2 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [depth_conv2d(256, 2 * 4, kernel=3, pad=1)]
|
||||
conf_layers += [depth_conv2d(256, 2 * num_classes, kernel=3, pad=1)]
|
||||
landm_layers += [depth_conv2d(256, 2 * 10, kernel=3, pad=1)]
|
||||
|
||||
loc_layers += [nn.Conv2d(256, 3 * 4, kernel_size=3, padding=1)]
|
||||
conf_layers += [nn.Conv2d(256, 3 * num_classes, kernel_size=3, padding=1)]
|
||||
landm_layers += [nn.Conv2d(256, 3 * 10, kernel_size=3, padding=1)]
|
||||
return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers), nn.Sequential(*landm_layers)
|
||||
|
||||
|
||||
def forward(self,inputs):
|
||||
detections = list()
|
||||
loc = list()
|
||||
conf = list()
|
||||
landm = list()
|
||||
|
||||
x1 = self.conv1(inputs)
|
||||
x2 = self.conv2(x1)
|
||||
x3 = self.conv3(x2)
|
||||
x4 = self.conv4(x3)
|
||||
x5 = self.conv5(x4)
|
||||
x6 = self.conv6(x5)
|
||||
x7 = self.conv7(x6)
|
||||
x8 = self.conv8(x7)
|
||||
detections.append(x8)
|
||||
|
||||
x9 = self.conv9(x8)
|
||||
x10 = self.conv10(x9)
|
||||
x11 = self.conv11(x10)
|
||||
detections.append(x11)
|
||||
|
||||
x12 = self.conv12(x11)
|
||||
x13 = self.conv13(x12)
|
||||
detections.append(x13)
|
||||
|
||||
x14= self.conv14(x13)
|
||||
detections.append(x14)
|
||||
|
||||
for (x, l, c, lam) in zip(detections, self.loc, self.conf, self.landm):
|
||||
loc.append(l(x).permute(0, 2, 3, 1).contiguous())
|
||||
conf.append(c(x).permute(0, 2, 3, 1).contiguous())
|
||||
landm.append(lam(x).permute(0, 2, 3, 1).contiguous())
|
||||
|
||||
bbox_regressions = torch.cat([o.view(o.size(0), -1, 4) for o in loc], 1)
|
||||
classifications = torch.cat([o.view(o.size(0), -1, 2) for o in conf], 1)
|
||||
ldm_regressions = torch.cat([o.view(o.size(0), -1, 10) for o in landm], 1)
|
||||
|
||||
|
||||
|
||||
if self.phase == 'train':
|
||||
output = (bbox_regressions, classifications, ldm_regressions)
|
||||
else:
|
||||
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
||||
return output
|
||||
@@ -0,0 +1,94 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Created on Mon Apr 1 15:43:29 2020
|
||||
@author: Sajjad Ayobbi
|
||||
"""
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from skimage import transform as trans
|
||||
|
||||
# reference facial points, a list of coordinates (x,y)
|
||||
REFERENCE_FACIAL_POINTS = [
|
||||
[30.29459953, 51.69630051],
|
||||
[65.53179932, 51.50139999],
|
||||
[48.02519989, 71.73660278],
|
||||
[33.54930115, 87],
|
||||
[62.72990036, 87]
|
||||
]
|
||||
|
||||
class FaceWarpException(Exception):
|
||||
def __str__(self):
|
||||
return 'In File {}:{}'.format(
|
||||
__file__, super.__str__(self))
|
||||
|
||||
|
||||
def get_reference_facial_points(output_size=(112, 112), crop_size=(96, 112)):
|
||||
|
||||
tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
|
||||
tmp_crop_size = np.array(crop_size)
|
||||
|
||||
# size_diff = max(tmp_crop_size) - tmp_crop_size
|
||||
# tmp_5pts += size_diff / 2
|
||||
# tmp_crop_size += size_diff
|
||||
# return tmp_5pts
|
||||
|
||||
x_scale = output_size[0]/tmp_crop_size[0]
|
||||
y_scale = output_size[1]/tmp_crop_size[1]
|
||||
tmp_5pts[:, 0] *= x_scale
|
||||
tmp_5pts[:, 1] *= y_scale
|
||||
|
||||
return tmp_5pts
|
||||
|
||||
|
||||
|
||||
|
||||
####################################################################################
|
||||
|
||||
|
||||
def findEuclideanDistance(source_representation, test_representation):
|
||||
euclidean_distance = source_representation - test_representation
|
||||
euclidean_distance = np.sum(np.multiply(euclidean_distance, euclidean_distance))
|
||||
euclidean_distance = np.sqrt(euclidean_distance)
|
||||
return euclidean_distance
|
||||
|
||||
|
||||
def alignment(img, left_eye, right_eye, nose):
|
||||
left_eye_x, left_eye_y = left_eye
|
||||
right_eye_x, right_eye_y = right_eye
|
||||
#-----------------------
|
||||
upside_down = False
|
||||
if nose[1] < left_eye[1] or nose[1] < right_eye[1]:
|
||||
upside_down = True
|
||||
#-----------------------
|
||||
#find rotation direction
|
||||
if left_eye_y > right_eye_y:
|
||||
point_3rd = (right_eye_x, left_eye_y)
|
||||
direction = -1 #rotate same direction to clock
|
||||
else:
|
||||
point_3rd = (left_eye_x, right_eye_y)
|
||||
direction = 1 #rotate inverse direction of clock
|
||||
|
||||
#-----------------------
|
||||
#find length of triangle edges
|
||||
a = findEuclideanDistance(np.array(left_eye), np.array(point_3rd))
|
||||
b = findEuclideanDistance(np.array(right_eye), np.array(point_3rd))
|
||||
c = findEuclideanDistance(np.array(right_eye), np.array(left_eye))
|
||||
#-----------------------
|
||||
|
||||
if b != 0 and c != 0: #this multiplication causes division by zero in cos_a calculation
|
||||
cos_a = (b*b + c*c - a*a)/(2*b*c)
|
||||
cos_a = min(1.0, max(-1.0, cos_a))
|
||||
angle = np.arccos(cos_a) #angle in radian
|
||||
angle = (angle * 180) / math.pi #radian to degree
|
||||
#-----------------------
|
||||
#rotate base image
|
||||
if direction == -1:
|
||||
angle = 90 - angle
|
||||
if upside_down == True:
|
||||
angle = angle + 90
|
||||
img = Image.fromarray(img)
|
||||
img = np.array(img.rotate(direction * angle))
|
||||
|
||||
return img
|
||||
@@ -0,0 +1,102 @@
|
||||
from itertools import product as product
|
||||
from math import ceil
|
||||
import torch
|
||||
|
||||
|
||||
# Original author: Francisco Massa:
|
||||
# https://github.com/fmassa/object-detection.torch
|
||||
# Ported to PyTorch by Max deGroot (02/01/2017)
|
||||
|
||||
def prior_box(cfg, image_size=None):
|
||||
steps = cfg['steps']
|
||||
feature_maps = [[ceil(image_size[0] / step), ceil(image_size[1] / step)] for step in steps]
|
||||
min_sizes_ = cfg['min_sizes']
|
||||
anchors = []
|
||||
|
||||
for k, f in enumerate(feature_maps):
|
||||
min_sizes = min_sizes_[k]
|
||||
for i, j in product(range(f[0]), range(f[1])):
|
||||
for min_size in min_sizes:
|
||||
s_kx = min_size / image_size[1]
|
||||
s_ky = min_size / image_size[0]
|
||||
dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
|
||||
dense_cy = [y * steps[k] / image_size[0] for y in [i + 0.5]]
|
||||
for cy, cx in product(dense_cy, dense_cx):
|
||||
anchors += [cx, cy, s_kx, s_ky]
|
||||
|
||||
# back to torch land
|
||||
output = torch.Tensor(anchors).view(-1, 4)
|
||||
return output
|
||||
|
||||
|
||||
# Adapted from https://github.com/Hakuyume/chainer-ssd
|
||||
def decode(loc, priors, variances):
|
||||
"""Decode locations from predictions using priors to undo
|
||||
the encoding we did for offset regression at train time.
|
||||
Args:
|
||||
loc (tensor): location predictions for loc layers,
|
||||
Shape: [num_priors,4]
|
||||
priors (tensor): Prior boxes in center-offset form.
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
decoded bounding box predictions
|
||||
"""
|
||||
|
||||
boxes = torch.cat((
|
||||
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
||||
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
|
||||
boxes[:, :2] -= boxes[:, 2:] / 2
|
||||
boxes[:, 2:] += boxes[:, :2]
|
||||
return boxes
|
||||
|
||||
|
||||
def decode_landmark(pre, priors, variances):
|
||||
"""Decode landm from predictions using priors to undo
|
||||
the encoding we did for offset regression at train time.
|
||||
Args:
|
||||
pre (tensor): landm predictions for loc layers,
|
||||
Shape: [num_priors,10]
|
||||
priors (tensor): Prior boxes in center-offset form.
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
decoded landm predictions
|
||||
"""
|
||||
landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
|
||||
priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
|
||||
priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
|
||||
priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
|
||||
priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
|
||||
), dim=1)
|
||||
return landms
|
||||
|
||||
|
||||
def nms(box, scores, thresh):
|
||||
x1 = box[:, 0]
|
||||
y1 = box[:, 1]
|
||||
x2 = box[:, 2]
|
||||
y2 = box[:, 3]
|
||||
zero = torch.tensor([0.0]).to(scores.device)
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort(descending=True)
|
||||
|
||||
keep = []
|
||||
while order.shape[0] > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = torch.max(x1[i], x1[order[1:]])
|
||||
yy1 = torch.max(y1[i], y1[order[1:]])
|
||||
xx2 = torch.min(x2[i], x2[order[1:]])
|
||||
yy2 = torch.min(y2[i], y2[order[1:]])
|
||||
|
||||
w = torch.max(zero, xx2 - xx1 + 1)
|
||||
h = torch.max(zero, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = torch.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
@@ -0,0 +1,74 @@
|
||||
# config.py
|
||||
|
||||
cfg_mnet = {
|
||||
'name': 'mobilenet0.25',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 640,
|
||||
'pretrain': True,
|
||||
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
|
||||
'in_channel': 32,
|
||||
'out_channel': 64
|
||||
}
|
||||
|
||||
cfg_re50 = {
|
||||
'name': 'Resnet50',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 24,
|
||||
'ngpu': 4,
|
||||
'epoch': 100,
|
||||
'decay1': 70,
|
||||
'decay2': 90,
|
||||
'image_size': 840,
|
||||
'pretrain': False,
|
||||
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
|
||||
'in_channel': 256,
|
||||
'out_channel': 256
|
||||
}
|
||||
|
||||
cfg_slim = {
|
||||
'name': 'slim',
|
||||
'min_sizes': [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]],
|
||||
'steps': [8, 16, 32, 64],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 300
|
||||
}
|
||||
|
||||
cfg_rfb = {
|
||||
'name': 'RFB',
|
||||
'min_sizes': [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]],
|
||||
'steps': [8, 16, 32, 64],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 300
|
||||
}
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
from .Retinaface.Retinaface import FaceDetector
|
||||
from .AgeGender.Detector import AgeGenderEstimator
|
||||
from .FacialExpression.FaceExpression import EmotionDetector
|
||||
from .InsightFace.models.Learner import FaceRecognizer
|
||||
from .InsightFace.models.utils import update_facebank, load_facebank, special_draw
|
||||
from .InsightFace.models.data.config import get_config
|
||||
from .InsightFace.add_face import add_from_webcam, add_from_folder
|
||||
# webcame classes
|
||||
from .InsightFace.verifier import WebcamVerify
|
||||
from .AgeGender.from_camera import WebcamAgeGenderEstimator
|
||||
from .Retinaface.from_camera import WebcamFaceDetector
|
||||
from .FacialExpression.from_camera import WebcamEmotionDetector
|
||||
|
||||
__all__ = ['FaceDetector', 'AgeGenderEstimator', 'EmotionDetector', 'FaceRecognizer', 'WebcamVerify', 'WebcamAgeGenderEstimator', 'WebcamFaceDetector',
|
||||
'get_config', 'update_facebank', 'load_facebank', 'special_draw', 'add_from_webcam', 'add_from_folder', 'WebcamEmotionDetector']
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
from tqdm import tqdm
|
||||
import requests
|
||||
import os
|
||||
|
||||
def download_weight(link, file_name, verbose=True):
|
||||
response = requests.get(link, stream=True)
|
||||
total_size_in_bytes= int(response.headers.get('content-length', 0))
|
||||
block_size = 1024 #1 Kibibyte
|
||||
progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True,
|
||||
desc='downloading defualt weights',
|
||||
disable=False if verbose else True)
|
||||
|
||||
with open(file_name, 'wb') as file:
|
||||
for data in response.iter_content(block_size):
|
||||
progress_bar.update(len(data))
|
||||
file.write(data)
|
||||
progress_bar.close()
|
||||
if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
|
||||
exit("ERROR, something went wrong (check your connection)")
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
import os
|
||||
import torch
|
||||
from PIL import Image
|
||||
from diffusers import UNet2DConditionModel as OriginalUNet2DConditionModel
|
||||
from utils.pipeline_sd15 import StableDiffusionControlNetPipeline
|
||||
from diffusers import DDIMScheduler, DPMSolverMultistepScheduler, ControlNetModel
|
||||
from diffusers.utils import load_image
|
||||
from detail_encoder.encoder_plus import detail_encoder
|
||||
from spiga_draw import *
|
||||
from PIL import Image
|
||||
from spiga.inference.config import ModelConfig
|
||||
from spiga.inference.framework import SPIGAFramework
|
||||
from facelib import FaceDetector
|
||||
|
||||
processor = SPIGAFramework(ModelConfig("300wpublic"))
|
||||
detector = FaceDetector(weight_path="./models/mobilenet0.25_Final.pth")
|
||||
|
||||
def get_draw(pil_img, size):
|
||||
spigas = spiga_process(pil_img, detector)
|
||||
if spigas == False:
|
||||
width, height = pil_img.size
|
||||
black_image_pil = Image.new('RGB', (width, height), color=(0, 0, 0))
|
||||
return black_image_pil
|
||||
else:
|
||||
spigas_faces = spiga_segmentation(spigas, size=size)
|
||||
return spigas_faces
|
||||
|
||||
def is_image_file(filename):
|
||||
return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])
|
||||
|
||||
def concatenate_images(image_files, output_file):
|
||||
images = image_files # list
|
||||
max_height = max(img.height for img in images)
|
||||
images = [img.resize((img.width, max_height)) for img in images]
|
||||
total_width = sum(img.width for img in images)
|
||||
combined = Image.new('RGB', (total_width, max_height))
|
||||
x_offset = 0
|
||||
for img in images:
|
||||
combined.paste(img, (x_offset, 0))
|
||||
x_offset += img.width
|
||||
combined.save(output_file)
|
||||
|
||||
model_id = "sd_model_v1-5". # your sdv1-5 path
|
||||
makeup_encoder_path = "./models/stablemakeup/pytorch_model.bin"
|
||||
id_encoder_path = "./models/stablemakeup/pytorch_model_1.bin"
|
||||
pose_encoder_path = "./models/stablemakeup/pytorch_model_2.bin"
|
||||
Unet = OriginalUNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to("cuda")
|
||||
|
||||
id_encoder = ControlNetModel.from_unet(Unet)
|
||||
pose_encoder = ControlNetModel.from_unet(Unet)
|
||||
makeup_encoder = detail_encoder(Unet, "./models/image_encoder_l", "cuda", dtype=torch.float32)
|
||||
makeup_state_dict = torch.load(makeup_encoder_path)
|
||||
id_state_dict = torch.load(id_encoder_path)
|
||||
id_encoder.load_state_dict(id_state_dict, strict=False)
|
||||
pose_state_dict = torch.load(pose_encoder_path)
|
||||
pose_encoder.load_state_dict(pose_state_dict, strict=False)
|
||||
makeup_encoder.load_state_dict(makeup_state_dict, strict=False)
|
||||
id_encoder.to("cuda")
|
||||
pose_encoder.to("cuda")
|
||||
makeup_encoder.to("cuda")
|
||||
|
||||
|
||||
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
||||
model_id,
|
||||
safety_checker=None,
|
||||
unet=Unet,
|
||||
controlnet=[id_encoder, pose_encoder],
|
||||
torch_dtype=torch.float32).to("cuda")
|
||||
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
def infer():
|
||||
id_folder = "./test_imgs/id"
|
||||
makeup_folder = "./test_imgs/makeup"
|
||||
out_folder = "./output"
|
||||
os.makedirs(out_folder, exist_ok=True)
|
||||
for name in os.listdir(id_folder):
|
||||
if not is_image_file(name):
|
||||
continue
|
||||
id_image = load_image(os.path.join(id_folder, name)).resize((512, 512))
|
||||
for mu in os.listdir(makeup_folder):
|
||||
if not is_image_file(mu):
|
||||
continue
|
||||
makeup_image = load_image(os.path.join(makeup_folder, mu)).resize((512, 512))
|
||||
pose_image = get_draw(id_image, size=512)
|
||||
for k in range(3):
|
||||
result_img = makeup_encoder.generate(id_image=[id_image, pose_image], makeup_image=makeup_image,
|
||||
pipe=pipe, guidance_scale=1.6)
|
||||
result_img.save(os.path.join(out_folder, name.split(".")[0] + "_" + mu.split(".")[0] + "_" + str(k) + ".jpg"))
|
||||
|
||||
if __name__ == '__main__':
|
||||
infer()
|
||||
@@ -0,0 +1,229 @@
|
||||
accelerate==0.23.0
|
||||
addict==2.4.0
|
||||
aiofiles==23.2.1
|
||||
aiohttp==3.8.6
|
||||
aiosignal==1.3.1
|
||||
albumentations==1.3.1
|
||||
altair==5.2.0
|
||||
annotated-types==0.6.0
|
||||
antlr4-python3-runtime==4.9.3
|
||||
anyio==4.2.0
|
||||
appdirs==1.4.4
|
||||
asgiref==3.7.2
|
||||
asttokens==2.4.1
|
||||
async-timeout==4.0.3
|
||||
attrs==23.1.0
|
||||
av==11.0.0
|
||||
boto3==1.29.5
|
||||
botocore==1.32.5
|
||||
cairocffi==1.6.1
|
||||
CairoSVG==2.7.1
|
||||
certifi==2023.7.22
|
||||
cffi==1.16.0
|
||||
chardet==4.0.0
|
||||
charset-normalizer==3.3.0
|
||||
click==8.1.7
|
||||
cmake==3.27.6
|
||||
coloredlogs==15.0.1
|
||||
comm==0.2.0
|
||||
contourpy==1.1.1
|
||||
cssselect2==0.7.0
|
||||
cssutils==2.9.0
|
||||
cycler==0.10.0
|
||||
Cython==3.0.8
|
||||
datasets==2.14.6
|
||||
decorator==5.1.1
|
||||
defusedxml==0.7.1
|
||||
Deprecated==1.2.14
|
||||
diffusers==0.25.1
|
||||
diffvg==0.0.1
|
||||
dill==0.3.7
|
||||
Django==4.2.7
|
||||
dlib==19.24.2
|
||||
dnspython==2.4.2
|
||||
docker-pycreds==0.4.0
|
||||
docopt==0.6.2
|
||||
easydict==1.11
|
||||
einops==0.7.0
|
||||
ete3==3.1.3
|
||||
exceptiongroup==1.2.0
|
||||
executing==2.0.1
|
||||
facelib @ git+https://github.com/sajjjadayobi/FaceLib.git@382841efec823f059f49754e5d378abb12cfb551
|
||||
fastapi==0.109.0
|
||||
ffmpy==0.3.1
|
||||
filelock==3.12.4
|
||||
flatbuffers==23.5.26
|
||||
fonttools==4.43.1
|
||||
frozenlist==1.4.0
|
||||
fsspec==2023.9.2
|
||||
fvcore==0.1.5.post20221221
|
||||
gitdb==4.0.10
|
||||
GitPython==3.1.37
|
||||
gradio==3.27.0
|
||||
gradio_client==0.8.0
|
||||
greenlet==3.0.1
|
||||
-e git+https://github.com/IDEA-Research/GroundingDINO.git@57535c5a79791cb76e36fdb64975271354f10251#egg=groundingdino
|
||||
h11==0.14.0
|
||||
hjson==3.1.0
|
||||
httpcore==1.0.2
|
||||
httpx==0.26.0
|
||||
huggingface-hub==0.17.3
|
||||
humanfriendly==10.0
|
||||
idna==2.10
|
||||
imagecorruptions==1.1.2
|
||||
imageio==2.31.5
|
||||
imageio-ffmpeg==0.4.9
|
||||
imgaug==0.4.0
|
||||
importlib-metadata==6.8.0
|
||||
insightface==0.7.3
|
||||
iopath==0.1.10
|
||||
ipdb==0.13.13
|
||||
ipython==8.17.2
|
||||
ipywidgets==8.1.1
|
||||
jedi==0.19.1
|
||||
Jinja2==3.1.2
|
||||
jmespath==1.0.1
|
||||
jmspath==0.0.1
|
||||
joblib==1.3.2
|
||||
jsonpatch==1.33
|
||||
jsonpointer==2.4
|
||||
jsonschema==4.20.0
|
||||
jsonschema-specifications==2023.12.1
|
||||
jupyterlab-widgets==3.0.9
|
||||
kiwisolver==1.4.5
|
||||
lazy_loader==0.3
|
||||
lightning-utilities==0.9.0
|
||||
linkify-it-py==2.0.2
|
||||
lit==17.0.2
|
||||
llvmlite==0.41.1
|
||||
markdown-it-py==2.2.0
|
||||
MarkupSafe==2.1.3
|
||||
matplotlib==3.8.0
|
||||
matplotlib-inline==0.1.6
|
||||
mdit-py-plugins==0.3.3
|
||||
mdurl==0.1.2
|
||||
moviepy==1.0.3
|
||||
mpmath==1.3.0
|
||||
multidict==6.0.4
|
||||
multiprocess==0.70.15
|
||||
networkx==3.1
|
||||
ninja==1.11.1.1
|
||||
numba==0.58.1
|
||||
numpy==1.26.0
|
||||
nvidia-cublas-cu11==11.10.3.66
|
||||
nvidia-cuda-cupti-cu11==11.7.101
|
||||
nvidia-cuda-nvrtc-cu11==11.7.99
|
||||
nvidia-cuda-runtime-cu11==11.7.99
|
||||
nvidia-cudnn-cu11==8.5.0.96
|
||||
nvidia-cufft-cu11==10.9.0.58
|
||||
nvidia-curand-cu11==10.2.10.91
|
||||
nvidia-cusolver-cu11==11.4.0.1
|
||||
nvidia-cusparse-cu11==11.7.4.91
|
||||
nvidia-nccl-cu11==2.14.3
|
||||
nvidia-nvtx-cu11==11.7.91
|
||||
omegaconf==2.3.0
|
||||
onnx==1.15.0
|
||||
onnxruntime==1.16.3
|
||||
opencv-python==4.8.1.78
|
||||
opencv-python-headless==4.8.0.74
|
||||
orderedset==2.0.3
|
||||
orjson==3.9.10
|
||||
packaging==23.2
|
||||
pandas==2.1.1
|
||||
parso==0.8.3
|
||||
pathtools==0.1.2
|
||||
pexpect==4.8.0
|
||||
Pillow==10.0.1
|
||||
platformdirs==4.1.0
|
||||
portalocker==2.8.2
|
||||
prettytable==3.9.0
|
||||
proglog==0.1.10
|
||||
progressbar==2.5
|
||||
prompt-toolkit==3.0.39
|
||||
protobuf==4.24.3
|
||||
psutil==5.9.5
|
||||
ptyprocess==0.7.0
|
||||
pure-eval==0.2.2
|
||||
pyaml==23.9.7
|
||||
pyarrow==14.0.0
|
||||
pycocotools==2.0.7
|
||||
pycparser==2.21
|
||||
pydantic==2.5.3
|
||||
pydantic_core==2.14.6
|
||||
pydub==0.25.1
|
||||
Pygments==2.17.2
|
||||
pymongo==4.6.0
|
||||
pyparsing==2.4.7
|
||||
python-dateutil==2.8.2
|
||||
python-dotenv==1.0.0
|
||||
python-magic==0.4.27
|
||||
python-multipart==0.0.6
|
||||
pytorch-lightning==2.1.1
|
||||
pytz==2023.3.post1
|
||||
PyYAML==6.0.1
|
||||
qudida==0.0.4
|
||||
referencing==0.32.1
|
||||
regex==2023.10.3
|
||||
requests==2.31.0
|
||||
requests-toolbelt==1.0.0
|
||||
roboflow==1.1.12
|
||||
rpds-py==0.17.1
|
||||
s3transfer==0.7.0
|
||||
safetensors==0.3.3
|
||||
scikit-image==0.22.0
|
||||
scikit-learn==1.4.0
|
||||
scipy==1.11.3
|
||||
seaborn==0.13.0
|
||||
segment-anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588
|
||||
semantic-version==2.10.0
|
||||
sentry-sdk==1.31.0
|
||||
setproctitle==1.3.3
|
||||
shapely==2.0.2
|
||||
simplejson==3.19.2
|
||||
six==1.16.0
|
||||
smmap==5.0.1
|
||||
sniffio==1.3.0
|
||||
SQLAlchemy==2.0.22
|
||||
sqlparse==0.4.4
|
||||
stack-data==0.6.3
|
||||
starlette==0.35.1
|
||||
supervision==0.6.0
|
||||
svgpathtools==1.6.1
|
||||
svgwrite==1.4.3
|
||||
sympy==1.12
|
||||
tabulate==0.9.0
|
||||
termcolor==2.4.0
|
||||
threadpoolctl==3.2.0
|
||||
tifffile==2023.9.26
|
||||
timm==0.9.12
|
||||
tinycss2==1.2.1
|
||||
tokenizers==0.14.0
|
||||
tomli==2.0.1
|
||||
toolz==0.12.0
|
||||
torch==2.0.1
|
||||
torch-tools==0.1.5
|
||||
torchmetrics==1.2.0
|
||||
torchsummary==1.5.1
|
||||
torchvision==0.15.2
|
||||
tornado==6.3.3
|
||||
tqdm==4.66.1
|
||||
traitlets==5.13.0
|
||||
transformers==4.34.0
|
||||
triton==2.0.0
|
||||
typing_extensions==4.8.0
|
||||
tzdata==2023.3
|
||||
uc-micro-py==1.0.2
|
||||
urllib3==2.0.6
|
||||
uvicorn==0.25.0
|
||||
visdom==0.2.4
|
||||
wandb==0.15.11
|
||||
wcwidth==0.2.12
|
||||
webencodings==0.5.1
|
||||
websocket-client==1.6.4
|
||||
websockets==11.0.3
|
||||
widgetsnbextension==4.0.9
|
||||
xxhash==3.4.1
|
||||
yacs==0.1.8
|
||||
yapf==0.40.2
|
||||
yarl==1.9.2
|
||||
zipp==3.17.0
|
||||
+1867
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,98 @@
|
||||
{
|
||||
"norm":
|
||||
{
|
||||
"lnd_bbox": 7,
|
||||
"corners": 8,
|
||||
"pupils": 8,
|
||||
"bbox": 7,
|
||||
"height": 4
|
||||
},
|
||||
|
||||
"ldm_ids": [101, 102, 103, 104, 105, 106, 107, 108, 24, 110, 111, 112,
|
||||
113, 114, 115, 116, 117, 1, 119, 2, 121, 3, 4, 124, 5, 126,
|
||||
6, 128, 129, 130, 17, 16, 133, 134, 135, 18, 7, 138, 139, 8,
|
||||
141, 142, 11, 144, 145, 12, 147, 148, 20, 150, 151, 22, 153,
|
||||
154, 21, 156, 157, 23, 159, 160, 161, 162, 163, 164, 165, 166,
|
||||
167, 168],
|
||||
|
||||
"ldm_flip_order": [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2,
|
||||
1, 0, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27,
|
||||
28, 29, 30, 35, 34, 33, 32, 31, 45, 44, 43, 42, 47,
|
||||
46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50, 49,
|
||||
48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66,
|
||||
65],
|
||||
"test_subsets":
|
||||
{
|
||||
"indoor": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299],
|
||||
"outdoor": [300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553, 554, 555, 556, 557, 558, 559, 560, 561, 562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573, 574, 575, 576, 577, 578, 579, 580, 581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, 596, 597, 598, 599]
|
||||
},
|
||||
|
||||
"ldm_edges_matrix": [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0]]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,98 @@
|
||||
{
|
||||
"norm":
|
||||
{
|
||||
"corners": 8,
|
||||
"pupils": 8,
|
||||
"height": 4
|
||||
},
|
||||
|
||||
"ldm_ids": [101, 102, 103, 104, 105, 106, 107, 108, 24, 110, 111, 112,
|
||||
113, 114, 115, 116, 117, 1, 119, 2, 121, 3, 4, 124, 5, 126,
|
||||
6, 128, 129, 130, 17, 16, 133, 134, 135, 18, 7, 138, 139, 8,
|
||||
141, 142, 11, 144, 145, 12, 147, 148, 20, 150, 151, 22, 153,
|
||||
154, 21, 156, 157, 23, 159, 160, 161, 162, 163, 164, 165, 166,
|
||||
167, 168],
|
||||
|
||||
"ldm_flip_order": [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2,
|
||||
1, 0, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27,
|
||||
28, 29, 30, 35, 34, 33, 32, 31, 45, 44, 43, 42, 47,
|
||||
46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50, 49,
|
||||
48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66,
|
||||
65],
|
||||
"test_subsets":
|
||||
{
|
||||
"helen": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329],
|
||||
"lfpw": [330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553],
|
||||
"common": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553],
|
||||
"ibug": [554, 555, 556, 557, 558, 559, 560, 561, 562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573, 574, 575, 576, 577, 578, 579, 580, 581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, 596, 597, 598, 599, 600, 601, 602, 603, 604, 605, 606, 607, 608, 609, 610, 611, 612, 613, 614, 615, 616, 617, 618, 619, 620, 621, 622, 623, 624, 625, 626, 627, 628, 629, 630, 631, 632, 633, 634, 635, 636, 637, 638, 639, 640, 641, 642, 643, 644, 645, 646, 647, 648, 649, 650, 651, 652, 653, 654, 655, 656, 657, 658, 659, 660, 661, 662, 663, 664, 665, 666, 667, 668, 669, 670, 671, 672, 673, 674, 675, 676, 677, 678, 679, 680, 681, 682, 683, 684, 685, 686, 687, 688]
|
||||
},
|
||||
|
||||
"ldm_edges_matrix": [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0]]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,92 @@
|
||||
{
|
||||
"norm":
|
||||
{
|
||||
"lnd_bbox": 7,
|
||||
"corners": 8,
|
||||
"pupils": 8
|
||||
},
|
||||
|
||||
"ldm_ids": [101, 102, 103, 104, 105, 106, 107, 108, 24, 110, 111, 112,
|
||||
113, 114, 115, 116, 117, 1, 119, 2, 121, 3, 4, 124, 5, 126,
|
||||
6, 128, 129, 130, 17, 16, 133, 134, 135, 18, 7, 138, 139, 8,
|
||||
141, 142, 11, 144, 145, 12, 147, 148, 20, 150, 151, 22, 153,
|
||||
154, 21, 156, 157, 23, 159, 160, 161, 162, 163, 164, 165, 166,
|
||||
167, 168],
|
||||
|
||||
"ldm_flip_order": [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2,
|
||||
1, 0, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27,
|
||||
28, 29, 30, 35, 34, 33, 32, 31, 45, 44, 43, 42, 47,
|
||||
46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50, 49,
|
||||
48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66,
|
||||
65],
|
||||
"test_subsets":{},
|
||||
|
||||
"ldm_edges_matrix": [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0]]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,288 @@
|
||||
{
|
||||
"norm":
|
||||
{
|
||||
"corners": 10,
|
||||
"pupils": 8
|
||||
},
|
||||
|
||||
"ldm_ids": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
|
||||
112, 113, 114, 115, 24, 117, 118, 119, 120, 121, 122, 123,
|
||||
124, 125, 126, 127, 128, 129, 130, 131, 132, 1, 134, 2, 136,
|
||||
3, 138, 139, 140, 141, 4, 143, 5, 145, 6, 147, 148, 149, 150,
|
||||
151, 152, 153, 17, 16, 156, 157, 158, 18, 7, 161, 9, 163, 8,
|
||||
165, 10, 167, 11, 169, 13, 171, 12, 173, 14, 175, 20, 177, 178,
|
||||
22, 180, 181, 21, 183, 184, 23, 186, 187, 188, 189, 190, 191,
|
||||
192, 193, 194, 195, 196, 197],
|
||||
|
||||
"ldm_flip_order": [32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20,
|
||||
19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6,
|
||||
5, 4, 3, 2, 1, 0,
|
||||
46, 45, 44, 43, 42, 50, 49, 48, 47,
|
||||
37, 36, 35, 34, 33, 41, 40, 39, 38,
|
||||
51, 52, 53, 54,
|
||||
59, 58, 57, 56, 55,
|
||||
72, 71, 70, 69, 68, 75, 74, 73,
|
||||
64, 63, 62, 61, 60, 67, 66, 65,
|
||||
82, 81, 80, 79, 78, 77, 76,
|
||||
87, 86, 85, 84, 83,
|
||||
92, 91, 90, 89, 88,
|
||||
95, 94, 93,
|
||||
97, 96],
|
||||
|
||||
"test_subsets":
|
||||
{
|
||||
"pose": [305, 310, 312, 328, 338, 392, 395, 403, 407, 415, 438, 478, 494, 498, 511, 529, 535, 564,
|
||||
572, 577, 582, 609, 619, 650, 655, 662, 669, 671, 673, 680, 704, 733, 765, 768, 804, 826, 837, 838, 839, 844, 845,
|
||||
878, 893, 935, 942, 966, 969, 972, 1007, 1009, 1015, 1018, 1023, 1100, 1121, 1125, 1148, 1172, 1180, 1186, 1221,
|
||||
1311, 1335, 1378, 1379, 1403, 1408, 1409, 1410, 1411, 1412, 1413, 1414, 1415, 1416, 1417, 1418, 1419, 1420, 1421,
|
||||
1422, 1423, 1424, 1425, 1426, 1427, 1428, 1429, 1430, 1431, 1432, 1433, 1434, 1435, 1436, 1437, 1438, 1439, 1440,
|
||||
1441, 1442, 1443, 1444, 1445, 1446, 1447, 1448, 1449, 1450, 1451, 1452, 1453, 1454, 1455, 1456, 1457, 1458, 1459,
|
||||
1460, 1461, 1462, 1463, 1464, 1465, 1466, 1467, 1468, 1469, 1470, 1471, 1472, 1473, 1474, 1475, 1476, 1477, 1478,
|
||||
1479, 1480, 1481, 1482, 1483, 1484, 1485, 1486, 1487, 1488, 1489, 1490, 1491, 1492, 1493, 1494, 1495, 1496, 1497,
|
||||
1498, 1499, 1500, 1501, 1502, 1503, 1504, 1505, 1506, 1507, 1508, 1509, 1510, 1511, 1512, 1513, 1514, 1515, 1516,
|
||||
1517, 1518, 1519, 1520, 1521, 1522, 1523, 1524, 1525, 1526, 1527, 1528, 1529, 1530, 1531, 1532, 1533, 1534, 1535,
|
||||
1536, 1537, 1538, 1539, 1540, 1541, 1542, 1543, 1544, 1545, 1546, 1547, 1548, 1549, 1550, 1551, 1552, 1553, 1554,
|
||||
1555, 1556, 1557, 1558, 1559, 1560, 1561, 1562, 1563, 1564, 1565, 1566, 1567, 1568, 1569, 1570, 1571, 1572, 1573,
|
||||
1574, 1575, 1576, 1577, 1578, 1579, 1580, 1581, 1582, 1583, 1584, 1585, 1586, 1587, 1588, 1589, 1590, 1591, 1592,
|
||||
1593, 1594, 1595, 1596, 1597, 1598, 1599, 1600, 1601, 1602, 1603, 1604, 1605, 1606, 1607, 1608, 1609, 1610, 1611,
|
||||
1612, 1613, 1614, 1615, 1616, 1617, 1618, 1619, 1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629, 1630,
|
||||
1631, 1632, 1633, 1634, 1635, 1636, 1637, 1638, 1639, 1640, 1641, 1668, 1711, 1719, 1723, 1822, 1907, 1912, 1933,
|
||||
1962, 1964, 1974, 2010, 2043, 2089, 2094, 2108, 2170, 2181, 2187, 2220, 2317, 2335, 2389, 2391, 2459, 2489],
|
||||
"expression": [310, 329, 357, 380, 396, 400, 408, 410, 417, 450, 452, 461, 469, 472, 501, 502, 505, 506, 516, 533,
|
||||
541, 568, 585, 588, 597, 600, 624, 649, 652, 655, 701, 724, 751, 789, 817, 874, 884, 918, 1010, 1032, 1066, 1091,
|
||||
1095, 1220, 1272, 1298, 1303, 1347, 1356, 1398, 1428, 1445, 1462, 1464, 1470, 1482, 1507, 1512, 1520, 1536, 1541,
|
||||
1543, 1545, 1548, 1565, 1587, 1598, 1638, 1642, 1643, 1644, 1645, 1646, 1647, 1648, 1649, 1650, 1651, 1652, 1653,
|
||||
1654, 1655, 1656, 1657, 1658, 1659, 1660, 1661, 1662, 1663, 1664, 1665, 1666, 1667, 1668, 1669, 1670, 1671, 1672,
|
||||
1673, 1674, 1675, 1676, 1677, 1678, 1679, 1680, 1681, 1682, 1683, 1684, 1685, 1686, 1687, 1688, 1689, 1690, 1691,
|
||||
1692, 1693, 1694, 1695, 1696, 1697, 1698, 1699, 1700, 1701, 1702, 1703, 1704, 1705, 1706, 1707, 1708, 1709, 1710,
|
||||
1711, 1712, 1713, 1714, 1715, 1716, 1717, 1718, 1719, 1720, 1721, 1722, 1723, 1724, 1725, 1726, 1727, 1728, 1729,
|
||||
1730, 1731, 1732, 1733, 1734, 1735, 1736, 1737, 1738, 1739, 1740, 1741, 1742, 1743, 1744, 1745, 1746, 1747, 1748,
|
||||
1749, 1750, 1751, 1752, 1753, 1754, 1755, 1756, 1757, 1758, 1759, 1760, 1761, 1762, 1763, 1764, 1765, 1766, 1767,
|
||||
1768, 1769, 1770, 1771, 1772, 1773, 1774, 1775, 1776, 1777, 1778, 1779, 1780, 1781, 1782, 1783, 1784, 1785, 1786,
|
||||
1787, 1788, 1789, 1790, 1791, 1792, 1793, 1794, 1795, 1796, 1797, 1798, 1799, 1800, 1801, 1802, 1803, 1804, 1805,
|
||||
1806, 1807, 1808, 1809, 1810, 1811, 1812, 1813, 1814, 1815, 1816, 1817, 1818, 1819, 1820, 1821, 1822, 1823, 1824,
|
||||
1825, 1826, 1827, 1828, 1829, 1830, 1831, 1832, 1833, 1834, 1835, 1836, 1837, 1838, 1839, 1840, 1841, 1842, 1843,
|
||||
1844, 1845, 1846, 1847, 1848, 1849, 1850, 1851, 1852, 1853, 1854, 1855, 1856, 1857, 1858, 1859, 1860, 1861, 1862,
|
||||
1863, 1864, 1865, 1866, 1867, 1868, 1869, 1870, 1871, 1872, 1873, 1910, 1930, 1965, 2013, 2024, 2069, 2263, 2272,
|
||||
2276, 2315, 2360, 2375, 2406, 2412],
|
||||
"illumination": [306, 309, 312, 314, 321, 325, 327, 330, 335, 336, 338, 351, 353, 360, 361, 364, 368, 370, 372, 373,
|
||||
374, 375, 381, 384, 388, 394, 404, 405, 410, 411, 415, 417, 418, 434, 435, 438, 448, 450, 454, 458, 460, 468, 469,
|
||||
473, 476, 479, 481, 485, 487, 490, 492, 497, 503, 504, 505, 508, 512, 513, 514, 516, 517, 519, 525, 532, 536, 541,
|
||||
544, 549, 553, 555, 558, 561, 562, 568, 573, 580, 582, 588, 592, 593, 596, 607, 611, 615, 619, 622, 624, 626, 627,
|
||||
628, 629, 631, 632, 635, 636, 643, 648, 661, 666, 671, 673, 676, 680, 686, 693, 695, 697, 701, 702, 706, 710, 715,
|
||||
719, 721, 723, 725, 726, 731, 732, 734, 738, 745, 750, 765, 770, 774, 781, 784, 787, 790, 796, 804, 805, 807, 810,
|
||||
821, 832, 834, 838, 841, 842, 843, 844, 849, 853, 856, 866, 874, 887, 899, 900, 901, 902, 906, 909, 910, 911, 919,
|
||||
921, 924, 929, 931, 939, 941, 942, 957, 959, 962, 964, 965, 968, 971, 972, 974, 983, 984, 987, 990, 996, 998,
|
||||
1001, 1003, 1004, 1021, 1022, 1036, 1048, 1053, 1056, 1057, 1058, 1064, 1068, 1072, 1079, 1094, 1095, 1096, 1102,
|
||||
1104, 1105, 1109, 1111, 1112, 1113, 1114, 1115, 1116, 1117, 1118, 1119, 1120, 1121, 1122, 1123, 1124, 1125, 1126,
|
||||
1127, 1128, 1129, 1130, 1131, 1132, 1133, 1134, 1135, 1136, 1137, 1138, 1139, 1140, 1141, 1142, 1143, 1144, 1145,
|
||||
1146, 1147, 1148, 1149, 1150, 1151, 1152, 1153, 1154, 1155, 1156, 1157, 1158, 1159, 1160, 1161, 1162, 1163, 1164,
|
||||
1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1174, 1175, 1176, 1177, 1178, 1179, 1180, 1181, 1182, 1183,
|
||||
1184, 1185, 1186, 1187, 1188, 1189, 1190, 1191, 1192, 1193, 1194, 1195, 1196, 1197, 1198, 1199, 1200, 1201, 1202,
|
||||
1203, 1204, 1205, 1206, 1207, 1208, 1209, 1210, 1211, 1212, 1213, 1214, 1215, 1216, 1217, 1218, 1219, 1220, 1221,
|
||||
1222, 1223, 1224, 1225, 1226, 1227, 1228, 1229, 1230, 1231, 1232, 1233, 1234, 1235, 1236, 1237, 1238, 1239, 1240,
|
||||
1241, 1242, 1243, 1244, 1245, 1246, 1247, 1248, 1249, 1250, 1251, 1252, 1253, 1254, 1255, 1256, 1257, 1258, 1259,
|
||||
1260, 1261, 1262, 1263, 1264, 1265, 1266, 1267, 1268, 1269, 1270, 1271, 1272, 1273, 1274, 1275, 1276, 1277, 1278,
|
||||
1279, 1280, 1281, 1282, 1283, 1284, 1285, 1286, 1287, 1288, 1289, 1290, 1291, 1292, 1293, 1294, 1295, 1296, 1297,
|
||||
1298, 1299, 1300, 1301, 1302, 1303, 1304, 1305, 1306, 1307, 1308, 1309, 1310, 1311, 1312, 1313, 1314, 1315, 1316,
|
||||
1317, 1318, 1319, 1320, 1321, 1322, 1323, 1324, 1325, 1326, 1327, 1328, 1329, 1330, 1331, 1332, 1333, 1334, 1335,
|
||||
1336, 1337, 1338, 1339, 1340, 1341, 1342, 1343, 1344, 1345, 1346, 1347, 1348, 1349, 1350, 1351, 1352, 1353, 1354,
|
||||
1355, 1356, 1357, 1358, 1359, 1360, 1361, 1362, 1363, 1364, 1365, 1366, 1367, 1368, 1369, 1370, 1371, 1372, 1373,
|
||||
1374, 1375, 1376, 1377, 1378, 1379, 1380, 1381, 1382, 1383, 1384, 1385, 1386, 1387, 1388, 1389, 1390, 1391, 1392,
|
||||
1393, 1394, 1395, 1396, 1397, 1398, 1399, 1400, 1401, 1402, 1403, 1404, 1405, 1406, 1407, 1408, 1418, 1441, 1448,
|
||||
1450, 1454, 1469, 1484, 1495, 1497, 1510, 1521, 1523, 1533, 1535, 1555, 1573, 1574, 1577, 1585, 1604, 1608, 1609,
|
||||
1613, 1616, 1619, 1621, 1622, 1635, 1637, 1649, 1661, 1664, 1668, 1671, 1691, 1698, 1721, 1731, 1738, 1740, 1741,
|
||||
1745, 1760, 1767, 1768, 1783, 1787, 1788, 1793, 1796, 1798, 1801, 1813, 1826, 1835, 1836, 1849, 1850, 1851, 1852,
|
||||
1853, 1865, 1869, 1872, 1877, 1886, 1904, 1907, 1908, 1910, 1914, 1921, 1923, 1942, 1947, 1953, 1954, 1955, 1956,
|
||||
1958, 1966, 1968, 1969, 1970, 1980, 1989, 1992, 2000, 2018, 2019, 2029, 2036, 2038, 2039, 2040, 2064, 2070, 2078,
|
||||
2080, 2081, 2084, 2088, 2090, 2098, 2106, 2112, 2123, 2128, 2135, 2145, 2151, 2163, 2171, 2172, 2176, 2180, 2183,
|
||||
2184, 2186, 2192, 2196, 2201, 2204, 2208, 2209, 2213, 2215, 2218, 2223, 2226, 2236, 2240, 2242, 2245, 2246, 2253,
|
||||
2264, 2270, 2273, 2280, 2281, 2285, 2287, 2288, 2289, 2290, 2301, 2313, 2314, 2321, 2323, 2331, 2334, 2336, 2337,
|
||||
2343, 2354, 2355, 2357, 2358, 2361, 2367, 2368, 2371, 2374, 2375, 2377, 2381, 2384, 2385, 2386, 2394, 2395, 2397,
|
||||
2407, 2409, 2413, 2414, 2425, 2429, 2434, 2435, 2439, 2444, 2445, 2448, 2449, 2451, 2454, 2466, 2467, 2480, 2481,
|
||||
2484, 2487, 2489, 2491, 2497],
|
||||
"makeup": [313, 425, 464, 468, 501, 551, 614, 621, 707, 769, 776, 846, 870, 878, 900, 910, 925, 952, 960, 961, 965,
|
||||
981, 1237, 1242, 1380, 1395, 1503, 1574, 1633, 1668, 1708, 1764, 1791, 1837, 1860, 1874, 1875, 1876, 1877, 1878,
|
||||
1879, 1880, 1881, 1882, 1883, 1884, 1885, 1886, 1887, 1888, 1889, 1890, 1891, 1892, 1893, 1894, 1895, 1896, 1897,
|
||||
1898, 1899, 1900, 1901, 1902, 1903, 1904, 1905, 1906, 1907, 1908, 1909, 1910, 1911, 1912, 1913, 1914, 1915, 1916,
|
||||
1917, 1918, 1919, 1920, 1921, 1922, 1923, 1924, 1925, 1926, 1927, 1928, 1929, 1930, 1931, 1932, 1933, 1934, 1935,
|
||||
1936, 1937, 1938, 1939, 1940, 1941, 1942, 1943, 1944, 1945, 1946, 1947, 1948, 1949, 1950, 1951, 1952, 1953, 1954,
|
||||
1955, 1956, 1957, 1958, 1959, 1960, 1961, 1962, 1963, 1964, 1965, 1966, 1967, 1968, 1969, 1970, 1971, 1972, 1973,
|
||||
1974, 1975, 1976, 1977, 1978, 1979, 1980, 1981, 1982, 1983, 1984, 1985, 1986, 1987, 1988, 1989, 1990, 1991, 1992,
|
||||
1993, 1994, 1995, 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011,
|
||||
2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022, 2023, 2024, 2025, 2026, 2027, 2028, 2029, 2030,
|
||||
2031, 2032, 2033, 2034, 2035, 2036, 2037, 2038, 2265, 2290, 2293, 2381, 2486, 2490],
|
||||
"occlusion": [300, 301, 310, 328, 329, 332, 333, 335, 336, 338, 339, 342, 343, 356, 362, 367, 370, 374, 377, 384,
|
||||
401, 406, 408, 417, 421, 423, 426, 427, 434, 435, 440, 453, 457, 461, 465, 466, 472, 475, 476, 479, 489, 497, 500,
|
||||
504, 506, 514, 515, 518, 520, 524, 525, 536, 543, 546, 558, 559, 563, 564, 567, 568, 574, 577, 584, 589, 590, 592,
|
||||
594, 602, 605, 609, 610, 614, 618, 627, 638, 640, 643, 645, 646, 652, 658, 661, 663, 666, 672, 674, 677, 682, 685,
|
||||
700, 701, 702, 703, 704, 705, 706, 707, 708, 709, 710, 711, 712, 713, 714, 715, 716, 717, 718, 719, 720, 721, 722,
|
||||
723, 724, 725, 726, 727, 728, 729, 730, 731, 732, 733, 734, 735, 736, 737, 738, 739, 740, 741, 742, 743, 744, 745,
|
||||
746, 747, 748, 749, 750, 751, 752, 753, 754, 755, 756, 757, 758, 759, 760, 761, 762, 763, 764, 765, 766, 767, 768,
|
||||
769, 770, 771, 772, 773, 774, 775, 776, 777, 778, 779, 780, 781, 782, 783, 784, 785, 786, 787, 788, 789, 790, 791,
|
||||
792, 793, 794, 795, 796, 797, 798, 799, 800, 801, 802, 803, 804, 805, 806, 807, 808, 809, 810, 811, 812, 813, 814,
|
||||
815, 816, 817, 818, 819, 820, 821, 822, 823, 824, 825, 826, 827, 828, 829, 830, 831, 832, 833, 834, 835, 836, 837,
|
||||
838, 839, 840, 841, 842, 843, 844, 845, 846, 847, 848, 849, 850, 851, 852, 853, 854, 855, 856, 857, 858, 859, 860,
|
||||
861, 862, 863, 864, 865, 866, 867, 868, 869, 870, 871, 872, 873, 874, 875, 876, 877, 878, 879, 880, 881, 882, 883,
|
||||
884, 885, 886, 887, 888, 889, 890, 891, 892, 893, 894, 895, 896, 897, 898, 899, 900, 901, 902, 903, 904, 905, 906,
|
||||
907, 908, 909, 910, 911, 912, 913, 914, 915, 916, 917, 918, 919, 920, 921, 922, 923, 924, 925, 926, 927, 928, 929,
|
||||
930, 931, 932, 933, 934, 935, 936, 937, 938, 939, 940, 941, 942, 943, 944, 945, 946, 947, 948, 949, 950, 951, 952,
|
||||
953, 954, 955, 956, 957, 958, 959, 960, 961, 962, 963, 964, 965, 966, 967, 968, 969, 970, 971, 972, 973, 974, 975,
|
||||
976, 977, 978, 979, 980, 981, 982, 983, 984, 985, 986, 987, 988, 989, 990, 991, 992, 993, 994, 995, 996, 997, 998,
|
||||
999, 1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014, 1015, 1016, 1017,
|
||||
1018, 1019, 1020, 1021, 1022, 1023, 1024, 1025, 1026, 1027, 1028, 1029, 1030, 1031, 1032, 1033, 1034, 1035, 1036,
|
||||
1037, 1038, 1039, 1040, 1041, 1042, 1043, 1044, 1045, 1046, 1047, 1048, 1049, 1050, 1051, 1052, 1053, 1054, 1055,
|
||||
1056, 1057, 1058, 1059, 1060, 1061, 1062, 1063, 1064, 1065, 1066, 1067, 1068, 1069, 1070, 1071, 1072, 1073, 1074,
|
||||
1075, 1076, 1077, 1078, 1079, 1080, 1081, 1082, 1083, 1084, 1085, 1086, 1087, 1088, 1089, 1090, 1091, 1092, 1093,
|
||||
1094, 1095, 1096, 1097, 1098, 1099, 1100, 1101, 1102, 1103, 1104, 1105, 1106, 1107, 1108, 1109, 1110, 1112, 1121,
|
||||
1127, 1129, 1135, 1136, 1143, 1146, 1147, 1148, 1151, 1160, 1161, 1166, 1170, 1177, 1180, 1190, 1191, 1195, 1201,
|
||||
1218, 1223, 1231, 1233, 1239, 1240, 1241, 1245, 1248, 1260, 1275, 1277, 1278, 1280, 1283, 1292, 1298, 1299, 1313,
|
||||
1320, 1328, 1333, 1338, 1339, 1342, 1347, 1355, 1363, 1368, 1373, 1377, 1382, 1383, 1384, 1389, 1390, 1391, 1399,
|
||||
1401, 1404, 1405, 1412, 1415, 1422, 1426, 1427, 1431, 1438, 1450, 1479, 1485, 1489, 1493, 1497, 1499, 1506, 1516,
|
||||
1524, 1526, 1531, 1533, 1534, 1547, 1555, 1560, 1561, 1592, 1595, 1603, 1616, 1622, 1625, 1632, 1647, 1650, 1651,
|
||||
1655, 1657, 1658, 1661, 1665, 1667, 1672, 1683, 1684, 1694, 1702, 1708, 1711, 1713, 1722, 1723, 1728, 1734, 1748,
|
||||
1750, 1752, 1753, 1757, 1764, 1775, 1793, 1794, 1801, 1809, 1813, 1820, 1832, 1838, 1841, 1856, 1864, 1866, 1869,
|
||||
1878, 1882, 1895, 1898, 1921, 1930, 1931, 1936, 1941, 1946, 1947, 1953, 1975, 1979, 1982, 1983, 1985, 1989, 1999,
|
||||
2013, 2021, 2026, 2042, 2044, 2047, 2048, 2059, 2069, 2077, 2083, 2085, 2086, 2087, 2094, 2110, 2112, 2122, 2128,
|
||||
2131, 2134, 2136, 2139, 2144, 2151, 2155, 2163, 2175, 2190, 2197, 2199, 2208, 2209, 2211, 2213, 2238, 2246, 2279,
|
||||
2301, 2305, 2316, 2321, 2324, 2328, 2329, 2331, 2335, 2339, 2340, 2348, 2351, 2354, 2364, 2367, 2370, 2371, 2372,
|
||||
2376, 2379, 2381, 2391, 2402, 2409, 2410, 2411, 2412, 2413, 2430, 2433, 2435, 2449, 2450, 2466, 2471, 2473, 2476,
|
||||
2478, 2484, 2485, 2487, 2491, 2494],
|
||||
"blur": [300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320,
|
||||
321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343,
|
||||
344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366,
|
||||
367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389,
|
||||
390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412,
|
||||
413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435,
|
||||
436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458,
|
||||
459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481,
|
||||
482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504,
|
||||
505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527,
|
||||
528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550,
|
||||
551, 552, 553, 554, 555, 556, 557, 558, 559, 560, 561, 562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573,
|
||||
574, 575, 576, 577, 578, 579, 580, 581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, 596,
|
||||
597, 598, 599, 600, 601, 602, 603, 604, 605, 606, 607, 608, 609, 610, 611, 612, 613, 614, 615, 616, 617, 618, 619,
|
||||
620, 621, 622, 623, 624, 625, 626, 627, 628, 629, 630, 631, 632, 633, 634, 635, 636, 637, 638, 639, 640, 641, 642,
|
||||
643, 644, 645, 646, 647, 648, 649, 650, 651, 652, 653, 654, 655, 656, 657, 658, 659, 660, 661, 662, 663, 664, 665,
|
||||
666, 667, 668, 669, 670, 671, 672, 673, 674, 675, 676, 677, 678, 679, 680, 681, 682, 683, 684, 685, 686, 687, 688,
|
||||
689, 690, 691, 692, 693, 694, 695, 696, 697, 698, 699, 702, 705, 714, 725, 731, 738, 740, 742, 745, 746, 748, 752,
|
||||
753, 754, 756, 760, 765, 766, 773, 780, 784, 785, 787, 796, 798, 799, 800, 804, 808, 809, 812, 813, 816, 817, 823,
|
||||
824, 827, 832, 834, 838, 865, 872, 874, 879, 882, 889, 891, 899, 901, 908, 911, 918, 924, 929, 931, 932, 938, 941,
|
||||
951, 953, 954, 959, 960, 965, 968, 969, 971, 977, 985, 989, 991, 995, 999, 1001, 1004, 1018, 1021, 1026, 1029,
|
||||
1031, 1033, 1040, 1050, 1056, 1060, 1061, 1069, 1074, 1076, 1084, 1085, 1099, 1103, 1109, 1115, 1117, 1118, 1120,
|
||||
1124, 1127, 1133, 1142, 1146, 1148, 1149, 1151, 1171, 1177, 1178, 1181, 1184, 1185, 1187, 1191, 1193, 1194, 1201,
|
||||
1202, 1206, 1217, 1224, 1226, 1227, 1231, 1232, 1236, 1237, 1246, 1250, 1254, 1255, 1258, 1260, 1262, 1263, 1266,
|
||||
1269, 1270, 1278, 1280, 1283, 1286, 1288, 1293, 1300, 1301, 1315, 1317, 1320, 1324, 1325, 1333, 1363, 1365, 1370,
|
||||
1373, 1380, 1382, 1386, 1392, 1396, 1400, 1402, 1403, 1413, 1421, 1422, 1424, 1434, 1442, 1450, 1451, 1459, 1464,
|
||||
1468, 1475, 1481, 1483, 1495, 1506, 1510, 1518, 1519, 1522, 1546, 1554, 1559, 1570, 1573, 1577, 1578, 1589, 1593,
|
||||
1601, 1603, 1608, 1609, 1612, 1613, 1614, 1616, 1618, 1619, 1633, 1642, 1644, 1647, 1649, 1652, 1653, 1658, 1660,
|
||||
1663, 1664, 1670, 1672, 1673, 1689, 1691, 1695, 1705, 1710, 1712, 1714, 1716, 1718, 1720, 1727, 1731, 1735, 1737,
|
||||
1739, 1742, 1750, 1761, 1765, 1768, 1770, 1776, 1779, 1782, 1785, 1788, 1796, 1811, 1813, 1815, 1826, 1828, 1831,
|
||||
1838, 1853, 1854, 1860, 1864, 1875, 1878, 1883, 1885, 1887, 1891, 1915, 1919, 1921, 1923, 1926, 1936, 1940, 1941,
|
||||
1942, 1951, 1954, 1957, 1975, 1978, 1981, 1997, 1998, 2005, 2015, 2021, 2028, 2031, 2036, 2040, 2042, 2046, 2047,
|
||||
2054, 2060, 2066, 2069, 2076, 2077, 2078, 2083, 2085, 2090, 2093, 2099, 2106, 2112, 2117, 2132, 2135, 2136, 2148,
|
||||
2151, 2158, 2162, 2163, 2167, 2169, 2172, 2173, 2184, 2186, 2194, 2201, 2203, 2225, 2229, 2235, 2237, 2242, 2259,
|
||||
2260, 2269, 2270, 2279, 2281, 2284, 2290, 2298, 2301, 2305, 2309, 2311, 2312, 2313, 2317, 2318, 2325, 2328, 2333,
|
||||
2335, 2336, 2337, 2341, 2344, 2357, 2371, 2372, 2374, 2375, 2380, 2390, 2396, 2397, 2410, 2418, 2419, 2423, 2428,
|
||||
2444, 2457, 2458, 2464, 2472, 2473, 2476, 2488, 2489]
|
||||
},
|
||||
"ldm_edges_matrix": [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,158 @@
|
||||
import os
|
||||
import json
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
|
||||
from .transforms import get_transformers
|
||||
|
||||
|
||||
class AlignmentsDataset(Dataset):
|
||||
'''Loads datasets of images with landmarks and bounding boxes.
|
||||
'''
|
||||
|
||||
def __init__(self,
|
||||
database,
|
||||
json_file,
|
||||
images_dir,
|
||||
image_size=(128, 128),
|
||||
transform=None,
|
||||
indices=None,
|
||||
debug=False):
|
||||
"""
|
||||
|
||||
:param database: class DatabaseStruct containing all the specifics of the database
|
||||
|
||||
:param json_file: path to the json file which contains the names of the images, landmarks, bounding boxes, etc
|
||||
|
||||
:param images_dir: path of the directory containing the images.
|
||||
|
||||
:param image_size: tuple like e.g. (128, 128)
|
||||
|
||||
:param transform: composition of transformations that will be applied to the samples.
|
||||
|
||||
:param debug_mode: bool if True, loads a very reduced_version of the dataset for debugging purposes.
|
||||
|
||||
:param indices: If it is a list of indices, allows to work with the subset of
|
||||
items specified by the list. If it is None, the whole set is used.
|
||||
"""
|
||||
|
||||
self.database = database
|
||||
self.images_dir = images_dir
|
||||
self.transform = transform
|
||||
self.image_size = image_size
|
||||
self.indices = indices
|
||||
self._imgs_dict = None
|
||||
self.debug = debug
|
||||
|
||||
with open(json_file) as jsonfile:
|
||||
self.data = json.load(jsonfile)
|
||||
|
||||
def __len__(self):
|
||||
'''Returns the length of the dataset
|
||||
'''
|
||||
if self.indices is None:
|
||||
return len(self.data)
|
||||
else:
|
||||
return len(self.indices)
|
||||
|
||||
def __getitem__(self, sample_idx):
|
||||
'''Returns sample of the dataset of index idx'''
|
||||
|
||||
# To allow work with a subset
|
||||
if self.indices is not None:
|
||||
sample_idx = self.indices[sample_idx]
|
||||
|
||||
# Load sample image
|
||||
img_name = os.path.join(self.images_dir, self.data[sample_idx]['imgpath'])
|
||||
if not self._imgs_dict:
|
||||
image_cv = cv2.imread(img_name)
|
||||
else:
|
||||
image_cv = self._imgs_dict[sample_idx]
|
||||
|
||||
# Some images are B&W. We make sure that any image has three channels.
|
||||
if len(image_cv.shape) == 2:
|
||||
image_cv = np.repeat(image_cv[:, :, np.newaxis], 3, axis=-1)
|
||||
|
||||
# Some images have alpha channel
|
||||
image_cv = image_cv[:, :, :3]
|
||||
|
||||
image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
|
||||
image = Image.fromarray(image_cv)
|
||||
|
||||
# Load sample anns
|
||||
ids = np.array(self.data[sample_idx]['ids'])
|
||||
landmarks = np.array(self.data[sample_idx]['landmarks'])
|
||||
bbox = np.array(self.data[sample_idx]['bbox'])
|
||||
vis = np.array(self.data[sample_idx]['visible'])
|
||||
headpose = self.data[sample_idx]['headpose']
|
||||
|
||||
# Generate bbox if need it
|
||||
if bbox is None:
|
||||
# Compute bbox using landmarks
|
||||
aux = landmarks[vis == 1.0]
|
||||
bbox = np.zeros(4)
|
||||
bbox[0] = min(aux[:, 0])
|
||||
bbox[1] = min(aux[:, 1])
|
||||
bbox[2] = max(aux[:, 0]) - bbox[0]
|
||||
bbox[3] = max(aux[:, 1]) - bbox[1]
|
||||
|
||||
# Clean and mask landmarks
|
||||
mask_ldm = np.ones(self.database.num_landmarks)
|
||||
if not self.database.ldm_ids == ids.tolist():
|
||||
new_ldm = np.zeros((self.database.num_landmarks, 2))
|
||||
new_vis = np.zeros(self.database.num_landmarks)
|
||||
xyv = np.hstack((landmarks, vis[np.newaxis,:].T))
|
||||
ids_dict = dict(zip(ids.astype(int).astype(str), xyv))
|
||||
|
||||
for pos, identifier in enumerate(self.database.ldm_ids):
|
||||
if str(identifier) in ids_dict:
|
||||
x, y, v = ids_dict[str(identifier)]
|
||||
new_ldm[pos] = [x,y]
|
||||
new_vis[pos] = v
|
||||
else:
|
||||
mask_ldm[pos] = 0
|
||||
landmarks = new_ldm
|
||||
vis = new_vis
|
||||
|
||||
sample = {'image': image,
|
||||
'sample_idx': sample_idx,
|
||||
'imgpath': img_name,
|
||||
'ids_ldm': np.array(self.database.ldm_ids),
|
||||
'bbox': bbox,
|
||||
'bbox_raw': bbox,
|
||||
'landmarks': landmarks,
|
||||
'visible': vis.astype(np.float64),
|
||||
'mask_ldm': mask_ldm,
|
||||
'imgpath_local': self.data[sample_idx]['imgpath'],
|
||||
}
|
||||
|
||||
if self.debug:
|
||||
sample['landmarks_ori'] = landmarks
|
||||
sample['visible_ori'] = vis.astype(np.float64)
|
||||
sample['mask_ldm_ori'] = mask_ldm
|
||||
if headpose is not None:
|
||||
sample['headpose_ori'] = np.array(headpose)
|
||||
|
||||
if self.transform:
|
||||
sample = self.transform(sample)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
def get_dataset(data_config, pretreat=None, debug=False):
|
||||
|
||||
augmentors = get_transformers(data_config)
|
||||
if pretreat is not None:
|
||||
augmentors.append(pretreat)
|
||||
|
||||
dataset = AlignmentsDataset(data_config.database,
|
||||
data_config.anns_file,
|
||||
data_config.image_dir,
|
||||
image_size=data_config.image_size,
|
||||
transform=transforms.Compose(augmentors),
|
||||
indices=data_config.ids,
|
||||
debug=debug)
|
||||
return dataset
|
||||
@@ -0,0 +1,122 @@
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
import cv2
|
||||
|
||||
|
||||
class AddBoundary(object):
|
||||
def __init__(self, num_landmarks=68, map_size=64, sigma=1, min_dpi=64):
|
||||
self.num_landmarks = num_landmarks
|
||||
self.sigma = sigma
|
||||
|
||||
if isinstance(map_size, (tuple, list)):
|
||||
self.width = map_size[0]
|
||||
self.height = map_size[1]
|
||||
else:
|
||||
self.width = map_size
|
||||
self.height = map_size
|
||||
|
||||
if max(map_size) > min_dpi:
|
||||
self.dpi = max(map_size)
|
||||
else:
|
||||
self.dpi = min_dpi
|
||||
|
||||
self.fig_size =[self.height/self.dpi, self.width/self.dpi]
|
||||
|
||||
def __call__(self, sample):
|
||||
landmarks = sample['landmarks_float']
|
||||
mask_lnd = sample['mask_ldm_float']
|
||||
boundaries = self.get_dataset_boundaries(landmarks, mask_lnd)
|
||||
functions = {}
|
||||
|
||||
for key, points in boundaries.items():
|
||||
if len(points) != 0:
|
||||
temp = points[0]
|
||||
new_points = points[0:1, :]
|
||||
for point in points[1:]:
|
||||
if point[0] == temp[0] and point[1] == temp[1]:
|
||||
continue
|
||||
else:
|
||||
new_points = np.concatenate((new_points, np.expand_dims(point, 0)), axis=0)
|
||||
temp = point
|
||||
|
||||
points = new_points
|
||||
if points.shape[0] == 1:
|
||||
points = np.concatenate((points, points+0.001), axis=0)
|
||||
k = min(4, points.shape[0])
|
||||
functions[key] = interpolate.splprep([points[:, 0], points[:, 1]], k=k-1,s=0)
|
||||
|
||||
boundary_maps = np.zeros((len(boundaries), self.height, self.width))
|
||||
for i_map, key in enumerate(functions.keys()):
|
||||
boundary_map = np.zeros((self.height, self.width))
|
||||
xnew = np.arange(0, 1, 1/self.dpi)
|
||||
out = interpolate.splev(xnew, functions[key][0], der=0)
|
||||
|
||||
out = np.round(out).astype(int).transpose()
|
||||
out = out[out[:, 0] < self.height]
|
||||
out = out[out[:, 1] < self.width]
|
||||
boundary_map[out[:,1], out[:,0]]= 255
|
||||
|
||||
# Smooth
|
||||
sigma = self.sigma
|
||||
temp = 255 - boundary_map.astype(np.uint8)
|
||||
temp = cv2.distanceTransform(temp, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
|
||||
temp = temp.astype(np.float32)
|
||||
temp = np.where(temp < 3*sigma, np.exp(-(temp*temp)/(2*sigma*sigma)), 0 )
|
||||
boundary_maps[i_map] = temp
|
||||
|
||||
sample['boundary'] = boundary_maps
|
||||
return sample
|
||||
|
||||
def get_dataset_boundaries(self, landmarks, mask_lnd):
|
||||
boundaries = {}
|
||||
if self.num_landmarks == 68:
|
||||
cheek = landmarks[0:17]
|
||||
boundaries['cheek'] = cheek[mask_lnd[0:17] > 0]
|
||||
left_eyebrow = landmarks[17:22]
|
||||
boundaries['left_eyebrow'] = left_eyebrow[mask_lnd[17:22] > 0]
|
||||
right_eyebrow = landmarks[22:27]
|
||||
boundaries['right_eyebrow'] = right_eyebrow[mask_lnd[22:27] > 0]
|
||||
nose = landmarks[27:31]
|
||||
boundaries['nose'] = nose[mask_lnd[27:31] > 0]
|
||||
nose_bot = landmarks[31:36]
|
||||
boundaries['nose_bot'] = nose_bot[mask_lnd[31:36] > 0]
|
||||
uper_left_eyelid = landmarks[36:40]
|
||||
boundaries['upper_left_eyelid'] = uper_left_eyelid[mask_lnd[36:40] > 0]
|
||||
lower_left_eyelid = np.array([landmarks[i] for i in [36, 41, 40, 39]])
|
||||
lower_left_eyelid_mask = np.array([mask_lnd[i] for i in [36, 41, 40, 39]])
|
||||
boundaries['lower_left_eyelid'] = lower_left_eyelid[lower_left_eyelid_mask > 0]
|
||||
upper_right_eyelid = landmarks[42:46]
|
||||
boundaries['upper_right_eyelid'] = upper_right_eyelid[mask_lnd[42:46] > 0]
|
||||
lower_right_eyelid = np.array([landmarks[i] for i in [42, 47, 46, 45]])
|
||||
lower_right_eyelid_mask = np.array([mask_lnd[i] for i in [42, 47, 46, 45]])
|
||||
boundaries['lower_right_eyelid'] = lower_right_eyelid[lower_right_eyelid_mask > 0]
|
||||
upper_outer_lip = landmarks[48:55]
|
||||
boundaries['upper_outer_lip'] = upper_outer_lip[mask_lnd[48:55] > 0]
|
||||
lower_outer_lip = np.array([landmarks[i] for i in [48, 59, 58, 57, 56, 55, 54]])
|
||||
lower_outer_lip_mask = np.array([mask_lnd[i] for i in [48, 59, 58, 57, 56, 55, 54]])
|
||||
boundaries['lower_outer_lip'] = lower_outer_lip[lower_outer_lip_mask > 0]
|
||||
upper_inner_lip = np.array([landmarks[i] for i in [60, 61, 62, 63, 64]])
|
||||
upper_inner_lip_mask = np.array([mask_lnd[i] for i in [60, 61, 62, 63, 64]])
|
||||
boundaries['upper_inner_lip'] = upper_inner_lip[upper_inner_lip_mask > 0]
|
||||
lower_inner_lip = np.array([landmarks[i] for i in [60, 67, 66, 65, 64]])
|
||||
lower_inner_lip_mask = np.array([mask_lnd[i] for i in [60, 67, 66, 65, 64]])
|
||||
boundaries['lower_inner_lip'] = lower_inner_lip[lower_inner_lip_mask > 0]
|
||||
|
||||
elif self.num_landmarks == 98:
|
||||
boundaries['cheek'] = landmarks[0:33]
|
||||
boundaries['upper_left_eyebrow'] = landmarks[33:38]
|
||||
boundaries['lower_left_eyebrow'] = np.array([landmarks[i] for i in [33, 41, 40, 39, 38]])
|
||||
boundaries['upper_right_eyebrow'] = landmarks[42:47]
|
||||
boundaries['lower_right_eyebrow'] = landmarks[46:51]
|
||||
boundaries['nose'] = landmarks[51:55]
|
||||
boundaries['nose_bot'] = landmarks[55:60]
|
||||
boundaries['upper_left_eyelid'] = landmarks[60:65]
|
||||
boundaries['lower_left_eyelid'] = np.array([landmarks[i] for i in [60, 67, 66, 65, 64]])
|
||||
boundaries['upper_right_eyelid'] = landmarks[68:73]
|
||||
boundaries['lower_right_eyelid'] = np.array([landmarks[i] for i in [68, 75, 74, 73, 72]])
|
||||
boundaries['upper_outer_lip'] = landmarks[76:83]
|
||||
boundaries['lower_outer_lip'] = np.array([landmarks[i] for i in [76, 87, 86, 85, 84, 83, 82]])
|
||||
boundaries['upper_inner_lip'] = np.array([landmarks[i] for i in [88, 89, 90, 91, 92]])
|
||||
boundaries['lower_inner_lip'] = np.array([landmarks[i] for i in [88, 95, 94, 93, 92]])
|
||||
|
||||
return boundaries
|
||||
@@ -0,0 +1,39 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Heatmaps:
|
||||
|
||||
def __init__(self, num_maps, map_size, sigma, stride=1, norm=True):
|
||||
self.num_maps = num_maps
|
||||
self.sigma = sigma
|
||||
self.double_sigma_pw2 = 2*sigma*sigma
|
||||
self.doublepi_sigma_pw2 = self.double_sigma_pw2 * np.pi
|
||||
self.stride = stride
|
||||
self.norm = norm
|
||||
|
||||
if isinstance(map_size, (tuple, list)):
|
||||
self.width = map_size[0]
|
||||
self.height = map_size[1]
|
||||
else:
|
||||
self.width = map_size
|
||||
self.height = map_size
|
||||
|
||||
grid_x = np.arange(self.width) * stride + stride / 2 - 0.5
|
||||
self.grid_x = np.repeat(grid_x.reshape(1, self.width), self.num_maps, axis=0)
|
||||
grid_y = np.arange(self.height) * stride + stride / 2 - 0.5
|
||||
self.grid_y = np.repeat(grid_y.reshape(1, self.height), self.num_maps, axis=0)
|
||||
|
||||
def __call__(self, sample):
|
||||
landmarks = sample['landmarks']
|
||||
landmarks = landmarks[-self.num_maps:]
|
||||
|
||||
# Heatmap generation
|
||||
exp_x = np.exp(-(self.grid_x - landmarks[:, 0].reshape(-1, 1)) ** 2 / self.double_sigma_pw2)
|
||||
exp_y = np.exp(-(self.grid_y - landmarks[:, 1].reshape(-1, 1)) ** 2 / self.double_sigma_pw2)
|
||||
heatmaps = np.matmul(exp_y.reshape(self.num_maps, self.height, 1), exp_x.reshape(self.num_maps, 1, self.width))
|
||||
|
||||
if self.norm:
|
||||
heatmaps = heatmaps/self.doublepi_sigma_pw2
|
||||
|
||||
sample['heatmap2D'] = heatmaps
|
||||
return sample
|
||||
@@ -0,0 +1,307 @@
|
||||
import random
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
# My libs
|
||||
from. import utils as dlu
|
||||
|
||||
|
||||
class HorizontalFlipAug:
|
||||
def __init__(self, ldm_flip_order, prob=0.5):
|
||||
self.prob = prob
|
||||
self.ldm_flip_order = ldm_flip_order
|
||||
|
||||
def __call__(self, sample):
|
||||
img = sample['image']
|
||||
landmarks = sample['landmarks']
|
||||
mask = sample['mask_ldm']
|
||||
vis = sample['visible']
|
||||
bbox = sample['bbox']
|
||||
|
||||
if random.random() < self.prob:
|
||||
new_img = transforms.functional.hflip(img)
|
||||
|
||||
lm_new_order = self.ldm_flip_order
|
||||
new_landmarks = landmarks[lm_new_order]
|
||||
new_landmarks = (new_landmarks - (img.size[0], 0)) * (-1, 1)
|
||||
new_mask = mask[lm_new_order]
|
||||
new_vis = vis[lm_new_order]
|
||||
|
||||
x, y, w, h = bbox
|
||||
new_x = img.size[0] - x - w
|
||||
new_bbox = np.array((new_x, y, w, h))
|
||||
|
||||
sample['image'] = new_img
|
||||
sample['landmarks'] = new_landmarks
|
||||
sample['mask_ldm'] = new_mask
|
||||
sample['visible'] = new_vis
|
||||
sample['bbox'] = new_bbox
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class GeometryBaseAug:
|
||||
|
||||
def __call__(self, sample):
|
||||
raise NotImplementedError('Inheritance __call__ not defined')
|
||||
|
||||
def map_affine_transformation(self, sample, affine_transf, new_size=None):
|
||||
sample['image'] = self._image_affine_trans(sample['image'], affine_transf, new_size)
|
||||
sample['bbox'] = self._bbox_affine_trans(sample['bbox'], affine_transf)
|
||||
if 'landmarks' in sample.keys():
|
||||
sample['landmarks'] = self._landmarks_affine_trans(sample['landmarks'], affine_transf)
|
||||
return sample
|
||||
|
||||
def clean_outbbox_landmarks(self, shape, landmarks, mask):
|
||||
filter_x1 = landmarks[:, 0] >= shape[0]
|
||||
filter_x2 = landmarks[:, 0] < (shape[0] + shape[2])
|
||||
filter_x = np.logical_and(filter_x1,filter_x2)
|
||||
|
||||
filter_y1 = landmarks[:, 1] >= shape[1]
|
||||
filter_y2 = landmarks[:, 1] < (shape[1] + shape[3])
|
||||
filter_y = np.logical_and(filter_y1, filter_y2)
|
||||
|
||||
filter_bbox = np.logical_and(filter_x, filter_y)
|
||||
new_mask = mask*filter_bbox
|
||||
new_landmarks = (landmarks.T * new_mask).T
|
||||
new_landmarks = new_landmarks.astype(int).astype(float)
|
||||
return new_mask, new_landmarks
|
||||
|
||||
def _image_affine_trans(self, image, affine_transf, new_size=None):
|
||||
|
||||
if not new_size:
|
||||
new_size = image.size
|
||||
|
||||
inv_affine_transf = dlu.get_inverse_transf(affine_transf)
|
||||
new_image = image.transform(new_size, Image.AFFINE, inv_affine_transf.flatten())
|
||||
return new_image
|
||||
|
||||
def _bbox_affine_trans(self, bbox, affine_transf):
|
||||
|
||||
x, y, w, h = bbox
|
||||
images_bb = []
|
||||
for point in ([x, y, 1], [x + w, y, 1],
|
||||
[x, y + h, 1], [x + w, y + h, 1]):
|
||||
images_bb.append(affine_transf.dot(point))
|
||||
images_bb = np.array(images_bb)
|
||||
|
||||
new_corner0 = np.min(images_bb, axis=0)
|
||||
new_corner1 = np.max(images_bb, axis=0)
|
||||
|
||||
new_x, new_y = new_corner0
|
||||
new_w, new_h = new_corner1 - new_corner0
|
||||
new_bbox = np.array((new_x, new_y, new_w, new_h))
|
||||
return new_bbox
|
||||
|
||||
def _landmarks_affine_trans(self, landmarks, affine_transf):
|
||||
|
||||
homog_landmarks = dlu.affine2homogeneous(landmarks)
|
||||
new_landmarks = affine_transf.dot(homog_landmarks.T).T
|
||||
return new_landmarks
|
||||
|
||||
|
||||
class RSTAug(GeometryBaseAug):
|
||||
|
||||
def __init__(self, angle_range=45., scale_min=-0.15, scale_max=0.15, trl_ratio=0.05):
|
||||
self.scale_max = scale_max
|
||||
self.scale_min = scale_min
|
||||
self.angle_range = angle_range
|
||||
self.trl_ratio = trl_ratio
|
||||
|
||||
def __call__(self, sample):
|
||||
x, y, w, h = sample['bbox']
|
||||
|
||||
x0, y0 = x + w/2, y + h/2 # center of the face, which will be the center of the rotation
|
||||
|
||||
# Bbox translation
|
||||
rnd_Tx = np.random.uniform(-self.trl_ratio, self.trl_ratio) * w
|
||||
rnd_Ty = np.random.uniform(-self.trl_ratio, self.trl_ratio) * h
|
||||
sample['bbox'][0] += rnd_Tx
|
||||
sample['bbox'][1] += rnd_Ty
|
||||
|
||||
scale = 1 + np.random.uniform(self.scale_min, self.scale_max)
|
||||
angle = np.random.uniform(-self.angle_range, self.angle_range)
|
||||
|
||||
similarity = dlu.get_similarity_matrix(angle, scale, center=(x0, y0))
|
||||
new_sample = self.map_affine_transformation(sample, similarity)
|
||||
return new_sample
|
||||
|
||||
|
||||
class TargetCropAug(GeometryBaseAug):
|
||||
def __init__(self, img_new_size=128, map_new_size=128, target_dist=1.3):
|
||||
|
||||
self.target_dist = target_dist
|
||||
self.new_size_x, self.new_size_y = self._convert_shapes(img_new_size)
|
||||
self.map_size_x, self.map_size_y = self._convert_shapes(map_new_size)
|
||||
self.img2map_scale = False
|
||||
|
||||
# Mismatch between img shape and featuremap shape
|
||||
if self.map_size_x != self.new_size_x or self.map_size_y != self.new_size_y:
|
||||
self.img2map_scale = True
|
||||
self.map_scale_x = self.map_size_x / self.new_size_x
|
||||
self.map_scale_y = self.map_size_y / self.new_size_y
|
||||
self.map_scale_xx = self.map_scale_x * self.map_scale_x
|
||||
self.map_scale_xy = self.map_scale_x * self.map_scale_y
|
||||
self.map_scale_yy = self.map_scale_y * self.map_scale_y
|
||||
|
||||
def _convert_shapes(self, new_size):
|
||||
if isinstance(new_size, (tuple, list)):
|
||||
new_size_x = new_size[0]
|
||||
new_size_y = new_size[1]
|
||||
else:
|
||||
new_size_x = new_size
|
||||
new_size_y = new_size
|
||||
return new_size_x, new_size_y
|
||||
|
||||
def __call__(self, sample):
|
||||
x, y, w, h = sample['bbox']
|
||||
# we enlarge the area taken around the bounding box
|
||||
# it is neccesary to change the botton left point of the bounding box
|
||||
# according to the previous enlargement. Note this will NOT be the new
|
||||
# bounding box!
|
||||
# We return square images, which is neccesary since
|
||||
# all the images must have the same size in order to form batches
|
||||
side = max(w, h) * self.target_dist
|
||||
x -= (side - w) / 2
|
||||
y -= (side - h) / 2
|
||||
|
||||
# center of the enlarged bounding box
|
||||
x0, y0 = x + side/2, y + side/2
|
||||
# homothety factor, chosen so the new horizontal dimension will
|
||||
# coincide with new_size
|
||||
mu_x = self.new_size_x / side
|
||||
mu_y = self.new_size_y / side
|
||||
|
||||
# new_w, new_h = new_size, int(h * mu)
|
||||
new_w = self.new_size_x
|
||||
new_h = self.new_size_y
|
||||
new_x0, new_y0 = new_w / 2, new_h / 2
|
||||
|
||||
# dilatation + translation
|
||||
affine_transf = np.array([[mu_x, 0, new_x0 - mu_x * x0],
|
||||
[0, mu_y, new_y0 - mu_y * y0]])
|
||||
|
||||
sample = self.map_affine_transformation(sample, affine_transf,(new_w, new_h))
|
||||
if 'landmarks' in sample.keys():
|
||||
img_shape = np.array([0, 0, self.new_size_x, self.new_size_y])
|
||||
sample['landmarks_float'] = sample['landmarks']
|
||||
sample['mask_ldm_float'] = sample['mask_ldm']
|
||||
sample['landmarks'] = np.round(sample['landmarks'])
|
||||
sample['mask_ldm'], sample['landmarks'] = self.clean_outbbox_landmarks(img_shape, sample['landmarks'],
|
||||
sample['mask_ldm'])
|
||||
|
||||
if self.img2map_scale:
|
||||
sample = self._rescale_map(sample)
|
||||
return sample
|
||||
|
||||
def _rescale_map(self, sample):
|
||||
|
||||
# Rescale
|
||||
lnd_float = sample['landmarks_float']
|
||||
lnd_float[:, 0] = self.map_scale_x * lnd_float[:, 0]
|
||||
lnd_float[:, 1] = self.map_scale_y * lnd_float[:, 1]
|
||||
|
||||
# Filter landmarks
|
||||
lnd = np.round(lnd_float)
|
||||
filter_x = lnd[:, 0] >= self.map_size_x
|
||||
filter_y = lnd[:, 1] >= self.map_size_y
|
||||
lnd[filter_x] = self.map_size_x - 1
|
||||
lnd[filter_y] = self.map_size_y - 1
|
||||
new_lnd = (lnd.T * sample['mask_ldm']).T
|
||||
new_lnd = new_lnd.astype(int).astype(float)
|
||||
|
||||
sample['landmarks_float'] = lnd_float
|
||||
sample['landmarks'] = new_lnd
|
||||
sample['img2map_scale'] = [self.map_scale_x, self.map_scale_y]
|
||||
return sample
|
||||
|
||||
|
||||
|
||||
class OcclusionAug:
|
||||
def __init__(self, min_length=0.1, max_length=0.4, num_maps=1):
|
||||
self.min_length = min_length
|
||||
self.max_length = max_length
|
||||
self.num_maps = num_maps
|
||||
|
||||
def __call__(self, sample):
|
||||
x, y, w, h = sample['bbox']
|
||||
image = sample['image']
|
||||
landmarks = sample['landmarks']
|
||||
vis = sample['visible']
|
||||
|
||||
min_ratio = self.min_length
|
||||
max_ratio = self.max_length
|
||||
rnd_width = np.random.randint(int(w * min_ratio), int(w * max_ratio))
|
||||
rnd_height = np.random.randint(int(h * min_ratio), int(h * max_ratio))
|
||||
|
||||
# (xi, yi) and (xf, yf) are, respectively, the lower left points of the
|
||||
# occlusion rectangle and the upper right point.
|
||||
xi = int(x + np.random.randint(0, w - rnd_width))
|
||||
xf = int(xi + rnd_width)
|
||||
yi = int(y + np.random.randint(0, h - rnd_height))
|
||||
yf = int(yi + rnd_height)
|
||||
|
||||
pixels = np.array(image)
|
||||
pixels[yi:yf, xi:xf, :] = np.random.uniform(0, 255, size=3)
|
||||
image = Image.fromarray(pixels)
|
||||
sample['image'] = image
|
||||
|
||||
# Update visibilities
|
||||
filter_x1 = landmarks[:, 0] >= xi
|
||||
filter_x2 = landmarks[:, 0] < xf
|
||||
filter_x = np.logical_and(filter_x1, filter_x2)
|
||||
|
||||
filter_y1 = landmarks[:, 1] >= yi
|
||||
filter_y2 = landmarks[:, 1] < yf
|
||||
filter_y = np.logical_and(filter_y1, filter_y2)
|
||||
|
||||
filter_novis = np.logical_and(filter_x, filter_y)
|
||||
filter_vis = np.logical_not(filter_novis)
|
||||
sample['visible'] = vis * filter_vis
|
||||
return sample
|
||||
|
||||
|
||||
class LightingAug:
|
||||
def __init__(self, hsv_range_min=(-0.5, -0.5, -0.5), hsv_range_max=(0.5, 0.5, 0.5)):
|
||||
self.hsv_range_min = hsv_range_min
|
||||
self.hsv_range_max = hsv_range_max
|
||||
|
||||
def __call__(self, sample):
|
||||
# Convert to HSV colorspace from RGB colorspace
|
||||
image = np.array(sample['image'])
|
||||
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
|
||||
|
||||
# Generate new random values
|
||||
H = 1 + np.random.uniform(self.hsv_range_min[0], self.hsv_range_max[0])
|
||||
S = 1 + np.random.uniform(self.hsv_range_min[1], self.hsv_range_max[1])
|
||||
V = 1 + np.random.uniform(self.hsv_range_min[2], self.hsv_range_max[2])
|
||||
hsv[:, :, 0] = np.clip(H*hsv[:, :, 0], 0, 179)
|
||||
hsv[:, :, 1] = np.clip(S*hsv[:, :, 1], 0, 255)
|
||||
hsv[:, :, 2] = np.clip(V*hsv[:, :, 2], 0, 255)
|
||||
# Convert back to BGR colorspace
|
||||
image = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
|
||||
sample['image'] = Image.fromarray(image)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class BlurAug:
|
||||
def __init__(self, blur_prob=0.5, blur_kernel_range=(0, 2)):
|
||||
self.blur_prob = blur_prob
|
||||
self.kernel_range = blur_kernel_range
|
||||
|
||||
def __call__(self, sample):
|
||||
# Smooth image
|
||||
image = np.array(sample['image'])
|
||||
if np.random.uniform(0.0, 1.0) < self.blur_prob:
|
||||
kernel = np.random.random_integers(self.kernel_range[0], self.kernel_range[1]) * 2 + 1
|
||||
image = cv2.GaussianBlur(image, (kernel, kernel), 0, 0)
|
||||
sample['image'] = Image.fromarray(image)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
# My libs
|
||||
from .utils import rotation_matrix_to_euler
|
||||
|
||||
current_path = os.path.dirname(os.path.abspath(__file__))
|
||||
current_path = os.path.dirname(current_path)
|
||||
current_path = os.path.dirname(current_path)
|
||||
# Model file nomenclature
|
||||
model_file_dft = os.path.join(current_path,"models3D/mean_face_3D_{num_ldm}.txt")
|
||||
|
||||
|
||||
class PositPose:
|
||||
|
||||
def __init__(self, ldm_ids, focal_ratio=1, selected_ids=None, max_iter=100,
|
||||
fix_bbox=True, model_file=model_file_dft):
|
||||
|
||||
# Load 3D face model
|
||||
model3d_world, model3d_ids = self._load_world_shape(ldm_ids, model_file)
|
||||
|
||||
# Generate id mask to pick only the robust landmarks for posit
|
||||
if selected_ids is None:
|
||||
model3d_mask = np.ones(len(ldm_ids))
|
||||
else:
|
||||
model3d_mask = np.zeros(len(ldm_ids))
|
||||
for index, posit_id in enumerate(model3d_ids):
|
||||
if posit_id in selected_ids:
|
||||
model3d_mask[index] = 1
|
||||
|
||||
self.ldm_ids = ldm_ids # Ids from the database
|
||||
self.model3d_world = model3d_world # Model data
|
||||
self.model3d_ids = model3d_ids # Model ids
|
||||
self.model3d_mask = model3d_mask # Model mask ids
|
||||
self.max_iter = max_iter # Refinement iterations
|
||||
self.focal_ratio = focal_ratio # Camera matrix focal length ratio
|
||||
self.fix_bbox = fix_bbox # Camera matrix centered on image (False to centered on bbox)
|
||||
|
||||
def __call__(self, sample):
|
||||
|
||||
landmarks = sample['landmarks']
|
||||
mask = sample['mask_ldm']
|
||||
|
||||
# Camera matrix
|
||||
img_shape = np.array(sample['image'].shape)[0:2]
|
||||
if 'img2map_scale' in sample.keys():
|
||||
img_shape = img_shape * sample['img2map_scale']
|
||||
|
||||
if self.fix_bbox:
|
||||
img_bbox = [0, 0, img_shape[1], img_shape[0]] # Shapes given are inverted (y,x)
|
||||
cam_matrix = self._camera_matrix(img_bbox)
|
||||
else:
|
||||
bbox = sample['bbox'] # Scale error when ftshape and img_shape mismatch
|
||||
cam_matrix = self._camera_matrix(bbox)
|
||||
|
||||
# Save intrinsic matrix and 3D model landmarks
|
||||
sample['cam_matrix'] = cam_matrix
|
||||
sample['model3d'] = self.model3d_world
|
||||
|
||||
world_pts, image_pts = self._set_correspondences(landmarks, mask)
|
||||
|
||||
if image_pts.shape[0] < 4:
|
||||
print('POSIT does not work without landmarks')
|
||||
rot_matrix, trl_matrix = np.eye(3, dtype=float), np.array([0, 0, 0])
|
||||
else:
|
||||
rot_matrix, trl_matrix = self._modern_posit(world_pts, image_pts, cam_matrix)
|
||||
|
||||
euler = rotation_matrix_to_euler(rot_matrix)
|
||||
sample['pose'] = np.array([euler[0], euler[1], euler[2], trl_matrix[0], trl_matrix[1], trl_matrix[2]])
|
||||
sample['model3d_proj'] = self._project_points(rot_matrix, trl_matrix, cam_matrix, norm=img_shape)
|
||||
return sample
|
||||
|
||||
def _load_world_shape(self, ldm_ids, model_file):
|
||||
return load_world_shape(ldm_ids, model_file=model_file)
|
||||
|
||||
def _camera_matrix(self, bbox):
|
||||
focal_length_x = bbox[2] * self.focal_ratio
|
||||
focal_length_y = bbox[3] * self.focal_ratio
|
||||
face_center = (bbox[0] + (bbox[2] * 0.5)), (bbox[1] + (bbox[3] * 0.5))
|
||||
|
||||
cam_matrix = np.array([[focal_length_x, 0, face_center[0]],
|
||||
[0, focal_length_y, face_center[1]],
|
||||
[0, 0, 1]])
|
||||
return cam_matrix
|
||||
|
||||
def _set_correspondences(self, landmarks, mask):
|
||||
# Correspondences using labelled and robust landmarks
|
||||
img_mask = np.logical_and(mask, self.model3d_mask)
|
||||
img_mask = img_mask.astype(bool)
|
||||
|
||||
image_pts = landmarks[img_mask]
|
||||
world_pts = self.model3d_world[img_mask]
|
||||
return world_pts, image_pts
|
||||
|
||||
def _modern_posit(self, world_pts, image_pts, cam_matrix):
|
||||
return modern_posit(world_pts, image_pts, cam_matrix, self.max_iter)
|
||||
|
||||
def _project_points(self, rot, trl, cam_matrix, norm=None):
|
||||
# Perspective projection model
|
||||
trl = np.expand_dims(trl, 1)
|
||||
extrinsics = np.concatenate((rot, trl), 1)
|
||||
proj_matrix = np.matmul(cam_matrix, extrinsics)
|
||||
|
||||
# Homogeneous landmarks
|
||||
pts = self.model3d_world
|
||||
ones = np.ones(pts.shape[0])
|
||||
ones = np.expand_dims(ones, 1)
|
||||
pts_hom = np.concatenate((pts, ones), 1)
|
||||
|
||||
# Project landmarks
|
||||
pts_proj = np.matmul(proj_matrix, pts_hom.T).T
|
||||
pts_proj = pts_proj / np.expand_dims(pts_proj[:, 2], 1) # Lambda = 1
|
||||
|
||||
if norm is not None:
|
||||
pts_proj[:, 0] /= norm[0]
|
||||
pts_proj[:, 1] /= norm[1]
|
||||
return pts_proj[:, :-1]
|
||||
|
||||
|
||||
def load_world_shape(db_landmarks, model_file=model_file_dft):
|
||||
|
||||
# Load 3D mean face coordinates
|
||||
num_ldm = len(db_landmarks)
|
||||
filename = model_file.format(num_ldm=num_ldm)
|
||||
if not os.path.exists(filename):
|
||||
raise ValueError('No 3D model find for %i landmarks' % num_ldm)
|
||||
|
||||
posit_landmarks = np.genfromtxt(filename, delimiter='|', dtype=int, usecols=0).tolist()
|
||||
mean_face_3D = np.genfromtxt(filename, delimiter='|', dtype=(float, float, float), usecols=(1, 2, 3)).tolist()
|
||||
world_all = len(mean_face_3D)*[None]
|
||||
index_all = len(mean_face_3D)*[None]
|
||||
|
||||
for cont, elem in enumerate(mean_face_3D):
|
||||
pt3d = [elem[2], -elem[0], -elem[1]]
|
||||
lnd_idx = db_landmarks.index(posit_landmarks[cont])
|
||||
world_all[lnd_idx] = pt3d
|
||||
index_all[lnd_idx] = posit_landmarks[cont]
|
||||
|
||||
return np.array(world_all), np.array(index_all)
|
||||
|
||||
|
||||
def modern_posit(world_pts, image_pts, cam_matrix, max_iters):
|
||||
# Homogeneous world points
|
||||
num_landmarks = image_pts.shape[0]
|
||||
one = np.ones((num_landmarks, 1))
|
||||
A = np.concatenate((world_pts, one), axis=1)
|
||||
B = np.linalg.pinv(A)
|
||||
|
||||
# Normalize image points
|
||||
focal_length = cam_matrix[0,0]
|
||||
img_center = (cam_matrix[0,2], cam_matrix[1,2])
|
||||
centered_pts = np.zeros((num_landmarks,2))
|
||||
centered_pts[:,0] = (image_pts[:,0]-img_center[0])/focal_length
|
||||
centered_pts[:,1] = (image_pts[:,1]-img_center[1])/focal_length
|
||||
Ui = centered_pts[:,0]
|
||||
Vi = centered_pts[:,1]
|
||||
|
||||
# POSIT loop
|
||||
Tx, Ty, Tz = 0.0, 0.0, 0.0
|
||||
r1, r2, r3 = [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]
|
||||
for iter in range(0, max_iters):
|
||||
I = np.dot(B,Ui)
|
||||
J = np.dot(B,Vi)
|
||||
|
||||
# Estimate translation vector and rotation matrix
|
||||
normI = 1.0 / np.sqrt(I[0] * I[0] + I[1] * I[1] + I[2] * I[2])
|
||||
normJ = 1.0 / np.sqrt(J[0] * J[0] + J[1] * J[1] + J[2] * J[2])
|
||||
Tz = np.sqrt(normI * normJ) # geometric average instead of arithmetic average of classicPosit
|
||||
r1N = I*Tz
|
||||
r2N = J*Tz
|
||||
r1 = r1N[0:3]
|
||||
r2 = r2N[0:3]
|
||||
r1 = np.clip(r1, -1, 1)
|
||||
r2 = np.clip(r2, -1, 1)
|
||||
r3 = np.cross(r1,r2)
|
||||
r3T = np.concatenate((r3, [Tz]), axis=0)
|
||||
Tx = r1N[3]
|
||||
Ty = r2N[3]
|
||||
|
||||
# Compute epsilon, update Ui and Vi and check convergence
|
||||
eps = np.dot(A, r3T)/Tz
|
||||
oldUi = Ui
|
||||
oldVi = Vi
|
||||
Ui = np.multiply(eps, centered_pts[:,0])
|
||||
Vi = np.multiply(eps, centered_pts[:,1])
|
||||
deltaUi = Ui - oldUi
|
||||
deltaVi = Vi - oldVi
|
||||
delta = focal_length * focal_length * (np.dot(np.transpose(deltaUi), deltaUi) + np.dot(np.transpose(deltaVi), deltaVi))
|
||||
if iter > 0 and delta < 0.01: # converged
|
||||
break
|
||||
|
||||
rot_matrix = np.array([np.transpose(r1), np.transpose(r2), np.transpose(r3)])
|
||||
trl_matrix = np.array([Tx, Ty, Tz])
|
||||
# Convert to the nearest orthogonal rotation matrix
|
||||
w, u, vt = cv2.SVDecomp(rot_matrix) # R = U*D*Vt
|
||||
rot_matrix = np.matmul(np.matmul(u, np.eye(3, dtype=float)), vt)
|
||||
return rot_matrix, trl_matrix
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
def affine2homogeneous(points):
|
||||
'''Returns the points completed with a new last coordinate
|
||||
equal to 1
|
||||
Arguments
|
||||
---------
|
||||
points: np.array of shape (num_points, dim)
|
||||
Returns
|
||||
-------
|
||||
hpoints: np.array of shape (num_points, dim + 1),
|
||||
of the points completed with ones'''
|
||||
|
||||
num_points = points.shape[0]
|
||||
hpoints = np.hstack(
|
||||
(points, np.repeat(1, num_points).reshape(num_points, 1)))
|
||||
return hpoints
|
||||
|
||||
|
||||
def get_similarity_matrix(deg_angle, scale, center):
|
||||
'''Similarity matrix.
|
||||
Arguments:
|
||||
---------
|
||||
deg_angle: rotation angle in degrees
|
||||
scale: factor scale
|
||||
center: coordinates of the rotation center
|
||||
|
||||
Returns:
|
||||
-------
|
||||
matrix: (2, 3) numpy array representing the
|
||||
similarity matrix.
|
||||
'''
|
||||
x0, y0 = center
|
||||
angle = np.radians(deg_angle)
|
||||
|
||||
matrix = np.zeros((2, 3))
|
||||
matrix[0:2, 0:2] = [[np.cos(angle), -np.sin(angle)],
|
||||
[np.sin(angle), np.cos(angle)]]
|
||||
matrix[0: 2, 0: 2] *= scale
|
||||
|
||||
matrix[:, 2] = [(1 - scale * np.cos(angle)) * x0 +
|
||||
scale * np.sin(angle) * y0,
|
||||
-scale * np.sin(angle) * x0 +
|
||||
(1 - scale * np.cos(angle)) * y0]
|
||||
return matrix
|
||||
|
||||
|
||||
def get_inverse_similarity_matrix(deg_angle, scale, center):
|
||||
'''Returns the inverse of the affine similarity
|
||||
Arguments
|
||||
---------
|
||||
deg_angle: angle in degrees of the rotation
|
||||
center: iterable of two components (x0, y0),
|
||||
center of the rotation
|
||||
scale: float, scale factor
|
||||
Returns
|
||||
-------
|
||||
matrix: np.array of shape (2, 3) with the coordinates of
|
||||
the inverse of the similarity'''
|
||||
|
||||
x0, y0 = center
|
||||
angle = np.radians(deg_angle)
|
||||
inv_scale = 1 / scale
|
||||
matrix = np.zeros((2, 3))
|
||||
matrix[0:2, 0:2] = [[np.cos(angle), np.sin(angle)],
|
||||
[-np.sin(angle), np.cos(angle)]]
|
||||
matrix[0:2, 0:2] *= inv_scale
|
||||
|
||||
matrix[:, 2] = [(1 - inv_scale * np.cos(angle)) * x0 -
|
||||
inv_scale * np.sin(angle) * y0,
|
||||
inv_scale * np.sin(angle) * x0 +
|
||||
(1 - inv_scale * np.cos(angle)) * y0]
|
||||
|
||||
return matrix
|
||||
|
||||
|
||||
def get_inverse_transf(affine_transf):
|
||||
A = affine_transf[0:2, 0:2]
|
||||
b = affine_transf[:, 2]
|
||||
|
||||
inv_A = np.linalg.inv(A) # we assume A invertible!
|
||||
|
||||
inv_affine = np.zeros((2, 3))
|
||||
inv_affine[0:2, 0:2] = inv_A
|
||||
inv_affine[:, 2] = -inv_A.dot(b)
|
||||
|
||||
return inv_affine
|
||||
|
||||
|
||||
def image2vect(image):
|
||||
'''
|
||||
Input:
|
||||
image[batch_size, num_channels, im_size_x, im_size_y]
|
||||
Output:
|
||||
vect[batch_size, num_channels, im_size_x*im_size_y]
|
||||
'''
|
||||
vect = image.reshape(*image.shape[0:-2], -1)
|
||||
return vect
|
||||
|
||||
|
||||
def rotation_matrix_to_euler(rot_matrix):
|
||||
# http://euclideanspace.com/maths/geometry/rotations/conversions/matrixToEuler/index.htm
|
||||
a00, a01, a02 = rot_matrix[0, 0], rot_matrix[0, 1], rot_matrix[0, 2]
|
||||
a10, a11, a12 = rot_matrix[1, 0], rot_matrix[1, 1], rot_matrix[1, 2]
|
||||
a20, a21, a22 = rot_matrix[2, 0], rot_matrix[2, 1], rot_matrix[2, 2]
|
||||
if abs(1.0 - a10) <= np.finfo(float).eps: # singularity at north pole / special case a10 == 1
|
||||
yaw = np.arctan2(a02, a22)
|
||||
pitch = np.pi/2.0
|
||||
roll = 0
|
||||
elif abs(-1.0 - a10) <= np.finfo(float).eps: # singularity at south pole / special case a10 == -1
|
||||
yaw = np.arctan2(a02, a22)
|
||||
pitch = -np.pi/2.0
|
||||
roll = 0
|
||||
else: # standard case
|
||||
yaw = np.arctan2(-a20, a00)
|
||||
pitch = np.arcsin(a10)
|
||||
roll = np.arctan2(-a12, a11)
|
||||
# Convert to degrees
|
||||
euler = np.array([yaw, pitch, roll])*(180.0/np.pi)
|
||||
# Change coordinates system
|
||||
euler = np.array([(-euler[0])+90, -euler[1], (-euler[2])-90])
|
||||
if euler[0] > 180: euler[0] -= 360
|
||||
elif euler[0] < -180: euler[0] += 360
|
||||
if euler[1] > 180: euler[1] -= 360
|
||||
elif euler[1] < -180: euler[1] += 360
|
||||
if euler[2] > 180: euler[2] -= 360
|
||||
elif euler[2] < -180: euler[2] += 360
|
||||
return euler
|
||||
|
||||
|
||||
def euler_to_rotation_matrix(headpose):
|
||||
# http://euclideanspace.com/maths/geometry/rotations/conversions/eulerToMatrix/index.htm
|
||||
# Change coordinates system
|
||||
euler = np.array([-(headpose[0]-90), -headpose[1], -(headpose[2]+90)])
|
||||
# Convert to radians
|
||||
rad = euler*(np.pi/180.0)
|
||||
cy = np.cos(rad[0])
|
||||
sy = np.sin(rad[0])
|
||||
cp = np.cos(rad[1])
|
||||
sp = np.sin(rad[1])
|
||||
cr = np.cos(rad[2])
|
||||
sr = np.sin(rad[2])
|
||||
Ry = np.array([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]]) # yaw
|
||||
Rp = np.array([[cp, -sp, 0.0], [sp, cp, 0.0], [0.0, 0.0, 1.0]]) # pitch
|
||||
Rr = np.array([[1.0, 0.0, 0.0], [0.0, cr, -sr], [0.0, sr, cr]]) # roll
|
||||
return np.matmul(np.matmul(Ry, Rp), Rr)
|
||||
@@ -0,0 +1,41 @@
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
|
||||
from .import alignments as zoo_alignments
|
||||
|
||||
zoos = [zoo_alignments]
|
||||
|
||||
|
||||
def get_dataset(data_config, pretreat=None, debug=False):
|
||||
|
||||
for zoo in zoos:
|
||||
dataset = zoo.get_dataset(data_config, pretreat=pretreat, debug=debug)
|
||||
if dataset is not None:
|
||||
return dataset
|
||||
raise NotImplementedError('Dataset not available')
|
||||
|
||||
|
||||
def get_dataloader(batch_size, data_config, pretreat=None, sampler_cfg=None, debug=False):
|
||||
|
||||
dataset = get_dataset(data_config, pretreat=pretreat, debug=debug)
|
||||
|
||||
if (len(dataset) % batch_size) == 1 and data_config.shuffle == True:
|
||||
drop_last_batch = True
|
||||
else:
|
||||
drop_last_batch = False
|
||||
|
||||
shuffle = data_config.shuffle
|
||||
sampler = None
|
||||
if sampler_cfg is not None:
|
||||
sampler = DistributedSampler(dataset, num_replicas=sampler_cfg.world_size, rank=sampler_cfg.rank)
|
||||
shuffle = False
|
||||
|
||||
dataloader = DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
num_workers=data_config.num_workers,
|
||||
pin_memory=True,
|
||||
drop_last=drop_last_batch,
|
||||
sampler=sampler)
|
||||
|
||||
return dataloader, dataset
|
||||
@@ -0,0 +1,173 @@
|
||||
import os
|
||||
import json
|
||||
from collections import OrderedDict
|
||||
|
||||
current_path = os.path.dirname(os.path.abspath(__file__))
|
||||
current_path= os.path.dirname(current_path)
|
||||
db_img_path =os.path.join(current_path,"databases")
|
||||
db_anns_path=os.path.join(current_path,"annotations","{database}/{file_name}.json")
|
||||
# Default data paths
|
||||
#db_img_path = "spiga/data/annotations/300wpublic"
|
||||
#db_anns_path = "spiga/data/annotations/300wpublic/db_info.json"
|
||||
|
||||
class AlignConfig:
|
||||
|
||||
def __init__(self, database_name, mode='train'):
|
||||
# Dataset
|
||||
self.database_name = database_name
|
||||
self.working_mode = mode
|
||||
self.database = None # Set at self._update_database()
|
||||
self.anns_file = None # Set at self._update_database()
|
||||
self.image_dir = None # Set at self._update_database()
|
||||
self._update_database()
|
||||
self.image_size = (256, 256)
|
||||
self.ftmap_size = (256, 256)
|
||||
|
||||
# Dataloaders
|
||||
self.ids = None # List of a subset if need it
|
||||
self.shuffle = True # Shuffle samples
|
||||
self.num_workers = 4 # Threads
|
||||
|
||||
# Posit
|
||||
self.generate_pose = True # Generate pose parameters from landmarks
|
||||
self.focal_ratio = 1.5 # Camera matrix focal length ratio
|
||||
self.posit_max_iter = 100 # Refinement iterations
|
||||
|
||||
# Subset of robust ids in the 3D model to use in posit.
|
||||
# 'None' to use all the available model landmarks.
|
||||
self.posit_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
|
||||
14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
|
||||
|
||||
# Data augmentation
|
||||
# Control augmentations with the following list, crop to self.img_size is mandatory, check target_dist param.
|
||||
if mode == 'train':
|
||||
self.aug_names = ['flip', 'rotate_scale', 'occlusion', 'lighting', 'blur']
|
||||
else:
|
||||
self.aug_names = []
|
||||
self.shuffle = False
|
||||
|
||||
# Flip
|
||||
self.hflip_prob = 0.5
|
||||
# Rotation
|
||||
self.angle_range = 45.
|
||||
# Scale
|
||||
self.scale_max = 0.15
|
||||
self.scale_min = -0.15
|
||||
# Translation
|
||||
self.trl_ratio = 0.05 # Translation augmentation
|
||||
# Crop target rescale
|
||||
self.target_dist = 1.6 # Target distance zoom in/out around face. Default: 1.
|
||||
# Occlusion
|
||||
self.occluded_max_len = 0.4
|
||||
self.occluded_min_len = 0.1
|
||||
self.occluded_covar_ratio = 2.25**0.5
|
||||
# Lighting
|
||||
self.hsv_range_min = [-0.5, -0.5, -0.5]
|
||||
self.hsv_range_max = [0.5, 0.5, 0.5]
|
||||
# Blur
|
||||
self.blur_prob = 0.5
|
||||
self.blur_kernel_range = [0, 2]
|
||||
|
||||
# Heatmaps 2D
|
||||
self.sigma2D = 1.5
|
||||
self.heatmap2D_norm = False
|
||||
|
||||
# Boundaries
|
||||
self.sigmaBD = 1
|
||||
|
||||
def update(self, params_dict):
|
||||
state_dict = self.state_dict()
|
||||
for k, v in params_dict.items():
|
||||
if k in state_dict or hasattr(self, k):
|
||||
setattr(self, k, v)
|
||||
else:
|
||||
Warning('Unknown option: {}: {}'.format(k, v))
|
||||
self._update_database()
|
||||
|
||||
def state_dict(self, tojson=False):
|
||||
state_dict = OrderedDict()
|
||||
for k in self.__dict__.keys():
|
||||
if not k.startswith('_'):
|
||||
if tojson and k in ['database']:
|
||||
continue
|
||||
state_dict[k] = getattr(self, k)
|
||||
return state_dict
|
||||
|
||||
def _update_database(self):
|
||||
self.database = DatabaseStruct(self.database_name)
|
||||
self.anns_file = db_anns_path.format(database=self.database_name, file_name=self.working_mode)
|
||||
self.image_dir = self._get_imgdb_path()
|
||||
|
||||
def _get_imgdb_path(self):
|
||||
img_dir = None
|
||||
if self.database_name in ['300wpublic', '300wprivate']:
|
||||
img_dir = db_img_path + '/300w/'
|
||||
elif self.database_name in ['aflw19', 'merlrav']:
|
||||
img_dir = db_img_path + '/aflw/data/'
|
||||
elif self.database_name in ['cofw', 'cofw68']:
|
||||
img_dir = db_img_path + '/cofw/'
|
||||
elif self.database_name in ['wflw']:
|
||||
img_dir = db_img_path + '/wflw/'
|
||||
return img_dir
|
||||
|
||||
def __str__(self):
|
||||
state_dict = self.state_dict()
|
||||
text = 'Dataloader {\n'
|
||||
for k, v in state_dict.items():
|
||||
if isinstance(v, DatabaseStruct):
|
||||
text += '\t{}: {}'.format(k, str(v).expandtabs(12))
|
||||
else:
|
||||
text += '\t{}: {}\n'.format(k, v)
|
||||
text += '\t}\n'
|
||||
return text
|
||||
|
||||
|
||||
class DatabaseStruct:
|
||||
|
||||
def __init__(self, database_name):
|
||||
|
||||
self.name = database_name
|
||||
self.ldm_ids, self.ldm_flip_order, self.ldm_edges_matrix = self._get_database_specifics()
|
||||
self.num_landmarks = len(self.ldm_ids)
|
||||
self.num_edges = len(self.ldm_edges_matrix[0])-1
|
||||
self.fields = ['imgpath', 'bbox', 'headpose', 'ids', 'landmarks', 'visible']
|
||||
|
||||
def _get_database_specifics(self):
|
||||
'''Returns specifics ids and horizontal flip reorder'''
|
||||
|
||||
database_name = self.name
|
||||
db_info_file = db_anns_path.format(database=database_name, file_name='db_info')
|
||||
ldm_edges_matrix = None
|
||||
|
||||
if os.path.exists(db_info_file):
|
||||
with open(db_info_file) as jsonfile:
|
||||
db_info = json.load(jsonfile)
|
||||
|
||||
ldm_ids = db_info['ldm_ids']
|
||||
ldm_flip_order = db_info['ldm_flip_order']
|
||||
if 'ldm_edges_matrix' in db_info.keys():
|
||||
ldm_edges_matrix = db_info['ldm_edges_matrix']
|
||||
|
||||
else:
|
||||
raise ValueError('Database ' + database_name + 'specifics not defined. Missing db_info.json')
|
||||
|
||||
return ldm_ids, ldm_flip_order, ldm_edges_matrix
|
||||
|
||||
def state_dict(self):
|
||||
state_dict = OrderedDict()
|
||||
for k in self.__dict__.keys():
|
||||
if not k.startswith('_'):
|
||||
state_dict[k] = getattr(self, k)
|
||||
|
||||
return state_dict
|
||||
|
||||
def __str__(self):
|
||||
state_dict = self.state_dict()
|
||||
text = 'Database {\n'
|
||||
for k, v in state_dict.items():
|
||||
text += '\t{}: {}\n'.format(k, v)
|
||||
text += '\t}\n'
|
||||
return text
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .augmentors.modern_posit import PositPose
|
||||
from .augmentors.heatmaps import Heatmaps
|
||||
from .augmentors.boundary import AddBoundary
|
||||
from .augmentors.landmarks import HorizontalFlipAug, RSTAug, OcclusionAug, \
|
||||
LightingAug, BlurAug, TargetCropAug
|
||||
|
||||
|
||||
def get_transformers(data_config):
|
||||
|
||||
# Data augmentation
|
||||
aug_names = data_config.aug_names
|
||||
augmentors = []
|
||||
|
||||
if 'flip' in aug_names:
|
||||
augmentors.append(HorizontalFlipAug(data_config.database.ldm_flip_order, data_config.hflip_prob))
|
||||
if 'rotate_scale' in aug_names:
|
||||
augmentors.append(RSTAug(data_config.angle_range, data_config.scale_min,
|
||||
data_config.scale_max, data_config.trl_ratio))
|
||||
if 'occlusion' in aug_names:
|
||||
augmentors.append(OcclusionAug(data_config.occluded_min_len,
|
||||
data_config.occluded_max_len,
|
||||
data_config.database.num_landmarks))
|
||||
if 'lighting' in aug_names:
|
||||
augmentors.append(LightingAug(data_config.hsv_range_min, data_config.hsv_range_max))
|
||||
if 'blur' in aug_names:
|
||||
augmentors.append(BlurAug(data_config.blur_prob, data_config.blur_kernel_range))
|
||||
|
||||
# Crop mandatory
|
||||
augmentors.append(TargetCropAug(data_config.image_size, data_config.ftmap_size, data_config.target_dist))
|
||||
# Opencv style
|
||||
augmentors.append(ToOpencv())
|
||||
|
||||
# Gaussian heatmaps
|
||||
if 'heatmaps2D' in aug_names:
|
||||
augmentors.append(Heatmaps(data_config.database.num_landmarks, data_config.ftmap_size,
|
||||
data_config.sigma2D, norm=data_config.heatmap2D_norm))
|
||||
|
||||
if 'boundaries' in aug_names:
|
||||
augmentors.append(AddBoundary(num_landmarks=data_config.database.num_landmarks,
|
||||
map_size=data_config.ftmap_size,
|
||||
sigma=data_config.sigmaBD))
|
||||
# Pose generator
|
||||
if data_config.generate_pose:
|
||||
augmentors.append(PositPose(data_config.database.ldm_ids,
|
||||
focal_ratio=data_config.focal_ratio,
|
||||
selected_ids=data_config.posit_ids,
|
||||
max_iter=data_config.posit_max_iter))
|
||||
|
||||
return augmentors
|
||||
|
||||
|
||||
class ToOpencv:
|
||||
def __call__(self, sample):
|
||||
# Convert in a numpy array and change to GBR
|
||||
image = np.array(sample['image'])
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
sample['image'] = image
|
||||
return sample
|
||||
|
||||
|
||||
class TargetCrop(TargetCropAug):
|
||||
def __init__(self, crop_size=256, target_dist=1.6):
|
||||
super(TargetCrop, self).__init__(crop_size, crop_size, target_dist)
|
||||
|
||||
|
||||
class AddModel3D(PositPose):
|
||||
def __init__(self, ldm_ids, ftmap_size=(256, 256), focal_ratio=1.5, totensor=False):
|
||||
super(AddModel3D, self).__init__(ldm_ids, focal_ratio=focal_ratio)
|
||||
img_bbox = [0, 0, ftmap_size[1], ftmap_size[0]] # Shapes given are inverted (y,x)
|
||||
self.cam_matrix = self._camera_matrix(img_bbox)
|
||||
|
||||
if totensor:
|
||||
self.cam_matrix = torch.tensor(self.cam_matrix, dtype=torch.float)
|
||||
self.model3d_world = torch.tensor(self.model3d_world, dtype=torch.float)
|
||||
|
||||
def __call__(self, sample={}):
|
||||
# Save intrinsic matrix and 3D model landmarks
|
||||
sample['cam_matrix'] = self.cam_matrix
|
||||
sample['model3d'] = self.model3d_world
|
||||
return sample
|
||||
@@ -0,0 +1,68 @@
|
||||
101|-0.71046061493|0.39334543762|-0.853184236267
|
||||
102|-0.685409656726|0.169750713576|-0.878623094747
|
||||
103|-0.607672440834|-0.0597969416711|-0.922355056609
|
||||
104|-0.545664962607|-0.262499453743|-0.920836599126
|
||||
105|-0.446214526601|-0.439045175771|-0.863139209
|
||||
106|-0.334435950087|-0.566049889292|-0.728913281016
|
||||
107|-0.236168305748|-0.653698223145|-0.585459059996
|
||||
108|-0.128307367184|-0.712045321067|-0.488221730877
|
||||
24|-0.00354849146401|-0.73799619898|-0.461364289002
|
||||
110|0.131199109787|-0.733135962679|-0.466262199442
|
||||
111|0.257373771199|-0.666487934012|-0.560781626341
|
||||
112|0.359823738861|-0.569594984479|-0.666664889106
|
||||
113|0.455878939364|-0.475272786995|-0.762199581604
|
||||
114|0.583151236654|-0.255441185628|-0.895445263128
|
||||
115|0.635035270608|-0.076593055883|-0.925024981148
|
||||
116|0.688562921266|0.127095701882|-0.90624669312
|
||||
117|0.711883056933|0.332170073984|-0.901990426153
|
||||
7|-0.406302776083|0.360389456414|-0.383328256737
|
||||
138|-0.34477414508|0.384294131116|-0.310253684249
|
||||
139|-0.261521910706|0.396020592391|-0.286645173372
|
||||
8|-0.160787895282|0.371046071664|-0.330243237009
|
||||
141|-0.245621576158|0.334421719594|-0.323105980143
|
||||
142|-0.328092143154|0.331973054751|-0.339942994735
|
||||
11|0.163880511996|0.34074733643|-0.314140711047
|
||||
144|0.250844605212|0.392490504074|-0.279030578851
|
||||
145|0.340770354713|0.378980645804|-0.285841921386
|
||||
12|0.432352000574|0.348665008376|-0.377158168878
|
||||
147|0.333467711264|0.328033806608|-0.318530865652
|
||||
148|0.246621391603|0.335053488011|-0.313137199532
|
||||
1|-0.552168877878|0.483792334207|-0.466456539274
|
||||
119|-0.459810924979|0.561780416181|-0.308160933539
|
||||
2|-0.360178576451|0.566290707558|-0.216753374881
|
||||
121|-0.249770054109|0.530336745189|-0.148694799479
|
||||
3|-0.101941089981|0.482471777751|-0.124216132105
|
||||
128|0.00843526194116|0.416310824525|-0.139142361862
|
||||
129|0.00731217835113|0.267155736618|-0.106471830332
|
||||
130|0.00190174006432|0.142359799355|-0.0544539305898
|
||||
17|0.0|0.0|0.0
|
||||
16|-0.118903311783|-0.0792784542251|-0.23553779698
|
||||
133|-0.0504300242939|-0.104665185108|-0.171210015652
|
||||
134|0.00312657092554|-0.111845126622|-0.154139340027
|
||||
135|0.0687543983424|-0.0995458563329|-0.173275169304
|
||||
18|0.148988810614|-0.0737848514841|-0.23851346235
|
||||
4|0.0979764884322|0.465935806416|-0.131673602172
|
||||
124|0.214281256931|0.497130488239|-0.135680573937
|
||||
5|0.373011222167|0.520379035113|-0.223000042464
|
||||
126|0.501607216527|0.500371303274|-0.31925158046
|
||||
6|0.57362522379|0.457842172776|-0.465084060802
|
||||
20|-0.2229832011|-0.277627584817|-0.354306324758
|
||||
150|-0.157062857603|-0.248737641674|-0.270218764688
|
||||
151|-0.0816354856174|-0.23618498428|-0.225237980338
|
||||
22|-0.00391577759471|-0.24499812323|-0.2126098208
|
||||
153|0.0629421066443|-0.230774971512|-0.208501270489
|
||||
154|0.135507933036|-0.240646328001|-0.235471195707
|
||||
21|0.240381217797|-0.283709533358|-0.332281891845
|
||||
156|0.143706630542|-0.329461104518|-0.283978333058
|
||||
157|0.084139951946|-0.355354733249|-0.256464077311
|
||||
23|0.00813704018925|-0.360357397837|-0.250423014186
|
||||
159|-0.0778260690655|-0.349041413147|-0.26966136614
|
||||
160|-0.157111645515|-0.311971743234|-0.322383820174
|
||||
161|-0.182660017757|-0.278377925956|-0.341170846656
|
||||
162|-0.075943475389|-0.283519721316|-0.270854173882
|
||||
163|-0.0026701322545|-0.286917503123|-0.254428375696
|
||||
164|0.0685062003389|-0.285630836758|-0.257845545665
|
||||
165|0.194170296105|-0.279641061937|-0.314605542849
|
||||
166|0.0689862474833|-0.288076969827|-0.258696545259
|
||||
167|-0.00273913414372|-0.28915723747|-0.25622946708
|
||||
168|-0.0766989697362|-0.286311968999|-0.273027820051
|
||||
@@ -0,0 +1,98 @@
|
||||
100|-0.710460614932|0.393345437621|-0.853184236268
|
||||
101|-0.697935135807|0.281548075586|-0.865903665506
|
||||
102|-0.685409656727|0.169750713576|-0.878623094748
|
||||
103|-0.646541048786|0.0549768859803|-0.90048907572
|
||||
104|-0.607672440835|-0.0597969416712|-0.92235505661
|
||||
105|-0.576668701754|-0.16114819775|-0.921595827864
|
||||
106|-0.545664962608|-0.262499453743|-0.920836599128
|
||||
107|-0.495939744597|-0.350772314723|-0.891987904048
|
||||
108|-0.446214526602|-0.439045175772|-0.863139209001
|
||||
109|-0.390325238349|-0.502547532515|-0.796026245014
|
||||
110|-0.334435950088|-0.566049889293|-0.728913281017
|
||||
111|-0.285302127936|-0.60987405623|-0.657186170535
|
||||
112|-0.236168305748|-0.653698223146|-0.585459059997
|
||||
113|-0.182237836481|-0.682871772143|-0.536840395466
|
||||
114|-0.128307367184|-0.712045321068|-0.488221730878
|
||||
115|-0.065927929287|-0.725020759991|-0.474793009911
|
||||
24|-0.00354849146405|-0.737996198981|-0.461364289002
|
||||
117|0.0638253091608|-0.735566080866|-0.463813244263
|
||||
118|0.131199109787|-0.73313596268|-0.466262199443
|
||||
119|0.194286440465|-0.699811948342|-0.513521912927
|
||||
120|0.257373771199|-0.666487934013|-0.560781626342
|
||||
121|0.308598755049|-0.61804145929|-0.613723257723
|
||||
122|0.359823738861|-0.569594984479|-0.666664889106
|
||||
123|0.407851339087|-0.522433885778|-0.714432235344
|
||||
124|0.455878939366|-0.475272786996|-0.762199581605
|
||||
125|0.519515088045|-0.365356986341|-0.828822422342
|
||||
126|0.583151236655|-0.255441185628|-0.89544526313
|
||||
127|0.609093253652|-0.166017120718|-0.910235122162
|
||||
128|0.635035270609|-0.0765930558831|-0.92502498115
|
||||
129|0.661799095895|0.0252513229862|-0.915635837119
|
||||
130|0.688562921268|0.127095701883|-0.906246693121
|
||||
131|0.700222989101|0.229632887961|-0.904118559673
|
||||
132|0.711883056935|0.332170073985|-0.901990426154
|
||||
1|-0.552168877879|0.483792334207|-0.466456539275
|
||||
134|-0.45981092498|0.561780416182|-0.30816093354
|
||||
2|-0.360178576451|0.566290707559|-0.216753374882
|
||||
136|-0.249770054109|0.53033674519|-0.148694799479
|
||||
3|-0.101941089981|0.482471777751|-0.124216132105
|
||||
138|-0.100790757462|0.429593539303|-0.13295377507
|
||||
139|-0.24861972159|0.477458506742|-0.157432442443
|
||||
140|-0.359028243932|0.513412469111|-0.225491017847
|
||||
141|-0.458660592461|0.508902177734|-0.316898576504
|
||||
4|0.0979764884324|0.465935806417|-0.131673602172
|
||||
143|0.214281256932|0.49713048824|-0.135680573937
|
||||
5|0.373011222167|0.520379035114|-0.223000042464
|
||||
145|0.501607216528|0.500371303275|-0.31925158046
|
||||
6|0.573625223791|0.457842172777|-0.465084060803
|
||||
147|0.500456884008|0.447493064827|-0.310513937496
|
||||
148|0.374161554686|0.467500796666|-0.231737685429
|
||||
149|0.215431589451|0.444252249792|-0.144418216902
|
||||
150|0.0991268209517|0.413057567969|-0.140411245136
|
||||
151|0.00843526194114|0.416310824526|-0.139142361862
|
||||
152|0.00731217835118|0.267155736618|-0.106471830333
|
||||
153|0.00190174006432|0.142359799356|-0.0544539305899
|
||||
17|0.0|0.0|0.0
|
||||
16|-0.118903311783|-0.0792784542252|-0.23553779698
|
||||
156|-0.0504300242939|-0.104665185108|-0.171210015652
|
||||
157|0.00312657092558|-0.111845126622|-0.154139340027
|
||||
158|0.0687543983424|-0.0995458563331|-0.173275169304
|
||||
18|0.148988810614|-0.0737848514843|-0.23851346235
|
||||
7|-0.406302776084|0.360389456414|-0.383328256738
|
||||
161|-0.34477414508|0.384294131116|-0.310253684249
|
||||
9|-0.303148027928|0.398970401478|-0.298449428843
|
||||
163|-0.261521910707|0.396020592392|-0.286645173373
|
||||
8|-0.160787895283|0.371046071665|-0.33024323701
|
||||
165|-0.245621576158|0.334421719595|-0.323105980144
|
||||
10|-0.286856859665|0.324384347457|-0.33152448747
|
||||
167|-0.328092143155|0.331973054752|-0.339942994736
|
||||
11|0.163880511996|0.340747336431|-0.314140711048
|
||||
169|0.250844605212|0.392490504075|-0.279030578852
|
||||
13|0.295807479997|0.39454861469|-0.282436250152
|
||||
171|0.340770354713|0.378980645804|-0.285841921386
|
||||
12|0.432352000574|0.348665008377|-0.377158168879
|
||||
173|0.333467711265|0.328033806608|-0.318530865652
|
||||
14|0.290044551442|0.322730607606|-0.315834032602
|
||||
175|0.246621391603|0.335053488012|-0.313137199533
|
||||
20|-0.2229832011|-0.277627584818|-0.354306324759
|
||||
177|-0.157062857603|-0.248737641675|-0.270218764688
|
||||
178|-0.0816354856176|-0.23618498428|-0.225237980339
|
||||
22|-0.00391577759461|-0.24499812323|-0.212609820801
|
||||
180|0.0629421066443|-0.230774971513|-0.20850127049
|
||||
181|0.135507933036|-0.240646328002|-0.235471195707
|
||||
21|0.240381217797|-0.283709533359|-0.332281891846
|
||||
183|0.143706630542|-0.329461104519|-0.283978333058
|
||||
184|0.0841399519463|-0.35535473325|-0.256464077312
|
||||
23|0.00813704018926|-0.360357397838|-0.250423014186
|
||||
186|-0.0778260690653|-0.349041413147|-0.269661366141
|
||||
187|-0.157111645515|-0.311971743235|-0.322383820174
|
||||
188|-0.182660017758|-0.278377925956|-0.341170846657
|
||||
189|-0.0759434753887|-0.283519721316|-0.270854173882
|
||||
190|-0.00267013225455|-0.286917503124|-0.254428375697
|
||||
191|0.068506200339|-0.285630836758|-0.257845545665
|
||||
192|0.194170296105|-0.279641061938|-0.31460554285
|
||||
193|0.0689862474834|-0.288076969827|-0.25869654526
|
||||
194|-0.00273913414368|-0.28915723747|-0.25622946708
|
||||
195|-0.0766989697367|-0.286311968999|-0.273027820052
|
||||
196|-0.291183407723|0.363024170951|-0.328919887677
|
||||
197|0.294656095888|0.353995131537|-0.31463990759
|
||||
@@ -0,0 +1,37 @@
|
||||
import argparse
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
def main():
|
||||
# Input arguments control
|
||||
pars = argparse.ArgumentParser(description='3D model visualization')
|
||||
pars.add_argument('file', type=str, help='File txt path')
|
||||
args = pars.parse_args()
|
||||
visualize_3Dmodel(args.file)
|
||||
|
||||
|
||||
def visualize_3Dmodel(input_file):
|
||||
|
||||
with open(input_file) as f:
|
||||
lines = f.readlines()
|
||||
|
||||
model = []
|
||||
for line in lines:
|
||||
line = line[:-1] # Remove \n
|
||||
line_split = line.split('|')
|
||||
values = np.array(line_split, dtype=float)
|
||||
model.append(values)
|
||||
|
||||
model = np.array(model)
|
||||
model_xyz = model[:, 1:]
|
||||
|
||||
# Show model
|
||||
fig = plt.figure()
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
ax.scatter(model_xyz[:, 0], model_xyz[:, 1], model_xyz[:, 2]+0.8)
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,72 @@
|
||||
# SPIGA: Dataloaders
|
||||
Currently, the repository provides a pytorch based image dataloader implementation for the task of:
|
||||
* **Facial Landmarks estimation**.
|
||||
* **Headpose estimation**.
|
||||
* **Facial landmarks visibilities**.
|
||||
|
||||
The dataloader can be used for training or testing the networks and it includes general and
|
||||
specifics data augmentation for each task, as it can be image partial occlusions
|
||||
or headpose generation from facial landmarks.
|
||||
|
||||
In addition, the framework provides a wide benchmark software in order to evaluate the
|
||||
different task along the following databases:
|
||||
* **WFLW**.
|
||||
* **MERLRAV (AFLW 68)**
|
||||
* **COFW68**.
|
||||
* **300W Public, Private**.
|
||||
<details>
|
||||
<summary> Coming soon... </summary>
|
||||
|
||||
* AFLW, AFLW19, AFLW2000 (test only).
|
||||
* Menpo and 3D Menpo.
|
||||
* COFW.
|
||||
* 300WLP
|
||||
* 300W Masked
|
||||
|
||||
</details>
|
||||
|
||||
***Note:*** All the callable files provide a detailed parser that describes the behaviour of
|
||||
the program and their inputs. Please, check the operational modes by using the extension ```--help```.
|
||||
|
||||
## Training/Testing
|
||||
The dataloader structure can be found in ```./data/loaders/aligments.py``` and it can be
|
||||
manually controlled by instantiating the class ```AlignmentsDataset()``` or by using
|
||||
the ```data_config``` structure available in ```./data/loaders/dl_config.py```.
|
||||
|
||||
Each image sample will follow the next configuration:
|
||||
```
|
||||
sample = {'image': Data augmented crop image,
|
||||
'sample_idx': Image ID,
|
||||
'imgpath': Absolute path to raw image,
|
||||
'imgpath_local': Relative path to raw image,
|
||||
'ids_ldm': Landmarks ids,
|
||||
'bbox': Face bbox [x,y,w,h] (ref crop),
|
||||
'bbox_raw': Face bbox [x,y,w,h] (ref image),
|
||||
'landmarks': Augmented landmarks [[x1,y1], [x2,y2], ...] (ref crop)
|
||||
'visible': Visibilities [0,1, ...] (1 == Visible)
|
||||
'mask_ldm': Available landmarks anns [True, False, ...] <- len(ids_ldm)
|
||||
'headpose': Augmented POSIT headpose [yaw, pithc, roll]
|
||||
}
|
||||
|
||||
Extra features while debugging:
|
||||
sample = { ...
|
||||
'landmarks_ori' = Ground truth landmarks before augmentation (ref image)
|
||||
'visible_ori' = Ground truth visibilities before augmentation
|
||||
'mask_ldm_ori' = Ground truth mask before augmentation
|
||||
'headpose_ori' = Ground truth headpose before augmentation (if available)
|
||||
}
|
||||
```
|
||||
|
||||
## Visualizers
|
||||
The dataloader framework provides complementary visualizers to further understand the databases,
|
||||
datasets and their difficulties:
|
||||
|
||||
* ```./data/visualize/inspect_dataset.py```
|
||||
Focus on the database annotations and the data augmentations, which allows us to
|
||||
understand the training, validation and test datasets.
|
||||
|
||||
* ```./data/visualize/inspect_heatmaps.py```
|
||||
Extended visualizer version focus on understanding the heatmaps and boundaries features used for training.
|
||||
|
||||
* ```./data/model3D/visualization.py```
|
||||
Visualize the rigid facial 3D models used by SPIGA to project the initial coordinates of the GAT regressor.
|
||||
@@ -0,0 +1,181 @@
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
from ..loaders import dl_config as dl_cfg
|
||||
from ..loaders import dataloader as dl
|
||||
from ..visualize import plotting as plot
|
||||
|
||||
|
||||
def inspect_parser():
|
||||
import argparse
|
||||
pars = argparse.ArgumentParser(description='Data augmentation and dataset visualization. '
|
||||
'Press Q to quit,'
|
||||
'N to visualize the next image'
|
||||
' and any other key to visualize the next default data.')
|
||||
pars.add_argument('database', type=str,
|
||||
choices=['wflw', '300wpublic', '300wprivate', 'cofw68', 'merlrav'], help='Database name')
|
||||
pars.add_argument('-a', '--anns', type=str, default='train', help='Annotation type: test, train or valid')
|
||||
pars.add_argument('-np', '--nopose', action='store_false', default=True, help='Avoid pose generation')
|
||||
pars.add_argument('-c', '--clean', action='store_true', help='Process without data augmentation for train')
|
||||
pars.add_argument('--shape', nargs='+', type=int, default=[256, 256], help='Image cropped shape (W,H)')
|
||||
pars.add_argument('--img', nargs='+', type=int, default=None, help='Select specific image ids')
|
||||
return pars.parse_args()
|
||||
|
||||
|
||||
class DatasetInspector:
|
||||
|
||||
def __init__(self, database, anns_type, data_aug=True, pose=True, image_shape=(256,256)):
|
||||
|
||||
data_config = dl_cfg.AlignConfig(database, anns_type)
|
||||
data_config.image_size = image_shape
|
||||
data_config.ftmap_size = image_shape
|
||||
data_config.generate_pose = pose
|
||||
|
||||
if not data_aug:
|
||||
data_config.aug_names = []
|
||||
|
||||
self.data_config = data_config
|
||||
dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
|
||||
self.dataset = dataset
|
||||
self.dataloader = dataloader
|
||||
self.colors_dft = {'lnd': (plot.GREEN, plot.RED), 'pose': (plot.BLUE, plot.GREEN, plot.RED)}
|
||||
|
||||
def show_dataset(self, ids_list=None):
|
||||
|
||||
if ids_list is None:
|
||||
ids = self.get_idx(shuffle=self.data_config.shuffle)
|
||||
else:
|
||||
ids = ids_list
|
||||
|
||||
for img_id in ids:
|
||||
data_dict = self.dataset[img_id]
|
||||
crop_imgs, full_img = self.plot_features(data_dict)
|
||||
|
||||
# Plot crop
|
||||
if 'merge' in crop_imgs.keys():
|
||||
crop = crop_imgs['merge']
|
||||
else:
|
||||
crop = crop_imgs['lnd']
|
||||
cv2.imshow('crop', crop)
|
||||
|
||||
# Plot full
|
||||
cv2.imshow('image', full_img['lnd'])
|
||||
|
||||
key = cv2.waitKey()
|
||||
if key == ord('q'):
|
||||
break
|
||||
|
||||
def plot_features(self, data_dict, colors=None):
|
||||
|
||||
# Init variables
|
||||
crop_imgs = {}
|
||||
full_imgs = {}
|
||||
if colors is None:
|
||||
colors = self.colors_dft
|
||||
|
||||
# Cropped image
|
||||
image = data_dict['image']
|
||||
landmarks = data_dict['landmarks']
|
||||
visible = data_dict['visible']
|
||||
if np.any(np.isnan(visible)):
|
||||
visible = None
|
||||
mask = data_dict['mask_ldm']
|
||||
|
||||
# Full image
|
||||
if 'image_ori' in data_dict.keys():
|
||||
image_ori = data_dict['image_ori']
|
||||
else:
|
||||
image_ori = cv2.imread(data_dict['imgpath'])
|
||||
landmarks_ori = data_dict['landmarks_ori']
|
||||
visible_ori = data_dict['visible_ori']
|
||||
if np.any(np.isnan(visible_ori)):
|
||||
visible_ori = None
|
||||
mask_ori = data_dict['mask_ldm_ori']
|
||||
|
||||
# Plot landmarks
|
||||
crop_imgs['lnd'] = self._plot_lnd(image, landmarks, visible, mask, colors=colors['lnd'])
|
||||
full_imgs['lnd'] = self._plot_lnd(image_ori, landmarks_ori, visible_ori, mask_ori, colors=colors['lnd'])
|
||||
|
||||
if self.data_config.generate_pose:
|
||||
rot, trl, cam_matrix = self._extract_pose(data_dict)
|
||||
|
||||
# Plot pose
|
||||
crop_imgs['pose'] = plot.draw_pose(image, rot, trl, cam_matrix, euler=True, colors=colors['pose'])
|
||||
|
||||
# Plot merge features
|
||||
crop_imgs['merge'] = plot.draw_pose(crop_imgs['lnd'], rot, trl, cam_matrix, euler=True, colors=colors['pose'])
|
||||
|
||||
return crop_imgs, full_imgs
|
||||
|
||||
def get_idx(self, shuffle=False):
|
||||
ids = list(range(len(self.dataset)))
|
||||
if shuffle:
|
||||
random.shuffle(ids)
|
||||
return ids
|
||||
|
||||
def reload_dataset(self, data_config=None):
|
||||
if data_config is None:
|
||||
data_config = self.data_config
|
||||
dataloader, dataset = dl.get_dataloader(1, data_config, debug=True)
|
||||
self.dataset = dataset
|
||||
self.dataloader = dataloader
|
||||
|
||||
def _extract_pose(self, data_dict):
|
||||
# Rotation and translation matrix
|
||||
pose = data_dict['pose']
|
||||
rot = pose[:3]
|
||||
trl = pose[3:]
|
||||
|
||||
# Camera matrix
|
||||
cam_matrix = data_dict['cam_matrix']
|
||||
|
||||
# Check for ground truth anns
|
||||
if 'headpose_ori' in data_dict.keys():
|
||||
if len(self.data_config.aug_names) == 0:
|
||||
print('Image headpose generated by ground truth data')
|
||||
pose_ori = data_dict['headpose_ori']
|
||||
rot = pose_ori
|
||||
|
||||
return rot, trl, cam_matrix
|
||||
|
||||
def _plot_lnd(self, image, landmarks, visible, mask, max_shape_thr=720, colors=None):
|
||||
|
||||
if colors is None:
|
||||
colors = self.colors_dft['lnd']
|
||||
|
||||
# Full image plots
|
||||
W, H, C = image.shape
|
||||
|
||||
# Original image resize if need it
|
||||
if W > max_shape_thr or H > max_shape_thr:
|
||||
max_shape = max(W, H)
|
||||
scale_factor = max_shape_thr / max_shape
|
||||
resize_shape = (int(H * scale_factor), int(W * scale_factor))
|
||||
image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask,
|
||||
thick_scale=1 / scale_factor, colors=colors)
|
||||
image_out = cv2.resize(image_out, resize_shape)
|
||||
else:
|
||||
image_out = plot.draw_landmarks(image, landmarks, visible=visible, mask=mask, colors=colors)
|
||||
|
||||
return image_out
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = inspect_parser()
|
||||
data_aug = True
|
||||
database = args.database
|
||||
anns_type = args.anns
|
||||
pose = args.nopose
|
||||
select_img = args.img
|
||||
if args.clean:
|
||||
data_aug = False
|
||||
|
||||
if len(args.shape) != 2:
|
||||
raise ValueError('--shape requires two values: width and height. Ej: --shape 256 256')
|
||||
else:
|
||||
img_shape = tuple(args.shape)
|
||||
|
||||
visualizer = DatasetInspector(database, anns_type, data_aug=data_aug, pose=pose, image_shape=img_shape)
|
||||
visualizer.show_dataset(ids_list=select_img)
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .inspect_dataset import DatasetInspector, inspect_parser
|
||||
|
||||
|
||||
class HeatmapInspector(DatasetInspector):
|
||||
|
||||
def __init__(self, database, anns_type, data_aug=True, image_shape=(256,256)):
|
||||
|
||||
super().__init__(database, anns_type, data_aug=data_aug, pose=False, image_shape=image_shape)
|
||||
|
||||
self.data_config.aug_names.append('heatmaps2D')
|
||||
self.data_config.heatmap2D_norm = False
|
||||
self.data_config.aug_names.append('boundaries')
|
||||
self.data_config.shuffle = False
|
||||
self.reload_dataset()
|
||||
|
||||
def show_dataset(self, ids_list=None):
|
||||
|
||||
if ids_list is None:
|
||||
ids = self.get_idx(shuffle=self.data_config.shuffle)
|
||||
else:
|
||||
ids = ids_list
|
||||
|
||||
for img_id in ids:
|
||||
data_dict = self.dataset[img_id]
|
||||
|
||||
crop_imgs, _ = self.plot_features(data_dict)
|
||||
|
||||
# Plot landmark crop
|
||||
cv2.imshow('crop', crop_imgs['lnd'])
|
||||
|
||||
# Plot landmarks 2D (group)
|
||||
crop_allheats = self._plot_heatmaps2D(data_dict)
|
||||
|
||||
# Plot boundaries shape
|
||||
cv2.imshow('boundary', np.max(data_dict['boundary'], axis=0))
|
||||
|
||||
for lnd_idx in range(self.data_config.database.num_landmarks):
|
||||
# Heatmaps 2D
|
||||
crop_heats = self._plot_heatmaps2D(data_dict, lnd_idx)
|
||||
maps = cv2.hconcat([crop_allheats['heatmaps2D'], crop_heats['heatmaps2D']])
|
||||
cv2.imshow('heatmaps', maps)
|
||||
|
||||
key = cv2.waitKey()
|
||||
if key == ord('q'):
|
||||
break
|
||||
if key == ord('n'):
|
||||
break
|
||||
|
||||
if key == ord('q'):
|
||||
break
|
||||
|
||||
def _plot_heatmaps2D(self, data_dict, heatmap_id=None):
|
||||
|
||||
# Variables
|
||||
heatmaps = {}
|
||||
image = data_dict['image']
|
||||
|
||||
if heatmap_id is None:
|
||||
heatmaps2D = data_dict['heatmap2D']
|
||||
heatmaps2D = np.max(heatmaps2D, axis=0)
|
||||
else:
|
||||
heatmaps2D = data_dict['heatmap2D'][heatmap_id]
|
||||
|
||||
# Plot maps
|
||||
heatmaps['heatmaps2D'] = self._merge_imgmap(image, heatmaps2D)
|
||||
return heatmaps
|
||||
|
||||
def _merge_imgmap(self, image, maps):
|
||||
crop_maps = cv2.applyColorMap(np.uint8(255 * maps), cv2.COLORMAP_JET)
|
||||
return cv2.addWeighted(image, 0.7, crop_maps, 0.3, 0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = inspect_parser()
|
||||
data_aug = True
|
||||
database = args.database
|
||||
anns_type = args.anns
|
||||
select_img = args.img
|
||||
if args.clean:
|
||||
data_aug = False
|
||||
|
||||
if len(args.shape) != 2:
|
||||
raise ValueError('--shape requires two values: width and height. Ej: --shape 256 256')
|
||||
else:
|
||||
img_shape = tuple(args.shape)
|
||||
|
||||
visualizer = HeatmapInspector(database, anns_type, data_aug, image_shape=img_shape)
|
||||
visualizer.show_dataset(ids_list=select_img)
|
||||
@@ -0,0 +1,97 @@
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from ..loaders.augmentors import utils as dlu
|
||||
|
||||
BLUE = (255, 0, 0)
|
||||
GREEN = (0, 255, 0)
|
||||
RED = (0, 0, 255)
|
||||
PURPLE = (128, 0, 128)
|
||||
|
||||
|
||||
def draw_landmarks(image, landmarks, visible=None, mask=None, thick_scale=1, colors=(GREEN, RED)):
|
||||
# Fix variable
|
||||
thick = int(2 * thick_scale + 0.5)
|
||||
# Initialize variables if need it
|
||||
if visible is None:
|
||||
visible = np.ones(len(landmarks))
|
||||
if mask is None:
|
||||
mask = np.ones(len(landmarks))
|
||||
|
||||
mask = np.array(mask, dtype=bool)
|
||||
visible = np.array(visible, dtype=bool)
|
||||
|
||||
# Clean and split landmarks
|
||||
landmarks = landmarks[mask]
|
||||
visible = visible[mask]
|
||||
ldm_vis = landmarks[visible]
|
||||
not_visible = np.logical_not(visible)
|
||||
ldm_notvis = landmarks[not_visible]
|
||||
|
||||
# Plot landmarks
|
||||
if image.shape[0] == 3:
|
||||
image = image.transpose(1, 2, 0)
|
||||
|
||||
canvas = image.copy()
|
||||
canvas = _write_circles(canvas, ldm_vis, color=colors[0], thick=thick)
|
||||
canvas = _write_circles(canvas, ldm_notvis, color=colors[1], thick=thick)
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def _write_circles(canvas, landmarks, color=RED, thick=2):
|
||||
for xy in landmarks:
|
||||
xy = np.array(xy+0.5, dtype=int)
|
||||
canvas = cv2.circle(canvas, (xy[0], xy[1]), thick, color, -1)
|
||||
return canvas
|
||||
|
||||
|
||||
def plot_landmarks_pil(image, landmarks, visible=None, mask=None):
|
||||
|
||||
# Initialize variables if need it
|
||||
if visible is None:
|
||||
visible = np.ones(len(landmarks))
|
||||
if mask is None:
|
||||
mask = np.ones(len(landmarks))
|
||||
|
||||
mask = np.array(mask, dtype=bool)
|
||||
visible = np.array(visible, dtype=bool)
|
||||
not_visible = np.logical_not(visible)
|
||||
|
||||
# Clean and split landmarks
|
||||
landmarks = landmarks[mask]
|
||||
ldm_vis = landmarks[visible]
|
||||
ldm_notvis = landmarks[not_visible]
|
||||
|
||||
# Plot landmarks
|
||||
if image.shape[0] == 3:
|
||||
image = image.transpose(1, 2, 0)
|
||||
|
||||
plt.imshow(image / 255)
|
||||
plt.scatter(ldm_vis[:, 0], ldm_vis[:, 1], s=10, marker='.', c='g')
|
||||
plt.scatter(ldm_notvis[:, 0], ldm_notvis[:, 1], s=10, marker='.', c='r')
|
||||
plt.show()
|
||||
|
||||
|
||||
def draw_pose(img, rot, trl, K, euler=False, size=0.5, colors=(BLUE, GREEN, RED)):
|
||||
if euler:
|
||||
rot = dlu.euler_to_rotation_matrix(rot)
|
||||
|
||||
canvas = img.copy()
|
||||
rotV, _ = cv2.Rodrigues(rot)
|
||||
points = np.float32([[size, 0, 0], [0, -size, 0], [0, 0, -size], [0, 0, 0]]).reshape(-1, 3)
|
||||
axisPoints, _ = cv2.projectPoints(points, rotV, trl, K, (0, 0, 0, 0))
|
||||
axisPoints = axisPoints.astype(int)
|
||||
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), colors[0], 3)
|
||||
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), colors[1], 3)
|
||||
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), colors[2], 3)
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def enhance_heatmap(heatmap):
|
||||
map_aux = heatmap - heatmap.min()
|
||||
map_aux = map_aux / map_aux.max()
|
||||
map_img = cv2.applyColorMap((map_aux * 255).astype(np.uint8), cv2.COLORMAP_BONE)
|
||||
return map_img
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user