(三)共用功能函数：AI实时抠图、AI实时抠像、PaddlePaddle模型、虚拟现实视频会议、沉浸式会议场景、人像去背景、图像去背景、视频背景消除
摘要：此文把一些共用的功能函数的源码公开，为适应实时处理视频，这些函数都经过反复的调试、优化、效率提升，使代码效率达到最优。本文与前2篇博文关联性较强，请事先阅读前2篇。
1、显卡检测
AI开发最好使用GPU，效率会提升10倍以上。由于飞桨系统只支持Nvidia显卡，所以要进行显卡信息检测，包括电脑是否有Nvidia卡、显卡数量、显卡型号、显存大小等。获得这些参数，主要是为了软件界面的CPU/GPU下拉框选择时判断这些参数。代码如下：
try: #检测显卡是否是Nvidia，并查看型号是否支持
pynvml.nvmlInit() #显卡初始化
GPU_nums = pynvml.nvmlDeviceGetCount() #显卡数量
handle = pynvml.nvmlDeviceGetHandleByIndex(GPU_nums-1) #显卡句柄
GPU_name = bytes.decode(pynvml.nvmlDeviceGetName(handle)) #显卡名称
GPU_memsize = pynvml.nvmlDeviceGetMemoryInfo(handle).total/1024/1024/1024 # 显卡总的显存大小
print("【显卡信息】Nvidia GPU数量：%d个 | 型号：%s | 显存：%dG" % (GPU_nums, GPU_name, GPU_memsize))
if ('GTX' in GPU_name or 'RTX' in GPU_name) and ('30' not in GPU_name): if_use_gpu = True; GPU_Ready = True
else:if_use_gpu = False; print("您的GPU不是本软件支持的Nvidia显卡，本软件将使用CPU进行后续处理！")
except: #没有Nvidia显卡
if_use_gpu = False; print("您的GPU不是Nvidia系列，本软件将使用CPU进行后续处理！"); GPU_memsize = 1.0
#if_use_gpu = True # 强制使用GPU
#if_use_gpu = False # 强制使用CPU
if if_use_gpu: os.environ['CUDA_VISIBLE_DEVICES'] = '0' # GPU环境变量
if GPU_memsize >= 3: if_good_model = True # 使用复杂的AI模型抠像
else: if_good_model = False
2、初始版本界面显示
由于调入AI模型的时间需要几秒，在等待时显示软件版本信息，启用一个新的线程专门处理图片显示，在主界面开始后，结束这个线程。代码如下：
pil_img = Image.open("start_img.jpg")
ImageDraw.Draw(pil_img).text((350,190), "图像抠图、背景替换工具v1.0", (255,255,255),font=ImageFont.truetype("msyh.ttc", 18))
ImageDraw.Draw(pil_img).text((410,330), "正在加载AI模型，请稍后 ......", (255,255,255),font=ImageFont.truetype("msyh.ttc", 16))
img_s = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
def showpic(): # 以下代码显示软件初始界面
while run_flag == 0:
cv2.imshow("iCANX VR Meeting System", img_s)
cv2.waitKey(100)
cv2.destroyAllWindows()
t = threading.Thread(target=showpic)
t.start()
3、格式转换
由于Python的图形库众多，它们的内部格式又都是各自专用的，所以需要一些函数进行格式转换。具体说明可以看代码注释，不懂的地方上网查询。
def CV2toPIL(self, img): # cv2转PIL
return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA))
def PILtoCV2(self, img): # PIL转cv2
return cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGRA)
def CvMatToQImage(self, ptr): # Converts an opencv MAT format into a QImage
ptr = cv2.cvtColor(ptr, cv2.COLOR_BGRA2RGBA) # 颜色格式转换
QtImg = QtGui.QImage(ptr.data, ptr.shape[1], ptr.shape[0], QtGui.QImage.Format_RGBA8888)
return QtGui.QPixmap.fromImage(QtImg)
def QImageToCvMat(self, my_image): # Converts a QImage into an opencv MAT format
my_image = my_image.convertToFormat(QImage.Format.Format_RGBA8888)
width = my_image.width()
height = my_image.height()
ptr = my_image.bits()
ptr.setsize(height * width * 4)
arr = np.frombuffer(ptr, np.uint8).reshape((height, width, 4))
return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
4、图像合并
此函数功能是将带alpha通道的图片叠加到背景图片，测试了几种方法，第3种效率最高。
def PngToGreen(back_img,fore_img):
scr_channels = cv2.split(fore_img)
dstt_channels = cv2.split(back_img)
b, g, r, a = cv2.split(fore_img)
for i in range(3):
dstt_channels[i][:,:] = dstt_channels[i][:,:] * (255.0 - a) / 255
dstt_channels[i][:,:] += np.array(scr_channels[i] * (a / 255), dtype=np.uint8)
return cv2.cvtColor(cv2.merge(dstt_channels), cv2.COLOR_BGRA2BGR)
def two_pic_combine_cv2(back_img,fore_img):
r, g, b, alpha = cv2.split(fore_img)
# Convert uint8 to float
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)
fore_img =cv2.cvtColor(fore_img, cv2.COLOR_BGRA2BGR)
back_img =cv2.cvtColor(back_img, cv2.COLOR_BGRA2BGR)
fore_img = fore_img.astype(float)
back_img = back_img.astype(float)
# Normalize the alpha mask to keep intensity between 0 and 1
alpha = alpha.astype(float) / 255
fore_img = cv2.multiply(alpha, fore_img)
back_img = cv2.multiply(1.0 - alpha, back_img)
return np.uint8(cv2.add(fore_img, back_img))
def two_pic_combine_PIL(self, back_img, fore_img):
back_img = self.CV2toPIL(back_img)
fore_img = self.CV2toPIL(fore_img)
r,g,b,alpha = fore_img.split()
return cv2.cvtColor(self.PILtoCV2(Image.composite(fore_img, back_img, alpha)), cv2.COLOR_BGRA2BGR)
以下是测试代码执行效率的代码：
import cv2,time
import numpy as np
from PIL import Image
f = cv2.imread("test.png",-1)
b = np.zeros((f.shape[0], f.shape[1], 3), np.uint8) # Creat a Image
b[:] = [88, 222, 118]
t1 = time.time(); res=PngToGreen(b,f)
print('消耗时间1：%.5f秒'%(time.time()-t1))
cv2.imshow('Img-1', res);cv2.waitKey(0)
t1 = time.time(); res=two_pic_combine_PIL(b,f)
print('消耗时间2：%.5f秒'%(time.time()-t1))
cv2.imshow('Img-2', res);cv2.waitKey(0)
t1 = time.time(); res=two_pic_combine_cv2(b,f)
print('消耗时间3：%.5f秒'%(time.time()-t1))
cv2.imshow('Img-3', res);cv2.waitKey(0)
cv2.destroyAllWindows()
5、AI人像分割
由于使用了飞桨，人像分割变得非常容易了，具体请参考飞桨网站。
6、音频处理
有时需要把原视频的音轨提取出来，与转换后的视频合并，用moviepy.editor库就非常容易了，代码如下：
audio = mpe.AudioFileClip(videofile) # 分离声轨
clip = mpe.VideoFileClip(tempfile)
videoclip = clip.set_audio(audio) # 写入声轨
videoclip.write_videofile(out_dir+'/'+os.path.splitext(os.path.split(videofile)[1])[0]+'_1.mp4')
7、窗口句柄获取
获取桌面上的窗口，如Zoom、腾讯会议等窗口，用于界面上的下拉框供用户选择。
def get_my_win(self): #获取桌面上的可见窗口，并存入列表
global win_list, win_list_t, hwnd, index_no, first_index
win_list =[]; hwnd_title = dict();
def get_all_hwnd(hwnd, mouse): # 获取全部窗口（句柄、窗口名称）
if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):
hwnd_title.update({
hwnd: win32gui.GetWindowText(hwnd)})
win32gui.EnumWindows(get_all_hwnd, 0)
for h, t in hwnd_title.items():
if t != '' and not ('PYTHON' in t) and not ('Microsoft' in t) and not ('Program' in t) and \
not ('iCANX' in t) and not ('python' in t) and not ('ZPTool' in t)and not ('设置' in t):
win_list.append([h, t[0:36]])
if t == 'Zoom 会议': first_index = index_no
elif t == '腾讯会议': first_index = index_no
index_no += 1
win_list.append([0, '请选择抓取的窗口...'])
if first_index==0: first_index = index_no
win_list_t = [i[1] for i in win_list]
8、屏幕抓取
用鼠标在需要抠像的程序窗口里抓取具体的人像矩形框，定义了类，方便使用。
class ScreenShot(QWidget): #定义抓取屏幕的类
def __init__(self, parent=None, butn_No=0):
super(ScreenShot, self).__init__(parent)
global get_win_run
get_win_run =1
self.butn_No1 = butn_No
self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
self.setStyleSheet('''background-color:black; ''')
self.setWindowOpacity(0.6)
desktop = QApplication.desktop()
winRect = win32gui.GetWindowRect(hwnd)
self.setGeometry(winRect[0],winRect[1],winRect[2]-winRect[0],winRect[3]-winRect[1])
self.setCursor(Qt.CrossCursor)
self.blackMask = QBitmap(winRect[2]-winRect[0],winRect[3]-winRect[1])
self.blackMask.fill(Qt.black)
self.mask = self.blackMask.copy()
self.isDrawing = False
self.startPoint = QPoint()
self.endPoint = QPoint()
def paintEvent(self, event):
if self.isDrawing:
self.mask = self.blackMask.copy()
pp = QPainter(self.mask)
pen = QPen()
pen.setStyle(Qt.NoPen)
pp.setPen(pen)
brush = QBrush(Qt.white)
pp.setBrush(brush)
pp.drawRect(QRect(self.startPoint, self.endPoint))
self.setMask(QBitmap(self.mask))
def mousePressEvent(self, event):
if event.button() == Qt.LeftButton:
self.startPoint = event.pos()
self.endPoint = self.startPoint
self.isDrawing = True
def mouseMoveEvent(self, event):
if self.isDrawing:
self.endPoint = event.pos()
self.update()
def mouseReleaseEvent(self, event):
global get_win_run
if event.button() == Qt.LeftButton:
self.endPoint = event.pos()
rect = QRect(self.startPoint, self.endPoint)
x = rect.left(); y = rect.top(); w=rect.width(); h=rect.height()
hang1=rect.top(); hang2 = rect.bottom(); lie1=rect.left();lie2=rect.right();
img_list[self.butn_No1][0] = hang1; img_list[self.butn_No1][1] = hang2;
img_list[self.butn_No1][2] = lie1; img_list[self.butn_No1][3] = lie2
get_win_run=0
self.close()
9、程序界面
用Pyqt5编写用户界面比较容易，网上有各种代码可以借鉴，Layout代码如下：
def createLayout(self):
mainLayout = QtWidgets.QVBoxLayout();topLayout1 = QtWidgets.QHBoxLayout();topLayout2 = QtWidgets.QHBoxLayout()
topLayout3 = QtWidgets.QHBoxLayout();topLayout4 = QtWidgets.QHBoxLayout()
self.my_label1 = QtWidgets.QLabel(); self.my_label2 = QtWidgets.QLabel()
topLayout1.addWidget(self.my_label1); topLayout1.addWidget(self.my_label2)
self.my_label1.setPixmap(QPixmap("ican_logo.png")); self.my_label2.setPixmap(QPixmap("ican_logo.png"))
self.my_label1.setFixedSize(427, 240); self.my_label2.setFixedSize(427, 240)
self.my_label1.setAlignment(Qt.AlignCenter); self.my_label2.setAlignment(Qt.AlignCenter)
self.my_label1.setToolTip("本区域，显示的是原始图像缩略图..."); self.my_label2.setToolTip("本区域，显示的是替换后的缩略图...")
self.GroupBox1 = QtWidgets.QGroupBox("软件设置")
self.GroupBox1.setFixedSize(280, 60)
self.lbl_1 = QtWidgets.QLabel("处理器：", self)
self.lbl_1.setFixedSize(45, 25)
self.comboBox1 = QtWidgets.QComboBox(self)
self.comboBox1.setFixedSize(50, 25)
self.comboBox1.addItem("CPU"); self.comboBox1.addItem("GPU")
if if_use_gpu: self.comboBox1.setCurrentIndex(1)
self.comboBox1.activated[str].connect(self.click_comboBox1)
self.lbl_2 = QtWidgets.QLabel("AI算法：", self)
self.lbl_2.setFixedSize(45, 25)
self.comboBox2 = QtWidgets.QComboBox(self)
self.comboBox2.setFixedSize(50, 25)
self.comboBox2.addItem("简易"); self.comboBox2.addItem("复杂")
if if_good_model:self.comboBox2.setCurrentIndex(1)
self.comboBox2.activated[str].connect(self.click_comboBox2)
GroupBox1Layout = QtWidgets.QHBoxLayout()
GroupBox1Layout.addWidget(self.lbl_2)
GroupBox1Layout.addWidget(self.comboBox2)
GroupBox1Layout.addWidget(self.lbl_1)
GroupBox1Layout.addWidget(self.comboBox1)
self.GroupBox1.setLayout(GroupBox1Layout)
if not if_use_gpu: self.comboBox1.setEnabled(False)
self.GroupBox2 = QtWidgets.QGroupBox("显示设置")
self.GroupBox2.setFixedSize(180, 60)
self.checkBox1 = QtWidgets.QCheckBox("原始图像")
self.checkBox2 = QtWidgets.QCheckBox("输出图像")
GroupBox2Layout = QtWidgets.QHBoxLayout()
GroupBox2Layout.addWidget(self.checkBox1)
GroupBox2Layout.addWidget(self.checkBox2)
self.GroupBox2.setLayout(GroupBox2Layout)
self.checkBox1.stateChanged.connect(self.box_choose)
self.checkBox2.stateChanged.connect(self.box_choose)
self.checkBox1.setChecked(True); self.checkBox2.setChecked(True)
self.GroupBox4 = QtWidgets.QGroupBox("文件设置")
self.GroupBox4.setFixedSize(850, 160)
self.filesButton = self.createButton("选择文件", self.filesButton_fuc)
self.outButton = self.createButton("输出目录", self.outButton_fuc)
self.bkfileButton = self.createButton("背景文件", self.bkfileButton_fuc)
self.filesButton.setToolTip("选择即将被替换背景的图片文件，可单选、多选...")
self.outButton.setToolTip("选择输出文件目录，替换后的文件将存在此目录...")
self.bkfileButton.setToolTip("选择可用作背景的图片文件，建议分辨率：1920x1080...")
self.filesButton.setFixedSize(80,23); self.outButton.setFixedSize(80,23)
self.bkfileButton.setFixedSize(80,23)
self.txt1 = QLabel('请选择图像文件[Ctrl+A全选、Ctrl/Shift+鼠标可多选]......', self); self.txt2 = QLabel('输出目录', self)
self.txt3 = QLabel('背景文件', self); self.txt4 = QLabel('纯色文件', self)
self.txt2.setText(out_dir); self.txt3.setText(bg_file);
self.txt4.setPixmap(self.CvMatToQImage(cv2.resize(bk_img, (50, 18))))
self.checkBox3_0 = QtWidgets.QCheckBox("只输出PNG文件(不换背景) ")
self.checkBox3_0.setToolTip("勾选：只输出透明的PNG文件，不更换背景；\n不勾选：同时输出PNG文件和更换背景的文件。");
self.checkBox3 = QtWidgets.QCheckBox("纯色背景")
self.checkBox3_0.stateChanged.connect(self.box_choose3_0)
self.checkBox3.stateChanged.connect(self.box_choose3)
self.checkBox3_0.setChecked(True)
self.red = QLabel(' 红', self); self.green= QLabel(' 绿', self); self.blue = QLabel(' 蓝', self)
self.red_e = QLineEdit(self); self.red_e.setText('8')
self.green_e = QLineEdit(self); self.green_e.setText('188')
self.blue_e = QLineEdit(self); self.blue_e.setText('8')
self.red_e.setValidator(QIntValidator(0, 254))
self.green_e.setValidator(QIntValidator(0, 254))
self.blue_e.setValidator(QIntValidator(0, 254))
self.red_e.setFixedSize(28, 20);self.green_e.setFixedSize(28, 20);self.blue_e.setFixedSize(28, 20)
self.red_e.textChanged[str].connect(self.red_e_fuc)
self.green_e.textChanged[str].connect(self.green_e_fuc)
self.blue_e.textChanged[str].connect(self.blue_e_fuc)
layout_box1 = QtWidgets.QHBoxLayout()
layout_box2 = QtWidgets.QHBoxLayout()
layout_box3 = QtWidgets.QHBoxLayout()
layout_box1.addWidget(self.filesButton, Qt.AlignLeft| Qt.AlignVCenter)
layout_box1.addWidget(self.txt1, Qt.AlignLeft| Qt.AlignVCenter)
layout_box2.addWidget(self.outButton)
layout_box2.addWidget(self.txt2)
layout_box2.addWidget(self.bkfileButton)
layout_box2.addWidget(self.txt3)
layout_box3.addWidget(self.checkBox3_0)
layout_box3.addWidget(self.checkBox3)
layout_box3.addWidget(self.txt4)
layout_box3.addWidget(self.red)
layout_box3.addWidget(self.red_e)
layout_box3.addWidget(self.green)
layout_box3.addWidget(self.green_e)
layout_box3.addWidget(self.blue)
layout_box3.addWidget(self.blue_e)
layout_box3.addStretch(1)
layout11 = QWidget(); layout21 = QWidget(); layout31 = QWidget()
layout11.setLayout(layout_box1); layout21.setLayout(layout_box2); layout31.setLayout(layout_box3)
GroupBoxmainLayout = QtWidgets.QVBoxLayout()
GroupBoxmainLayout.addWidget(layout11, Qt.AlignLeft | Qt.AlignVCenter)
GroupBoxmainLayout.addWidget(layout21, Qt.AlignLeft | Qt.AlignVCenter)
GroupBoxmainLayout.addWidget(layout31, Qt.AlignLeft | Qt.AlignVCenter)
self.GroupBox4.setLayout(GroupBoxmainLayout)
self.GroupBox5 = QtWidgets.QGroupBox("信息统计")
self.GroupBox5.setFixedSize(850, 90)
self.txt11 = QLabel('【图像信息】', self)
self.txt12 = QLabel('【运行信息】', self)
GroupBox5Layout = QtWidgets.QGridLayout()
GroupBox5Layout.addWidget(self.txt11, 0, 1)
GroupBox5Layout.addWidget(self.txt12, 1, 1)
self.GroupBox5.setLayout(GroupBox5Layout)
self.startButton = self.createButton("开始处理", self.startrun)
self.stopButton = self.createButton("停止", self.stoprun)
self.helpButton = self.createButton("帮助", self.helpWin)
self.quitButton = self.createButton("退出", self.quitWin)
self.startButton.setFixedSize(80,25)
self.stopButton.setFixedSize(55, 25)
self.helpButton.setFixedSize(55,25)
self.quitButton.setFixedSize(55,25)
topLayout2.addWidget(self.GroupBox4)
topLayout3.addWidget(self.GroupBox5)
topLayout4.addWidget(self.GroupBox1)
topLayout4.addWidget(self.GroupBox2)
topLayout4.addWidget(self.startButton)
topLayout4.addWidget(self.stopButton)
topLayout4.addWidget(self.helpButton)
topLayout4.addWidget(self.quitButton)
topLayout4.setSpacing(20)
layout1 = QWidget(); layout2 = QWidget()
layout3 = QWidget(); layout4 = QWidget()
layout1.setLayout(topLayout1); layout2.setLayout(topLayout2)
layout3.setLayout(topLayout3); layout4.setLayout(topLayout4)
mainLayout.addWidget(layout1, Qt.AlignLeft | Qt.AlignTop)
mainLayout.addWidget(layout2, Qt.AlignLeft | Qt.AlignBottom)
mainLayout.addWidget(layout3, Qt.AlignLeft | Qt.AlignBottom)
mainLayout.addWidget(layout4, Qt.AlignLeft | Qt.AlignBottom)
self.setLayout(mainLayout)
def createButton(self, text, member):
button = QtWidgets.QPushButton(text)
button.clicked.connect(member)
return button
10、主程序
由于软件的计算任务较重，使用多线程技术，避免用户界面卡顿，
class Winshot(QWidget):
def __init__(self):
super(Winshot, self).__init__()
global hwnd, run_flag
import paddlehub as hub # 导入模型很消耗时间，大概5秒
self.humanseg1 = hub.Module('deeplabv3p_xception65_humanseg') # 抠图复杂模型
self.humanseg2 = hub.Module('humanseg_mobile') # 抠图简易模型
palette1 = QtGui.QPalette()
palette1.setColor(palette1.Background, QtGui.QColor(200, 200, 200))
self.setPalette(palette1)
self.createLayout()
self.setWindowTitle("iCANX虚拟会议系统(图像抠图、背景替换工具)")
self.setFixedSize(my_width, my_height)
self.show(); run_flag = 1
class removebg(QThread):
def __init__(self):
super().__init__()
def run(self):
#main_run() #执行繁重图形计算任务的主函数
self.exec() # 进入事件循环
if __name__ == '__main__':
QApplication.setAttribute(Qt.AA_EnableHighDpiScaling) #适应高分辨率显示器、大字体适配
app = QtWidgets.QApplication(sys.argv)
Winshot = Winshot()
removebg_thread = removebg()
removebg_thread.start()
sys.exit(app.exec_())