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

def try_except(func):
    def handler(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            logger.error(e)
            return None

    return handler

保存视频

class MyVideoWriter:
    def __init__(self, save_path, height=None, width=None, fps=25):
        self.save_path = save_path
        self.height = height
        self.width = width
        self.fourcc = cv2.VideoWriter_fourcc(*'MP4V')
        self.cur_frame_id = -1
        self.fps = fps
        self.videoWriter = None

    def save(self, frame):
        self.cur_frame_id += 1
        if self.cur_frame_id == 0:
            if self.height is None or self.width is None:
                self.height, self.width, _ = frame.shape
            self.videoWriter = cv2.VideoWriter(
                self.save_path, self.fourcc, self.fps, (self.width, self.height))
        self.videoWriter.write(frame)

    def release(self):
        if self.videoWriter:
            self.videoWriter.release()
            self.videoWriter = None

    def __del__(self):
        self.release()


class MyVideoWiterImgIo:
    def __init__(self, save_path, fps=25):
        self.videowriter = imageio.get_writer(save_path, fps=fps)

    def save(self, frame):
        self.videowriter.append_data(frame[..., ::-1])

    def __del__(self):
        self.close()

    def close(self):
        if self.videowriter is not None:
            self.videowriter.close()
            self.videowriter = None

图片画多个roi

from itertools import chain
import os
import cv2
import numpy as np

cur_folder = os.path.dirname(__file__)


class RoiDrawer:
    def __init__(self, img):
        self.img = img

    def draw(self):
        def onMouse(event, x, y, flags, param):
            if event == cv2.EVENT_LBUTTONDOWN:
                cache_imgs.append(img.copy())
                cv2.circle(img, (x, y), 3, (0, 0, 255), -1)
                if len(polygon) != 0:
                    last_pt = polygon[-1]
                    cv2.line(img, last_pt, [x, y], (255, 0, 0), 1)
                polygon.append([x, y])
                cv2.imshow(winname, img)

        cache_imgs = []
        polygons = []
        polygon = []
        winname = "ori"
        cv2.namedWindow(winname)
        cv2.setMouseCallback(winname, onMouse)
        img = self.img.copy()
        img_h, img_w, *_ = img.shape

        while True:
            cv2.imshow(winname, img)
            key = cv2.waitKey()

            if key == 27:
                if len(polygon) > 0:
                    polygon.pop(-1)
                    img = cache_imgs.pop(-1)

            elif key == 13:  # enter
                pt1 = polygon[0]
                pt2 = polygon[-1]
                cv2.line(img, pt1, pt2, (255, 0, 0), 1)
                contour = np.array(polygon) / [img_w, img_h] * [640, 480]
                contour = contour.astype(np.int32).tolist()
                contour = list(chain.from_iterable(contour))
                contour = ",".join(map(str, contour)) + ','
                polygons.append({"Id": len(polygons), "Point": contour})
                cache_imgs.clear()
                polygon.clear()
                print(polygons)

            elif key == ord('q'):
                break

        cv2.destroyWindow(winname)
        return {"Polygons": polygons}


if __name__ == '__main__':

    root_folder = "test_imgs"
    img_names = os.listdir(root_folder)
    img_names.sort()

    for img_name in img_names:
        print(img_name)
        frame = cv2.imread(os.path.join(root_folder, img_name), cv2.IMREAD_COLOR)
        drawer = RoiDrawer(frame)
        polygons = drawer.draw()
        print(polygons)
        # for polygon in polygons:
        #     print(polygon)

        # # polygon = [[536, 188], [1239, 697], [1399, 906], [1018, 690], [537, 277]]
        #
        # img_h, img_w, *_ = img.shape
        # contour = np.array(polygon) / [img_w, img_h] * [640, 480]
        # contour = contour.astype(np.int32).tolist()
        # contour = list(chain.from_iterable(contour))
        # contour = ",".join(map(str, contour)) + ','
        # mon_parm = {"Polygons": [{"Id": 999999, "Point": contour}]}
        # result = processTest(mon_parm, None, None, frame)

操作minio

import minio
import os
import cv2
from loguru import logger


class MinioClient:
    def __init__(self,
                 endpoint="192.168.4.78:9000",
                 access_key="minioadmin",
                 secret_key="minioadmin",
                 secure=False
                 ):
        self.endpoint = endpoint
        try:
            self.client = minio.Minio(
                endpoint=endpoint, access_key=access_key, secret_key=secret_key, secure=secure)
        except Exception as e:
            self.client = None
            logger.error(e)

    def upload(self, bucket, file_path, dst_path=None, retry_times=3):
        if self.client is None:
            return False

        try:
            if not self.client.bucket_exists(bucket):
                self.client.make_bucket(bucket)
        except Exception as e:
            logger.error("创建bucket失败")
            logger.error(e)
            return False

        filename = os.path.basename(file_path)
        if dst_path is None:
            dst_path = filename

        for _ in range(retry_times):
            try:
                self.client.fput_object(
                    bucket, dst_path, file_path)
                return True

            except Exception as e:
                logger.error(f"{file_path} 上传失败")
                logger.error(e)

        return False

    def download(self, bucket, filename, dst_folder="/tmp", retry_times=3, only_path=True):
        if self.client is None:
            return False, ""

        if only_path:
            dst_path = os.path.join(self.endpoint, bucket, filename)
            dst_path = dst_path.replace("\\", "/")
            if "http" not in dst_path:
                dst_path = "http://" + dst_path
            return True, dst_path

        dst_path = os.path.join(dst_folder, filename)

        for _ in range(retry_times):
            try:
                self.client.fget_object(bucket, filename, dst_path)
                return True, dst_path

            except Exception as e:
                logger.error(f"{dst_path} 下载失败")
                logger.error(e)

        return False, ""


if __name__ == '__main__':

    client = MinioClient()

    file_path = "/workspace_wjr/develop/projects/continuous_model_server/test.py"
    # s = time()
    print(client.upload("test",file_path))
    # print(time()-s)

    # ret, dst_path = client.download("test", "1.mp4",only_path=False)
    # print(dst_path)
    # cap = cv2.VideoCapture(dst_path)
    # # while cap.isOpened():
    # #     ret, frame = cap.read()
    # #     cv2.imshow("frame", frame)
    # #     cv2.waitKey(5)
    # logger.info(cap.isOpened())

    # from utils.db_utils import load_config

    # config = load_config()
    # minio_config = config["minio"]

    # client = MinioClient(minio_config["endpoint"],
    #                      minio_config["access_key"],
    #                      minio_config["secret_key"],
    #                      minio_config["secure"])

    # client.upload("test", "minio_client.py")
 

操纵mq

import pika
from loguru import logger
from pika.exceptions import AMQPConnectionError
from container_utils.common import try_except


class MqSender:

    @try_except
    def __init__(self, config):
        host = config.get("host", "127.0.0.1")
        username = config.get("username", "admin")
        pwd = config.get("pwd", "admin")
        port = config.get("port", 5672)
        virtual_host = config.get("virtual_host", "/")
        target_queue = config.get("target_queue", "test_queue")

        self.target_queue = target_queue

        try:
            credentials = pika.PlainCredentials(username, pwd)
            self.conn = pika.BlockingConnection(
                parameters=pika.ConnectionParameters(host=host, port=port, virtual_host=virtual_host,
                                                     credentials=credentials))
            self.channel = self.conn.channel()

            result = self.channel.queue_declare(queue=self.target_queue)  # exclusive=True 会在mq断开时自动删除队列
            # self.queue_name = result.method.queue
        except AMQPConnectionError as e:
            logger.error("mq连接异常！")
            for se in e.args:
                logger.error(se.exception)
            self.conn = None
            self.channel = None


    def __call__(self, body):
        if self.channel is None or self.conn is None:
            logger.warning("mq写入异常！")
            logger.warning(f"{self.target_queue}: {body}")
            return False

        try:
            self.channel.basic_publish(
                exchange="",
                routing_key=self.target_queue,
                body=body.encode(encoding="utf-8")
            )
            return True

        except Exception as e:
            logger.error("mq写入异常！")
            logger.error(e)
            logger.error(f"{self.target_queue}: {body}")
            return False

    def __del__(self):
        self.close()

    def close(self):
        if self.conn is not None:
            self.conn.close()
            self.conn = None

if __name__ == '__main__':
    from container_utils.db_utils import load_config
    import datetime
    import json

    mq_config = load_config("../container_config/config.yml").get("mq", {"host": "127.0.0.1", "port": 5672, "username": "admin",
                                                                     "pwd": "admin"})

    sender = MqSender("test_ch", mq_config)
    for content in ["good morning!", "good afternoon!", "good evening!"]:
        # 处理完一条才会处理下一条
        sender(content)

操作mysql

import pymysql
from loguru import logger


class MysqlDb:
    def __init__(self, config):
        host = config.get("host", "127.0.0.1")
        user = config.get("user", "root")
        pwd = config.get("pwd", "admin")
        port = config.get("port", 3306)
        db = config.get("db", "mvp")

        # 如果不带cursorclass，返回值是一个元组
        try:
            self.db = pymysql.connect(host=host, user=user, password=pwd, database=db,
                                      port=port, charset="utf8mb4", cursorclass=pymysql.cursors.DictCursor,autocommit=True)
        except Exception as e:
            logger.error(f"mysql连接异常！")
            logger.error(e)
            self.db = None

    def __del__(self):
        self.close()
    
    def close(self):
        if self.db is not None:
            self.db.close()
            self.db = None

    # 查询数据库，type表示每次查询一条还是查询所有
    def query(self, cmd, type="one"):
        if self.db is None:
            logger.warning("mysql连接异常！")
            return None
        try:
            with self.db.cursor() as cur:
                row = cur.execute(cmd)
                if type == "one":
                    data = cur.fetchone()
                else:
                    data = cur.fetchall()
                return data
        except Exception as e:
            logger.error(e)
            return None

    # 增/删/改数据库，args不为None表示一次更新多条数据
    def update(self, cmd, args=None):
        if self.db is None:
            logger.warning("mysql连接异常！")
            return False
        try:
            with self.db.cursor() as cur:
                if args is not None:
                    cur.executemany(cmd, args)
                else:
                    cur.execute(cmd)
                self.db.commit()
                return True
        except Exception as e:
            logger.error(e)
            self.db.rollback()
            return False


if __name__ == '__main__':
    from container_utils.db_utils import load_config
    import datetime
    import json

    mysql_config = load_config("../container_config/config.yml").get("mysql", {"host": "127.0.0.1", "port": 3306, "user": "admin",
                                                                     "pwd": "", "db": "mvp"})

    db = MysqlDb(mysql_config)

    # 模型服务刚启动，需要查找所有status为0或1的任务
    select_sql = 'SELECT * FROM continuous_tasks where status in(0,1)'
    # # 模型启动后续只需要查找status=0的任务
    # select_sql = 'SELECT * FROM continuous_tasks where status=0'

    # 每次更新一条数据
    result_tmp = [[1,2,3,4,0.5],[5,6,7,8,0.9]]
    resultUrl = "a/b/c/d.avi"
    # 一次更新一条任务
    update_sql1 = f'UPDATE continuous_tasks SET status="1", updateTime="{datetime.datetime.now()}",timestamp="{datetime.datetime.now()}",resultUrl="{resultUrl}", result="{json.dumps(result_tmp)}" WHERE id=4'

    # 一次更新多条任务
    # 要注意的是里面的参数， 不管什么类型，统一使用%s作为占位符
    update_sql2 = 'UPDATE continuous_tasks SET status=%s, updateTime=%s, result=%s,timestamp=%s,resultUrl=%s WHERE id=%s'


    "***************************************"
    data = db.query(select_sql)
    print(data)
    print(">"*30)

    data = db.query(select_sql, "all")
    if data is not None:
        for d in data:
            print(d)
    print(">" * 30)
    "***************************************"
    res = db.update(update_sql1)
    print(res)
    print(">" * 30)

    res = db.update(update_sql2,[(2, datetime.datetime.now(),json.dumps(result_tmp), datetime.datetime.now(),"/data/result/res.mp4",4),
                                 (2, datetime.datetime.now(),json.dumps(result_tmp), datetime.datetime.now(),"/data/result/res2.mp4",5)])
    print(">" * 30)


    #
    # # insert_sql = 'INSERT INTO continuous_tasks(id, username, password) VALUES(11, "王五", "333333")'
    # # delete_sql = 'DELETE FROM continuous_tasks WHERE id = 11'
    #
    # # data = db.query(select_sql)
    # # print(data)
    # data = db.query(select_sql, type="all")
    # for d in data:
    #     print(d)
    # # db.update(update_sql1)
    # # db.update(update_sql2, [("aaa", 123456, 0), ("bbb", 654321, 1)])
    # # db.update(update_sql2,[("smoke",json.dumps({"a":10, "b":20}),4)])
    #
    # # db.update(update_sql1.format("modelName","fire","modelConfig",json.dumps({"a":10, "b":20}),"id", 5))
    # # modelName = "fire"
    # # modelConfig = {'a':10, 'b':20}
    # # db.update(f'UPDATE continuous_tasks SET modelName="{modelName}", modelConfig="{modelConfig}" WHERE id=5')
    # # data = db.query(select_sql, type="all")
    # # print(data)
    # # print(json.loads(data[0]["modelConfig"])["a"])
    # # print(json.loads(data[0]["modelConfig"])["b"])

操作redis

import redis
from loguru import logger

class RedisDb:
    def __init__(self, redis_config):
        host = redis_config.get("host", "127.0.0.1")
        port = redis_config.get("port", 6379)
        pwd = redis_config.get("pwd", "")
        db_id = redis_config.get("db_id", 0)

        try:
            self.conn = redis.Redis(host=host, port=port, db=db_id, password=pwd)
        except Exception as e:
            logger.error(e)
            logger.error("redis连接异常！")
            self.conn = None

    def __del__(self):
        self.close()

    def close(self):
        if self.conn is not None:
            self.conn.close()
            self.conn = None

    def set(self, func_name, *args, **kwargs):
        if self.conn is None:
            return False

        if not hasattr(self.conn, func_name):
            logger.warning(f"redis不存在{func_name}函数")
            return False

        func = getattr(self.conn, func_name)
        try:
            res = func(*args, **kwargs)
            return res
        except Exception as e:
            logger.error(f"redis执行{func_name}错误！")
            logger.error(e)
            return False

    def get(self, func_name, *args, **kwargs):
        if self.conn is None:
            return None

        if not hasattr(self.conn, func_name):
            logger.warning(f"redis不存在{func_name}函数")
            return None

        func = getattr(self.conn, func_name)
        try:
            res = func(*args, **kwargs)
            if res is not None:
                res = res.decode("utf8")
            return res
        except Exception as e:
            logger.error(f"redis执行{func_name}错误！")
            logger.error(e)
            return None


if __name__ == '__main__':
    from container_utils.db_utils import load_config
    redis_config = load_config("../container_config/config.yml").get("redis", {"host": "127.0.0.1", "port": 6379, "pwd": "", "db_id": 0})
    db = RedisDb(redis_config)
    # res = db.conn.setex("name", 60, "wjr")
    # print(db.conn.get("name").decode("utf8"))
    print(db.set("setex", "name", 60, "wjr"))
    print(db.get("get", "name"))
    print(db.get("get", "name2"))

马士兵全栈mysql封装

from pymysql import cursors, connect


class MysqlHelper(object):
    # 数据库链接参数，可以定义多个，比如conn_params1,conn_params2，用于连接多个数据库，在类实例化时指定
    conn_params1 = {'host': '192.168.4.60', 'port': 3306, 'user': 'root',
                    'passwd': 'admin', 'db': 'testdb', 'charset': 'utf8'}
    conn_params2 = {'host': 'localhost', 'port': 3306, 'user': 'root',
                    'passwd': 'root', 'db': 'mytestdb2', 'charset': 'utf8'}

    # 类的构造函数，主要用于类的初始化
    def __init__(self, conn_params):
        self.__host = conn_params['host']
        self.__port = conn_params['port']
        self.__db = conn_params['db']
        self.__user = conn_params['user']
        self.__passwd = conn_params['passwd']
        self.__charset = conn_params['charset']

    # 建立数据库连接和打开游标
    def __connect(self):
        self.__conn = connect(host=self.__host, port=self.__port, db=self.__db,
                              user=self.__user, passwd=self.__passwd, charset=self.__charset,
                              cursorclass=cursors.DictCursor,    # 结果以字典形式返回，否则以元组形式返回
                              #   autocommit=True   # autocommit会自动提交，刷新内存，同步数据库
                              )
        self.__cursor = self.__conn.cursor()

    # 关闭游标和关闭连接
    def __close(self):
        self.__cursor.close()
        self.__conn.close()

    # 取一条数据
    def get_one(self, sql, params):
        result = None
        try:
            self.__connect()
            self.__cursor.execute(sql, params)
            result = self.__cursor.fetchone()
            self.__close()
        except Exception as e:
            print(e)
        return result

    # 取所有数据
    def get_all(self, sql, params):
        lst = []
        try:
            self.__connect()
            self.__cursor.execute(sql, params)
            lst = self.__cursor.fetchall()
            self.__close()
        except Exception as e:
            print(e)
        return lst

    # 增加数据
    def insert(self, sql, params):
        return self.__edit(sql, params)

    # 修改数据
    def update(self, sql, params):
        return self.__edit(sql, params)

    # 删除数据
    def delete(self, sql, params):
        return self.__edit(sql, params)

    def __edit(self, sql, params):
        count = 0
        try:
            self.__connect()
        except Exception as e:
            print(e)
            return count
        try:
            if isinstance(params, list):
                count = self.__cursor.executemany(sql, params)
            else:
                count = self.__cursor.execute(sql, params)
            self.__conn.commit()
            self.__conn.rollback()
            self.__close()
        except Exception as e:
            print(e)

        return count


if __name__ == "__main__":

    conn = MysqlHelper(MysqlHelper.conn_params1)
    # 插入一条数据
    sql = "insert into student values (%s,%s,%s,%s)"
    params1 = [(3, "wjr", "900120", "male"), (4, "wjr", "900120", "male")]
    params2 = (3, "wjr", "900120", "male")
    count = conn.insert(sql, params1)
    print("插入了"+str(count)+"条数据")
    count = conn.insert(sql, params2)
    print("插入了"+str(count)+"条数据")

    # 修改数据
    sql = "update student set name=%s where idx=%s"
    params = ("wm", 1)
    count = conn.update(sql, params)
    print("更新了" + str(count)+"条数据")

    # 删除数据
    sql = "delete from student where name=%s"
    params = ("wm")
    count = conn.delete(sql, params)
    print("删除了"+str(count)+"条数据")

    # 查询一条数据
    sql = "select * from student where name=%s"
    params = ("wjr")
    data = conn.get_one(sql, params)
    print(data)

    # 查询所有数据
    sql = "select * from student where name=%s"
    params = ("wjr")
    data = conn.get_all(sql, params)
    print(data)
    
    # 只查询某些关键字
    sql = "select idx, passwd from student where name=%s"
    params = ("wjr")
    data = conn.get_all(sql, params)
    print(data)