Upload 2 files

app.py

@@ -2,24 +2,24 @@ from ultralytics import YOLO #行人识别，采用YoloV8模型
 import math #用于四舍五入取整
 import cv2 #opencv图像处理库
 import cvzone #在图像上绘画
-#import numpy as np #忘记取消注释了...
+import numpy as np
 import gradio as gr #GUI界面
 from sort import * #运动检测，采用sort算法
-#from deep_sort_realtime.deepsort_tracker import DeepSort #不会用QAQ
 import tempfile #创建输出临时文件夹
 import os
 from detectMotion import * #单独的运动检测
 import time
-#from matplotlib.pyplot import MultipleLocator
 
 
-#导入YoloV8模型,第一次使用会下载模型到当前文件夹当中
-model=YOLO("yolov8n.pt")
-global number_of_people_in_one_frame_list, flag
-number_of_people_in_one_frame_list = None
-flag = False
+#导入YoloV8模型,模型名称存储在model_list.txt中
+with open("model_list.txt") as file:
+    lines = file.readlines()
+    model_list = [line.rstrip() for line in lines] #逐行导入
 
-#YoloV8标签数据，本次项目只使用了'person'
+#用于人流量统计
+number_of_people_in_one_frame_list = []
+
+#YoloV8官方模型标签数据，本次项目只使用了'person'
 classNames=['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
     'fire hydrant',
     'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
@@ -36,16 +36,19 @@ classNames=['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train'
 #运动检测算法参数
 tracker=Sort(max_age=20,min_hits=3,iou_threshold=0.3)
 
+#实时显示人流量统计图表
 def plot_number():
-    global number_of_people_in_one_frame_list, flag
-    if flag == False:
-        time.sleep(15)
-        flag = True
+    global number_of_people_in_one_frame_list
     plt.close()
     fig = plt.figure()
     plt.xlabel('Second')
     plt.ylabel('Traffic')
     plt.grid(linestyle='--', alpha=0.3, linewidth=2)
+    #若函数在YOLO模型生成bounding box前运行，则返回空图表
+    if number_of_people_in_one_frame_list == []:
+        time.sleep(3) #延迟3秒
+        plt.plot([])
+        return fig
     plt.plot(number_of_people_in_one_frame_list)
     return fig
 
@@ -95,9 +98,11 @@ def processImg(img,sigma,median_filter,ISR):
     return res1, fig1, fig2
 
 #视频处理
-def processVideo(inputPath,codec):
-    global number_of_people_in_one_frame_list, flag
-    flag = False
+def processVideo(inputPath, codec, model):
+    global number_of_people_in_one_frame_list
+    if inputPath == None:
+        raise gr.Error("请先上传视频")
+    model=YOLO(model)
     number_of_people_in_one_frame = 0
     number_of_people_in_one_frame_list = []
     number_of_people = 0
@@ -112,10 +117,10 @@ def processVideo(inputPath,codec):
     #输出格式
     if codec == "mp4":
         #h264:avc1,此处为了兼容性换成mp4v
-        fourcc = cv2.VideoWriter_fourcc('m','p','4','v')#视频编码：h264,只有h264格式的mp4文件才能在浏览器直接播放
+        fourcc = cv2.VideoWriter_fourcc('m','p','4','v')#视频编码：mp4v,此编码不需要openh264-1.8.0-win64.dll
         video_save_path = os.path.join(outputPath,"output.mp4")#创建输出视频路径
     elif codec == "avi":
-        fourcc = cv2.VideoWriter_fourcc('h','2','6','4')#视频编码：h264,只能保存为avi格式且不能在浏览器直接播放
+        fourcc = cv2.VideoWriter_fourcc('h','2','6','4')#视频编码：h264,只能保存为avi格式,本地运行时不能在浏览器播放
         video_save_path = os.path.join(outputPath,"output.avi")#创建输出视频路径
     elif codec == "mkv":
         fourcc = cv2.VideoWriter_fourcc('X','V','I','D')#视频编码：XVID,此编码不需要openh264-1.8.0-win64.dll
@@ -125,14 +130,14 @@ def processVideo(inputPath,codec):
         video_save_path = os.path.join(outputPath,"output.wmv")#创建输出视频路径
     elif codec == "mp4(h264)":
         #h264:avc1, 在本地运行时可以在浏览器播放
-        fourcc = cv2.VideoWriter_fourcc('a','v','c','1')#视频编码：h264,只有h264格式的mp4文件才能在浏览器直接播放
-        video_save_path = os.path.join(outputPath,"output.mp4")#创建输出视频路径
+        fourcc = cv2.VideoWriter_fourcc('a','v','c','1')#视频编码：h264,本地运行时可直接在浏览器播放
+        video_save_path = os.path.join(outputPath,"output_h264.mp4")#创建输出视频路径
 
     output_viedo.open(video_save_path , fourcc, fps, size, True)
     #对每一帧图片进行读取和处理
     while True:
         ret, img = cap.read()#将每一帧图片读取到img当中
-        results=model(img,stream=True)#使用YoloV8模型进行推理，stream=True提高流式处理速度
+        results=model(img,stream=True)#使用YoloV8模型进行推理
         detections=np.empty((0, 5))#初始化运动检测
         if not(ret):#当视频全部读完，ret返回false，终止循环，视频帧读取和写入结束
             break
@@ -157,7 +162,6 @@ def processVideo(inputPath,codec):
         sum_of_frame += 1
         if sum_of_frame % fps == 0:
             number_of_people_in_one_frame_list.append(number_of_people_in_one_frame)
-        #print(type(number_of_people_in_one_frame_list))
         number_of_people_in_one_frame = 0
         #运动检测
         resultsTracker=tracker.update(detections)
@@ -165,13 +169,11 @@ def processVideo(inputPath,codec):
             x1,y1,x2,y2,Id=result
             number_of_people=max(str(int(Id)),str(number_of_people))
             x1,y1,x2,y2=int(x1),int(y1),int(x2),int(y2)#将浮点数转变为整型
-            #print(result)
             cvzone.putTextRect(img,f'{int(Id)}',(max(0,x1),max(30,y1)),scale=0.7,thickness=1)
         #image_np = np.squeeze(img.render())#用np.squeeze将输出结果降维
         output_viedo.write(img)#将处理后的图像写入视频
     output_viedo.release()#释放
     cap.release()#释放
-    flag = False
     return video_save_path,video_save_path
 
 
@@ -198,6 +200,7 @@ with gr.Blocks() as demo:
                 text_inputPath = gr.Video()
                 codec = gr.Radio(["mp4","avi","mkv","wmv",'mp4(h264)'], label="输出视频格式",
                                  value="mp4")
+                model = gr.Dropdown(model_list, value="yolov8n.pt", label="模型", info="模型列表存储在model_list.txt中")
                 videoProcess_button = gr.Button("处理")
             with gr.Column():
                 text_output = gr.Video()
@@ -205,7 +208,6 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 figure_number_output = gr.Plot(label="人流量")
-                #plotProcess_button = gr.Button("人流量显示")
     with gr.Tab("图像增强"):
         with gr.Row():
             with gr.Column():
@@ -226,14 +228,15 @@ with gr.Blocks() as demo:
                     motionProcess_button = gr.Button("处理")
                 motion_output_frame = gr.Video()
                 motion_output_fmask = gr.Video()
-                frame_output_path = gr.Text(label="frame输出路径")
-                fmask_output_path = gr.Text(label="mask输出路径")
+                with gr.Column():
+                    frame_output_path = gr.Text(label="frame输出路径")
+                    fmask_output_path = gr.Text(label="mask输出路径")
             with gr.Accordion("算法："):
                 gr.Markdown("高斯混合模型（GMM）")
 
     
 
-    videoProcess_button.click(processVideo, inputs=[text_inputPath,codec],
+    videoProcess_button.click(processVideo, inputs=[text_inputPath, codec, model],
                                outputs=[text_output,text_output_path])
     image_button.click(processImg, inputs=[image_input,image_sigma,median_filter,Add_ISR], 
                        outputs=[image_output,figure_pre_output,figure_post_output])

model_list.txt

@@ -0,0 +1,5 @@
+yolov8n.pt
+yolov8s.pt
+yolov8m.pt
+yolov8l.pt
+yolov8x.pt