Untitled

import cv2
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
from p_tqdm import p_map

app = FaceAnalysis()
app.prepare(ctx_id=0, det_size=(1024,1024))

def save_cropped_face(cname):
    faces = app.get(cname)
    if len(faces) != 0:
        return
    im = cv2.imread(cname)
    box = faces[0].bbox
    angle = faces[0].bbox
    
    x1 = box[0]
    y1 = box[1]
    x2 = box[2]
    y2 = box[3]
    center = [(x2+x1)/2, (y2+y1)/2]
    w = x2 - x1
    h = y2 - y1
    d = max(w,h)
    delta_d = 1.5
    
    if angle[1] < -50:
        x1 = max(center[0] - d*0.55, 0)
        delta_d = 1.8
    elif angle[1] > 50:
        x1 = max(center[0] - d*1.25, 0)
        delta_d = 1.8
    else:
        x1 = max(center[0] - d*0.75, 0)
        
    x2 = x1 + d*delta_d

    if delta_d == 1.5:
        y1 = max(center[1] - d*0.875, 0)
    else:
        y1 = max(center[1] - d*1.05, 0)
        
    y2 = y1 + d*delta_d
    small_im = im[int(y1):int(y2), int(x1):int(x2), :]
    
    cname = 'face_crops/model_' + cname.split('/')[-1]
    small_im = cv2.resize(small_im, (512, 512))
    cv2.imwrite(cname, small_im)
        
        

info = []

for name in names:
    info.append(name)

p_map(info)