图像畸变-径向切向畸变实时图像RTSP推流

实验环境

注意：ffmpeg进程stdin写入两张图片的时间间隔不能太长，否则mediamtx会出现对应的推流session超时退出。

实验效果

全部代码

my_util.py

#进度条
import os
import sys
import time
import shutil
import logging
import time
from datetime import datetime

def print_progress_bar(iteration, total, prefix='', suffix='', decimals=1, length=100, fill='█', print_end="\r"):
    """
    调用在Python终端中打印自定义进度条的函数
    iteration - 当前迭代（Int）
    total - 总迭代（Int）
    prefix - 前缀字符串（Str）
    suffix - 后缀字符串（Str）
    decimals - 正数的小数位数（Int）
    length - 进度条的长度（Int）
    fill - 进度条填充字符（Str）
    print_end - 行尾字符（Str）
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    print(f'\r{prefix} |{bar}| {percent}% {suffix}', end=print_end)
    # 打印新行，完成进度条
    if iteration == total:
        print()

class Logger(object):
    """
    终端打印不同颜色的日志
    """
    ch = logging.StreamHandler()  # 创建日志处理器对象，在__init__外创建，是类当中的静态属性，不是__init__中的实例属性

    # #创建静态的日志处理器可以减少内存消耗

    # # 创建 FileHandler 实例，指定日志文件路径
    # ch = logging.FileHandler(filename='app1.log')

    def __init__(self):
        self.logger = logging.getLogger()  # 创建日志记录对象
        self.logger.setLevel(logging.DEBUG)  # 设置日志等级info，其他低于此等级的不打印

    def debug(self, message):
        self.fontColor('\033[0;37m%s\033[0m')
        self.logger.debug(message)

    def info(self, message):
        self.fontColor('\033[0;32m%s\033[0m')
        self.logger.info(message)

    def warning(self, message):
        self.fontColor('\033[0;33m%s\033[0m')
        self.logger.warning(message)

    def error(self, message):
        self.fontColor('\033[0;31m%s\033[0m')
        self.logger.error(message)

    def fontColor(self, color):
        formatter = logging.Formatter(color % '%(asctime)s - %(name)s - %(levelname)s - %(message)s')  # 控制日志输出颜色
        self.ch.setFormatter(formatter)
        self.logger.addHandler(self.ch)  # 向日志记录对象中加入日志处理器对象


def delete_files(folder_path, max_files):
    """
    监控指定文件夹中的文件数量，并在超过max_files时删除最旧的文件。
    """
    print("进入删除图片文件夹"+folder_path)
    print("需要删除文件数量")
    print(max_files)
    if True:
        # 获取文件夹中的文件列表
        files = os.listdir(folder_path)
        file_count = len(files)
        print(f"当前文件夹 {folder_path} 中的文件数量: {file_count}")

        # 如果文件数量超过max_files，则删除最旧的文件
        if file_count > max_files:
            # 获取文件夹中所有文件的完整路径，并带上修改时间
            file_paths_with_mtime = [(os.path.join(folder_path, f), os.path.getmtime(os.path.join(folder_path, f))) for
                                     f in files]
            # 按修改时间排序
            sorted_files = sorted(file_paths_with_mtime, key=lambda x: x[1])

            # 删除最旧的文件，直到文件数量在阈值以下
            for file_path, mtime in sorted_files[:file_count - max_files]:
                try:
                    os.remove(file_path)
                    print(f"已删除文件: {file_path}")
                except OSError as e:
                    print(f"删除文件时出错: {e.strerror}")

def copy_file(src, dst):
    shutil.copy2(src, dst)  # copy2会尝试保留文件的元数据


def end_sentence(text, max_length):
    '''
    保证在max_length长度前以句号或点号结束文本
    :param text: 文本
    :param max_length: 最大长度
    :return:
    '''
    # 如果文本长度已经超过最大长度，则直接截断
    if len(text) > max_length:
        text = text[:max_length]

    # print("结果长度 {}".format(len(text)))

    # 查找句号的位置(en)
    period_index = max(text.rfind('.'), text.rfind(','),
                       text.rfind(':'), text.rfind(';'),
                       text.rfind('!'), text.rfind('?'))  # 从后往前找，找到最后一个句号
    # 如果找到了句号且它在最大长度内
    if period_index != -1 and (period_index + 1 < max_length or
                               max_length == -1):
        # 如果需要替换，则替换句号
        text = text[:period_index] + '.'

    # 查找句号的位置(cn)
    period_index = max(text.rfind('。'), text.rfind('，'),
                       text.rfind('：'), text.rfind('；'),
                       text.rfind('！'), text.rfind('？'))  # 从后往前找，找到最后一个句号
    # 如果找到了句号且它在最大长度内
    if period_index != -1 and (period_index + 1 < max_length or
                               max_length == -1):
        # 如果需要替换，则替换句号
        text = text[:period_index] + '。'

    return text


import base64


def encode_base64(input_string):
    """
    对字符串进行Base64编码
    """
    encoded_bytes = base64.b64encode(input_string.encode('utf-8'))
    encoded_string = encoded_bytes.decode('utf-8')
    return encoded_string


def decode_base64(input_string):
    """
    对Base64编码的字符串进行解码
    """
    decoded_bytes = base64.b64decode(input_string.encode('utf-8'))
    decoded_string = decoded_bytes.decode('utf-8')
    return decoded_string


import socket

def get_local_ip():
    try:
        # 创建一个 UDP 套接字
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 连接到一个公共的 IP 地址和端口
        s.connect(("8.8.8.8", 80))
        # 获取本地 IP 地址
        local_ip = s.getsockname()[0]
        s.close()
        return local_ip
    except Exception as e:
        print(f"获取本地 IP 地址时出错: {e}")
        return None

make_pics.py

import numpy as np
import cv2
import math
import time
from PIL import Image, ImageDraw, ImageFont

def distort_image(image, k1, k2, p1, p2):
    """
    对图像应用径向和切向畸变
    :param image: 输入图像
    :param k1: 径向畸变系数
    :param k2: 径向畸变系数
    :param p1: 切向畸变系数
    :param p2: 切向畸变系数
    :return: 畸变后的图像
    """
    h, w = image.shape[:2]

    camera_matrix = np.array([[w, 0, w / 2],
                              [0, h, h / 2],
                              [0, 0, 1]], dtype=np.float32)
    distort_coeffs = np.array([k1, k2, p1, p2, 0], dtype=np.float32)

    # 生成畸变映射
    map1, map2 = cv2.initUndistortRectifyMap(camera_matrix, distort_coeffs, np.eye(3), camera_matrix, (w, h), cv2.CV_32FC1)
    # 应用畸变映射
    distorted_img = cv2.remap(image, map1, map2, cv2.INTER_LINEAR)
    
    return distorted_img

def put_chinese_text(img, text, position, font_path, font_size, color):
    """
    在图像上添加中文文字
    :param img: 输入图像
    :param text: 要添加的文字
    :param position: 文字位置
    :param font_path: 字体文件路径
    :param font_size: 字体大小
    :param color: 文字颜色
    :return: 添加文字后的图像
    """
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    font = ImageFont.truetype(font_path, font_size)
    draw.text(position, text, font=font, fill=color)
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

def make_pics(pic_path="picture.jpg", k1=-0.5, k2=0.0, p1=0.0, p2=0.0):
    # 生成棋盘图像
    chessboard = np.zeros((400, 400, 3), dtype=np.uint8)
    for i in range(0, 400, 40):
        for j in range(0, 400, 40):
            if (i // 40 + j // 40) % 2 == 0:
                chessboard[i:i + 40, j:j + 40] = [255, 255, 255]

    # 生成雷达图
    radar = chessboard.copy()
    x0, y0 = radar.shape[1] // 2, radar.shape[0] // 2
    for radius in range(0, 400, 40):
        cv2.circle(radar, (x0, y0), radius, (0, 0, 255), 1)
    # 绘制径向线
    for angle in range(0, 360, 40):
        # 使用最大半径 400 计算径向线的终点坐标
        x = int(x0 + 400 * math.cos(math.radians(angle)))
        y = int(y0 + 400 * math.sin(math.radians(angle)))
        cv2.line(radar, (x0, y0), (x, y), (0, 0, 255), 1)

    font_size = 15
    font_color = (250, 100, 0)


    combined_distorted_chessboard = distort_image(radar, k1, k2, p1, p2)
    text1 = "k1={:.2f},k2={:.2f},p1={:.2f},p2={:.2f}".format(k1,k2,p1,p2)
    text2 = "图像畸变"
    combined_distorted_chessboard = put_chinese_text(combined_distorted_chessboard, text1, (10, 30), 'simhei.ttf', font_size, font_color)
    combined_distorted_chessboard = put_chinese_text(combined_distorted_chessboard, text2, (10, 60), 'simhei.ttf', font_size, font_color)

    # 保存图像
    cv2.imwrite(pic_path, combined_distorted_chessboard)
    # cv2.imshow(pic_path, combined_distorted_chessboard)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    return

if False:
    for k1 in np.arange(-100,100,0.1):
        for k2 in np.arange(-100, 100, 0.1):
            for p1 in np.arange(-100, 100, 0.1):
                for p2 in np.arange(-100, 100, 0.1):
                    make_pics("picture.jpg", k1, k2, p1, p2)

pic_2_rtsp.py

import numpy as np
import make_pics
import sys
import msvcrt
import subprocess
import time
import shlex
import my_util
from PIL import Image, ImageDraw
import random
import os

log = my_util.Logger()
# RTSP_DEF_IP = "192.168.31.185"
RTSP_DEF_IP = my_util.get_local_ip()
RTSP_PORT = 8554
local_ip = my_util.get_local_ip()
if local_ip:
    RTSP_URL = "rtsp://{}:{}/live".format(local_ip, RTSP_PORT)
else:
    RTSP_URL = "rtsp://{}:{}/live".format(RTSP_DEF_IP, RTSP_PORT)
frame_duration = 1/25  # 每张图片显示的时长（秒），process.stdin.wirte写入速度需要够快，否则可能接收端接受数据不足无法获取解码信息（抓包看看）
frame_num = 0

old_picname = "past.jpg"
new_picname = "now.jpg"
k1 = k2 = p1 = p2 = 0.0


def generate_orig(old_picname):
    """生成默认图片"""
    image = Image.new('RGB', (640, 480), color=(
        random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    draw = ImageDraw.Draw(image)
    draw.text((100, 100), 'No Signal', fill=(255, 255, 255))
    image.save(old_picname)
    return


def generate_image(new_picname, k1, k2, p1, p2):
    """生成图片"""
    make_pics.make_pics(new_picname, k1, k2, p1, p2)
    return


# 构建 ffmpeg 命令使用图片序列
command_line = 'ffmpeg -loglevel error -re -i - -c:v libx264 -pix_fmt yuv420p -r {} -f rtsp {}'.format(1/frame_duration, RTSP_URL)
command_list = shlex.split(command_line)

log.debug(command_list)


def start_process():
    global process
    process = subprocess.Popen(command_list, stdin=subprocess.PIPE, text=False)
    log.info("Process started.")


start_time = time.time()

class QuitException(Exception):
    pass

try:
    # 默认先生成初始图
    step_length = 0.125
    generate_orig(old_picname)
    start_process()
    while True:
        for p1 in np.arange(-1 * step_length, step_length, step_length):
            for p2 in np.arange(-1 * step_length, step_length, step_length):
                for k1 in np.arange(-1, 1, step_length):
                    for k2 in np.arange(-1, 1, step_length):
                        log.debug("畸变系数 k1={}, k2={}, p1={}, p2={}".format(k1, k2, p1, p2))
                        generate_image(new_picname, k1, k2, p1, p2)

                        if msvcrt.kbhit():  # 检查是否有键盘输入
                            input_char = msvcrt.getch().decode('utf-8')
                            if input_char == 'q' or input_char == 'Q':
                                try:
                                    # 向进程的标准输入发送 'q' 并换行
                                    if process.stdin:
                                        process.stdin.write('q\n'.encode())
                                        process.stdin.flush()
                                except Exception as e:
                                    pass
                                raise QuitException()

                        # 持续生成新图片替换旧图片
                        try:
                            if os.path.exists(new_picname):
                                with open(new_picname, 'rb') as f:
                                    process.stdin.write(f.read())
                            else:
                                with open(old_picname, 'rb') as f:
                                    process.stdin.write(f.read())
                        except Exception as e:
                            log.error(f"Error writing to process stdin: {e}")
                            log.info("Restarting process...")
                            process.terminate()
                            try:
                                process.wait(timeout=1)
                            except subprocess.TimeoutExpired:
                                process.kill()
                            start_process()

                        time.sleep(frame_duration)

except QuitException:
    pass
finally:
    try:
        process.terminate()
        try:
            process.wait(timeout=1)
        except subprocess.TimeoutExpired:
            process.kill()
    except Exception:
        pass

    try:
        if os.path.exists(new_picname):
            os.remove(new_picname)
    except Exception as e:
        log.error(f"Error removing {new_picname}: {e}")

    end_time = time.time()
    time_cnt = end_time - start_time
    log.info("FFmpeg进程已执行{}秒并通过输入 'q' 退出。".format(round(time_cnt)))