脚本集

发布时间 2024-01-09 17:37:13作者: scrutiny-span

批量修改YOLO格式的txt文件内的类别值

  1. 脚本如下:
import fileinput
from pathlib import Path

'''
修改txt文件的类别值,例如0 0.407532 0.697235 0.13095 0.59244更改为1 0.407532 0.697235 0.13095 0.59244
:param path: 标注文件所在的目录
:param dict_map: 一个包含了修改关系的字典
'''
def alterClassName(path='.', dict_map = {}):
    files = Path(path).glob('*.txt')
    with fileinput.input(files, inplace=True) as f:
        for line in f:
            className, separator, boxValue = line.partition(' ')
            line = f'{dict_map.get(className, className)}{separator}{boxValue}'
            print(line, end='')



if __name__ == '__main__':
    dict_map = {'0': '1'}   # 将txt文件内的类别值从0改为1
    alterClassName('C:/Users/root/test/', dict_map)

将给定的一个视频文件或者一个实时流地址切割成帧

  1. 脚本如下:
import cv2
from loguru import logger

'''
将给定的视频文件或者实时流地址按照客户指定的截取间隔切割成图片
:param source:待切割的原视频路径
:param frameRate: 截取间隔,默认为25表示每隔25帧截取一帧
:param destination:保存切割图片的路径
'''
def cuttingVideoToImage(source = '.', frameRate = 25, destination = '.'):
	cap = cv2.VideoCapture(source)
	if (False == cap.isOpened()):
		logger.error(f"can not open {source}")
		exit(-1)

	count = 1
	while(True):
		ret, frame = cap.read()
		if ret:
			if(count % frameRate == 0):
				logger.info(f"开始截取视频第{str(count)}帧")
				if (destination.endswith('/') == False):
					destination = destination + '/'
				cv2.imwrite(destination + str(count) + '.jpg', frame)  # 这里是将截取的图像保存在本地
			count += 1
			cv2.waitKey(0)
		else:
			logger.info("视频文件已经播放完毕")
			break
	cap.release()


if __name__ == '__main__':
	# cuttingVideoToImage('C:/Users/root/Desktop/test+data/3109fcd7aaa941b4ae489d3466d9e088.mp4', 20, 'C:/Users/root/Desktop/test')
	cuttingVideoToImage('rtsp://xxx', 25, 'C:/Users/root/Desktop/test')

文件重命名

  1. 脚本如下:
import os
from loguru import logger

'''
对指定目录下的文件进行重命名
:param path:需要重命名的文件所在的目录
:param count:文件名称,和prefix一起组成文件名'prefix{count}'
:param prefix:文件名称
:param fileType:文件后缀
'''
def fileRename(path = '.', count = 0, prefix = '', fileType = 'jpg'):
    file_names = os.listdir(path)             
    for file in file_names:  
        old_path_name = os.path.join(path, file)     
        if os.path.isdir(old_path_name):             
            continue

        new_path_name = os.path.join(path, prefix + str(count) + '.' + fileType)  
        logger.info(f'old_path_name:{old_path_name}---->new_path_name:{new_path_name}')
        
        if os.path.exists(new_path_name):
            logger.error(f'{path}目录下已经存在同名的文件,请小心')
            exit(-1)
        
        os.rename(old_path_name, new_path_name)    
        count += 1                 

if __name__ == '__main__':
    # fileRename(r'C:\Users\root\Desktop\test', count=1, prefix='b', fileType='png')
    fileRename(r'C:\Users\root\Desktop\test', count=100, prefix='b')

统计指定目录下的文件个数

  1. 脚本如下:
import os
from loguru import logger

'''
统计指定目录下的指定文件后缀的文件个数
:param path:目录
:param suffix:文件后缀
:param recursive:是否递归子目录
:return count:文件个数
'''
def statFileCount(path = '.', suffix = 'txt', recursive = True):
    count = 0
    for dirpath, dirnames, filenames in os.walk(path):    #遍历统计
        # logger.info(f'dirpath:{dirpath},dirname:{dirnames},filenames:{filenames}')
        for filename in filenames:
            lst = filename.split('.')
            if lst[1] == suffix:
                count += 1   
        if recursive == False:
            break
    return count


if __name__ == '__main__':
    count = statFileCount(r'C:\Users\root\Desktop\test', suffix='jpg', recursive=False)
    logger.info(f'file count:{count}')

读取单个rtsp流地址,并使用YOLOv8模型进行推理预测

  1. 脚本如下:
import cv2
import queue
import threading
from loguru import logger

from ultralytics import YOLO

q=queue.Queue()

model = YOLO('weights/xxx.pt')  

def Receive(streamUrl):
    cap = cv2.VideoCapture(streamUrl)
    if (False == cap.isOpened()):
        logger.error(f"can not open {streamUrl}")
        exit(-1)
    
    count = 0
    while True:
        ret, frame = cap.read()
        if ret:
            count = 0
            q.put(frame)
        else:
            count = count + 1
            if (count > 25):
                # 连续25帧读取不到则重新初始化cap
                cap = cv2.VideoCapture(streamUrl)
            continue
    
    # cap.release()
    
 
def Display():
    while True:
        if q.empty() != True:
            frame=q.get()
            model.predict(source=frame, conf = 0.35,save = True, save_txt=True)
        
        # cv2.imshow("frame", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cv2.destroyAllWindows()

if __name__=='__main__':
    streamUrl = "rtsp://xxx"

    p1=threading.Thread(target=Receive, args = (streamUrl,))
    p2 = threading.Thread(target=Display)
    
    p1.start()
    p2.start()

本文持续更新...........