python常用工具类

# 获取当前机器gpu的数量
def get_gpu_count():
    # return len(os.popen("nvidia-smi -L").read().strip().split("n"))
    # num_default = len(os.popen("nvidia-smi -L").read().strip().split("n"))
    pattern = "(GPU d+?): NVIDIA"
    num_default = len(re.findall(pattern, os.popen("nvidia-smi -L").read()))
    if "CUDA_VISIBLE_DEVICES" not in os.environ:
        return num_default
    num_specified = len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))
    return min(num_default, num_specified)

def get_gpu_memory(device_id=0):
    try:
        result = os.popen("nvidia-smi").read()
        pattern = "MiB.+?(d+)+?MiB"
        results = re.findall(pattern, result)
        return int(results[device_id])
    except Exception as e:
        logger.error(e)
        return 0


def get_available_memory(device_id=0):
    try:
        result = os.popen("nvidia-smi").read()
        pattern = ".+?(d+)MiB.+?(d+)?MiB"
        results = re.findall(pattern, result)
        used, total = map(int, results[device_id])
        left = total - used
        return left
    except Exception as e:
        logger.error(e)
        return 0

# 提交结果
def submit_result(api, msg, retry=3, timeout=5):
    i = 0
    while i < retry:
        try:
            r = requests.post(api, json=msg, timeout=timeout)
            logger.info(r.text)
            return
        except Exception as e:
            i += 1
            logger.error(e)

    logger.error("结果提交失败!")
    logger.error(f"{api},{msg}")

def cv2_base64(image):
    base64_str = cv2.imencode('.jpg', image)[1].tostring()
    base64_str = base64.b64encode(base64_str)
    return base64_str


def base64_cv2(base64_str):
    imgString = base64.b64decode(base64_str)
    nparr = np.fromstring(imgString, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
    return image

def align_faces(dets, img_raw):
    '''
    dets对应retinaface的原生结果，img_raw为原始图片
    '''
    def align_face(img, bb, landmark, image_size):
        M = None
        if landmark is not None:
            src = np.array([
                [30.2946, 51.6963],
                [65.5318, 51.5014],
                [48.0252, 71.7366],
                [33.5493, 92.3655],
                [62.7299, 92.2041]], dtype=np.float32)
            if image_size[1] == 112:
                src[:, 0] += 8.0
            dst = landmark.astype(np.float32)

            tform = trans.SimilarityTransform()
            tform.estimate(dst, src)
            M = tform.params[0:2, :]

        if M is None:
            ret = img[bb[1]:bb[3], bb[0]:bb[2], :]
            if len(image_size) > 0:
                ret = cv2.resize(
                    ret, (image_size[1], image_size[0]), interpolation=cv2.INTER_CUBIC)
            return ret
        else:
            warped = cv2.warpAffine(
                img, M, (image_size[1], image_size[0]), borderValue=0.0)
            return warped

    face_boxes = dets[:, :4]
    face_landmarks = dets[:, 5:]
    face_cropped = []
    boxes = []
    # print("len(boxes):", len(face_boxes))
    for i in range(len(face_boxes)):
        face_box = face_boxes[i]
        face_landmark = face_landmarks[i].reshape((5, 2))
        face_aligned = align_face(
            img_raw, face_box, face_landmark, (112, 112))
        face_cropped.append(face_aligned)
        boxes.append(face_box)
    return face_cropped