github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

added sort by "face rect size in source image"

Browse source 
small faces from source image will be placed at the end
This commit is contained in:
Colombo 2020-03-03 21:47:49 +04:00
parent 780ad5679f
commit f56d583cb5
2 changed files with 32 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
main.py
View file

@ -112,7 +112,7 @@ if __name__ == "__main__":
p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.set_defaults (func=process_sort)
def process_util(arguments):

35
mainscripts/Sorter.py
View file

@ -1,5 +1,5 @@
import multiprocessing
import math
import math
import multiprocessing
import operator
import os
import sys
@ -12,7 +12,7 @@ import cv2
import numpy as np
from numpy import linalg as npla
from core import imagelib, pathex
from core import imagelib, mathlib, pathex
from core.cv2ex import *
from core.imagelib import estimate_sharpness
from core.interact import interact as io
@ -146,6 +146,32 @@ def sort_by_face_pitch(input_path):
return img_list, trash_img_list
def sort_by_face_source_rect_size(input_path):
io.log_info ("Sorting by face rect size...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
filepath = Path(filepath)
dflimg = DFLIMG.load (filepath)
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
trash_img_list.append ( [str(filepath)] )
continue
source_rect = dflimg.get_source_rect()
rect_area = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))
img_list.append( [str(filepath), rect_area ] )
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
class HistSsimSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
#override
@ -837,6 +863,7 @@ sort_func_methods = {
'blur': ("blur", sort_by_blur),
'face-yaw': ("face yaw direction", sort_by_face_yaw),
'face-pitch': ("face pitch direction", sort_by_face_pitch),
'face-source-rect-size' : ("face rect size in source image", sort_by_face_source_rect_size),
'hist': ("histogram similarity", sort_by_hist),
'hist-dissim': ("histogram dissimilarity", sort_by_hist_dissim),
'brightness': ("brightness", sort_by_brightness),
@ -860,7 +887,7 @@ def main (input_path, sort_by_method=None):
io.log_info(f"[{i}] {desc}")
io.log_info("")
id = io.input_int("", 3, valid_list=[*range(len(key_list))] )
id = io.input_int("", 4, valid_list=[*range(len(key_list))] )
sort_by_method = key_list[id]
else: