Python常用工具函数汇总

常用工具函数汇总

[TOC]

绘制矩形框

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# -------------------------------- #
# 绘制矩形框
# -------------------------------- #
def drawn_rectangle(img_path, bboxs, map_bboxs):
    # bboxs = [x1, y1, x2, y2]

    save_path = '/home/hjy/DeepLearning/OC_SORT/YOLOX_outputs/dance_val_eval/vis_embedding/'
    img_path = '/data/hjy/Datasets/DanceTrack/images/val/' + img_path[0]
    img0 = cv2.imread(img_path)

    # 绘制GT
    bboxs = np.array(bboxs, dtype=np.int32)
    for bbox in bboxs:
        xy1 = (bbox[0], bbox[1])
        xy2 = (bbox[2], bbox[3])
        cv2.rectangle(img0, xy1, xy2, (50,205,50), thickness=2)

    # 绘制映射后的bbox
    for bbox in map_bboxs:
        xy1 = (bbox[0], bbox[1])
        xy2 = (bbox[2], bbox[3])
        cv2.rectangle(img0, xy1, xy2, (0,191,255), thickness=2)
    save_path = save_path + img_path.split('/')[-1]
    cv2.imwrite(save_path, img0)
    print(save_path)
    print('writting .......')

可视化Heatmap

1
2
3
4
5
6
7
8
from PIL import Image
import matplotlib.pyplot as plt

def vis_heatmap(heatmaps):
    for ii, heatmap in enumerate(heatmaps):
        heatmap = heatmap.cpu().numpy()
        plt.imshow(heatmap, cmap=plt.cm.jet)
        plt.show()

相关滤波

1
2
3
4
5
import torch.nn.functional as F
    res = []
    for ii in range(len(c_x)):
        out = F.conv2d(id_features, f).view(152, 272)
        res.append(out)

边框等比缩放并填充

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def letterbox(img, height=608, width=1088,
              color=(127.5, 127.5, 127.5)):  # resize a rectangular image to a padded rectangular
    shape = img.shape[:2]  # shape = [height, width]
    ratio = min(float(height) / shape[0], float(width) / shape[1])
    new_shape = (round(shape[1] * ratio), round(shape[0] * ratio))  # new_shape = [width, height]
    print(new_shape)
    dw = (width - new_shape[0]) / 2  # width padding
    dh = (height - new_shape[1]) / 2  # height padding
    top, bottom = round(dh - 0.1), round(dh + 0.1)
    left, right = round(dw - 0.1), round(dw + 0.1)

    img = cv2.resize(img, new_shape, interpolation=cv2.INTER_AREA)  # resized, no border
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # padded rectangular
    return img, ratio, dw, dh

OpenCV读取显示视频

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import numpy as np
import cv2

def check_video():

    cap = cv2.VideoCapture('videos/test_dance.avi')
    print('video W:{}, video H:{}, video fps:{}, video count:{}'.format(
        cap.get(cv2.CAP_PROP_FRAME_WIDTH),
        cap.get(cv2.CAP_PROP_FRAME_HEIGHT),
        cap.get(cv2.CAP_PROP_FPS),
        cap.get(cv2.CAP_PROP_FRAME_COUNT)

    ))
    while(cap.isOpened()):
        ret, frame = cap.read()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        cv2.imshow('frame',gray)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

OpenCV图片生成视频

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def visualization_one_test_video():
    PATH = '/data/hjy/Datasets/DanceTrack/train2/dancetrack0066/img1'
    FPS = 20
    SIZE = (1920, 1080)
    # SIZE = (1280, 720)
    save_path = './vis_result/' + 'test_dance.avi'
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    videoWriter = cv2.VideoWriter(save_path,fourcc,FPS, SIZE, True)#最后一个是保存图片的尺寸
    for f in tqdm.tqdm(sorted(glob.glob(PATH + '/*'))):
        frame = cv2.imread(f)
        videoWriter.write(frame)

    videoWriter.release()
    cv2.destroyAllWindows()

Parser 简易使用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import argparse 

#定义一个求和函数
def sum(x,y):
    return x+y

#创建解析器
parser = argparse.ArgumentParser(description='test')
#添加参数
parser.add_argument('--x', type=int, default=10, help='Random seed.')
parser.add_argument('--y', type=int, default=20, help='Random seed.')
#解析参数
args = parser.parse_args()
#使用参数
print(sum(args.x,args.y)) #结果：30

绘制跟踪结果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def get_color(idx):
    idx = idx * 3
    color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
    return color



def plot_tracking(image, tlwhs, obj_ids, scores=None, frame_id=0, fps=0., ids2=None):
    im = np.ascontiguousarray(np.copy(image))
    im_h, im_w = im.shape[:2]
    top_view = np.zeros([im_w, im_w, 3], dtype=np.uint8) + 255

    #text_scale = max(1, image.shape[1] / 1600.)
    #text_thickness = 2
    #line_thickness = max(1, int(image.shape[1] / 500.))
    text_scale = 2
    text_thickness = 2
    line_thickness = 3

    radius = max(5, int(im_w/140.))
    cv2.putText(im, 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
                (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), thickness=2)

    for i, tlwh in enumerate(tlwhs):
        x1, y1, w, h = tlwh
        intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
        obj_id = int(obj_ids[i])
        id_text = '{}'.format(int(obj_id))
        if ids2 is not None:
            id_text = id_text + ', {}'.format(int(ids2[i]))
        color = get_color(abs(obj_id))
        cv2.rectangle(im, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness)
        cv2.putText(im, id_text, (intbox[0], intbox[1]), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255),
                    thickness=text_thickness)
    return im

生成树状目录结构

版本1

代码来源暂时找到不了，也是CSDN上一个博主的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# tree.py
import re
from pathlib import Path
from pathlib import WindowsPath
from typing import Optional, List


class DirectionTree:
    def __init__(self,
                 direction_name: str = 'WorkingDirection',
                 direction_path: str = '.',
                 ignore_list: Optional[List[str]] = None):
        self.owner: WindowsPath = Path(direction_path)
        self.tree: str = direction_name + '/\n'
        self.ignore_list = ignore_list
        if ignore_list is None:
            self.ignore_list = []
        self.direction_ergodic(path_object=self.owner, n=0)

    def tree_add(self, path_object: WindowsPath, n=0, last=False):
        if n > 0:
            if last:
                self.tree += '│' + ('    │' * (n - 1)) + '    └────' + path_object.name
            else:
                self.tree += '│' + ('    │' * (n - 1)) + '    ├────' + path_object.name
        else:
            if last:
                self.tree += '└' + ('──' * 2) + path_object.name
            else:
                self.tree += '├' + ('──' * 2) + path_object.name
        if path_object.is_file():
            self.tree += '\n'
            return False
        elif path_object.is_dir():
            self.tree += '/\n'
            return True

    def filter_file(self, file):
        for item in self.ignore_list:
            if re.fullmatch(item, file.name):
                return False
        return True

    def direction_ergodic(self, path_object: WindowsPath, n=0):
        dir_file: list = list(path_object.iterdir())
        dir_file.sort(key=lambda x: x.name.lower())
        dir_file = [f for f in filter(self.filter_file, dir_file)]
        for i, item in enumerate(dir_file):
            if i + 1 == len(dir_file):
                if self.tree_add(item, n, last=True):
                    self.direction_ergodic(item, n + 1)
            else:
                if self.tree_add(item, n, last=False):
                    self.direction_ergodic(item, n + 1)


if __name__ == '__main__':
    i_l = [
        '\.git', '__pycache__', '', 'venv', '\.idea', '.+\.jpg', '.+\.png',
        'image', 'css', 'admin', 'beizhu.txt',
          ]
    tree = DirectionTree(ignore_list=i_l, direction_path='/home/hjy/DeepLearning/summer_practice_baseline/')
    print(tree.tree)

输出样式：

image-20220813145725737

版本2

代码来源：https://blog.csdn.net/zichen_ziqi/article/details/107775556

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 显示文件夹树状目录
import os
import os.path
 
def dfs_showdir(path, depth):
    if depth == 0:
        print("root:[" + path + "]")
 
    for item in os.listdir(path):
        if '.git' not in item:
            print("|      " * depth + "+--" + item)
            newitem = path +'/'+ item
            if os.path.isdir(newitem):
                dfs_showdir(newitem, depth +1)
 
 
if __name__ == '__main__':
    path = ""             # 文件夹路径
    dfs_showdir(path, 0)  # 显示文件夹的树状结构

输出样式：

image-20220813145521942

Python判断文件或者目录是否存在，如果不存在则创建

转载至CSDN，原文链接：https://blog.csdn.net/u013247765/article/details/79050947

判断目录是否存在

1
2
3
4
5
import os
dirs = '/Users/joseph/work/python/'

if not os.path.exists(dirs):
    os.makedirs(dirs)

判断文件是否存在

1
2
3
4
5
import os
filename = '/Users/joseph/work/python/poem.txt'

if not os.path.exists(filename):
    os.system(r"touch {}".format(path))#调用系统命令行来创建文件