|
| 1 | +""" |
| 2 | + Code related to face detection and manipulation |
| 3 | +""" |
| 4 | + |
| 5 | +#pip install facenet_pytorch |
| 6 | + |
| 7 | +from facenet_pytorch import MTCNN |
| 8 | +mtcnn = MTCNN(image_size=256, margin=80) |
| 9 | + |
| 10 | +# simplest ye olde trustworthy MTCNN for face detection with landmarks |
| 11 | +def detect(img): |
| 12 | + |
| 13 | + # Detect faces |
| 14 | + batch_boxes, batch_probs, batch_points = mtcnn.detect(img, landmarks=True) |
| 15 | + # Select faces |
| 16 | + if not mtcnn.keep_all: |
| 17 | + batch_boxes, batch_probs, batch_points = mtcnn.select_boxes( |
| 18 | + batch_boxes, batch_probs, batch_points, img, method=mtcnn.selection_method |
| 19 | + ) |
| 20 | + |
| 21 | + return batch_boxes, batch_points |
| 22 | + |
| 23 | +# my version of isOdd, should make a separate repo for it :D |
| 24 | +def makeEven(_x): |
| 25 | + return _x if (_x % 2 == 0) else _x+1 |
| 26 | + |
| 27 | +# the actual scaler function |
| 28 | +def scale(boxes, _img, max_res=1_500_000, target_face=256, fixed_ratio=0, max_upscale=2, VERBOSE=False): |
| 29 | + |
| 30 | + x, y = _img.size |
| 31 | + |
| 32 | + ratio = 2 #initial ratio |
| 33 | + |
| 34 | + #scale to desired face size |
| 35 | + if (boxes is not None): |
| 36 | + if len(boxes)>0: |
| 37 | + ratio = target_face/max(boxes[0][2:]-boxes[0][:2]); |
| 38 | + ratio = min(ratio, max_upscale) |
| 39 | + if VERBOSE: print('up by', ratio) |
| 40 | + |
| 41 | + if fixed_ratio>0: |
| 42 | + if VERBOSE: print('fixed ratio') |
| 43 | + ratio = fixed_ratio |
| 44 | + |
| 45 | + x*=ratio |
| 46 | + y*=ratio |
| 47 | + |
| 48 | + #downscale to fit into max res |
| 49 | + res = x*y |
| 50 | + if res > max_res: |
| 51 | + ratio = pow(res/max_res,1/2); |
| 52 | + if VERBOSE: print(ratio) |
| 53 | + x=int(x/ratio) |
| 54 | + y=int(y/ratio) |
| 55 | + |
| 56 | + #make dimensions even, because usually NNs fail on uneven dimensions due skip connection size mismatch |
| 57 | + x = makeEven(int(x)) |
| 58 | + y = makeEven(int(y)) |
| 59 | + |
| 60 | + size = (x, y) |
| 61 | + |
| 62 | + return _img.resize(size) |
| 63 | + |
| 64 | +""" |
| 65 | + A useful scaler algorithm, based on face detection. |
| 66 | + Takes PIL.Image, returns a uniformly scaled PIL.Image |
| 67 | +
|
| 68 | + boxes: a list of detected bboxes |
| 69 | + _img: PIL.Image |
| 70 | + max_res: maximum pixel area to fit into. Use to stay below the VRAM limits of your GPU. |
| 71 | + target_face: desired face size. Upscale or downscale the whole image to fit the detected face into that dimension. |
| 72 | + fixed_ratio: fixed scale. Ignores the face size, but doesn't ignore the max_res limit. |
| 73 | + max_upscale: maximum upscale ratio. Prevents from scaling images with tiny faces to a blurry mess. |
| 74 | +""" |
| 75 | + |
| 76 | +def scale_by_face_size(_img, max_res=1_500_000, target_face=256, fix_ratio=0, max_upscale=2, VERBOSE=False): |
| 77 | + boxes = None |
| 78 | + boxes, _ = detect(_img) |
| 79 | + if VERBOSE: print('boxes',boxes) |
| 80 | + img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE) |
| 81 | + return img_resized |
| 82 | + |
0 commit comments