|
@@ -3,12 +3,13 @@ import cv2
|
|
import threading
|
|
import threading
|
|
from datetime import datetime
|
|
from datetime import datetime
|
|
from ultralytics import YOLO
|
|
from ultralytics import YOLO
|
|
-from globals import stop_event,redis_client
|
|
|
|
-from config import SAVE_IMG_PATH,POST_IMG_PATH1,VIDEO_SOURCE,MODEL_PATH
|
|
|
|
|
|
+from globals import stop_event,redis_client,steps,hand_box,head_box
|
|
|
|
+from config import VIDEO_SOURCE,MODEL_PATH
|
|
|
|
|
|
|
|
|
|
def init_compressed_oxygen_detection():
|
|
def init_compressed_oxygen_detection():
|
|
- pass
|
|
|
|
|
|
+ for i, step in enumerate(steps):
|
|
|
|
+ redis_client.set(f"compressed_oxygen_step_{i+1}",'False')
|
|
|
|
|
|
def start_compressed_oxygen_detection(start_events):
|
|
def start_compressed_oxygen_detection(start_events):
|
|
|
|
|
|
@@ -19,6 +20,24 @@ def start_compressed_oxygen_detection(start_events):
|
|
thread.start()
|
|
thread.start()
|
|
thread.join()
|
|
thread.join()
|
|
|
|
|
|
|
|
+def IoU(box1, box2):
|
|
|
|
+ '''
|
|
|
|
+ 计算两个矩形框的交并比
|
|
|
|
+ :param box1: list,第一个矩形框的左上角和右下角坐标
|
|
|
|
+ :param box2: list,第二个矩形框的左上角和右下角坐标
|
|
|
|
+ :return: 两个矩形框的交并比iou
|
|
|
|
+ '''
|
|
|
|
+ x1 = max(box1[0], box2[0]) # 交集左上角x
|
|
|
|
+ x2 = min(box1[2], box2[2]) # 交集右下角x
|
|
|
|
+ y1 = max(box1[1], box2[1]) # 交集左上角y
|
|
|
|
+ y2 = min(box1[3], box2[3]) # 交集右下角y
|
|
|
|
+
|
|
|
|
+ overlap = max(0., x2-x1) * max(0., y2-y1)
|
|
|
|
+ union = (box1[2]-box1[0]) * (box1[3]-box1[1]) \
|
|
|
|
+ + (box2[2]-box2[0]) * (box2[3]-box2[1]) \
|
|
|
|
+ - overlap
|
|
|
|
+
|
|
|
|
+ return overlap/union
|
|
|
|
|
|
def process_video(model_path, video_source, start_event):
|
|
def process_video(model_path, video_source, start_event):
|
|
|
|
|
|
@@ -35,7 +54,7 @@ def process_video(model_path, video_source, start_event):
|
|
if cap.get(cv2.CAP_PROP_POS_FRAMES) % 10 != 0:#跳帧检测,
|
|
if cap.get(cv2.CAP_PROP_POS_FRAMES) % 10 != 0:#跳帧检测,
|
|
continue
|
|
continue
|
|
|
|
|
|
-
|
|
|
|
|
|
+ global step,hand_box,head_box
|
|
results = model.predict(frame,conf=0.6,verbose=False)
|
|
results = model.predict(frame,conf=0.6,verbose=False)
|
|
|
|
|
|
for r in results:
|
|
for r in results:
|
|
@@ -43,16 +62,42 @@ def process_video(model_path, video_source, start_event):
|
|
confidences = r.boxes.conf # 提取所有检测到的置信度
|
|
confidences = r.boxes.conf # 提取所有检测到的置信度
|
|
classes = r.boxes.cls # 提取所有检测到的类别索引
|
|
classes = r.boxes.cls # 提取所有检测到的类别索引
|
|
|
|
|
|
-
|
|
|
|
|
|
+ head_box=[]
|
|
|
|
+ hand_box=[]
|
|
for i in range(len(boxes)):
|
|
for i in range(len(boxes)):
|
|
x1, y1, x2, y2 = boxes[i].tolist()
|
|
x1, y1, x2, y2 = boxes[i].tolist()
|
|
confidence = confidences[i].item()
|
|
confidence = confidences[i].item()
|
|
cls = int(classes[i].item())
|
|
cls = int(classes[i].item())
|
|
label = model.names[cls]
|
|
label = model.names[cls]
|
|
print("label:",label)
|
|
print("label:",label)
|
|
- #pass
|
|
|
|
-
|
|
|
|
|
|
+
|
|
|
|
+ if(label=='head'):head_box=[x1,y1,x2,y2]
|
|
|
|
+ if(label=='hand'):hand_box=[x1,y1,x2,y2]
|
|
|
|
+
|
|
|
|
+ if(label=='aerostat_gasbag'):
|
|
|
|
+ steps[0]=True
|
|
|
|
+ if(IoU(head_box,[x1,y1,x2,y2])>0.1):steps[2]=True
|
|
|
|
+
|
|
|
|
+ if(label=='neckband'):
|
|
|
|
+ if(IoU(head_box,[x1,y1,x2,y2])>0.1):steps[1]=True
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+ if(label=='valve'):
|
|
|
|
+ if(IoU(hand_box,[x1,y1,x2,y2])>0.1):steps[2]=True
|
|
|
|
|
|
|
|
+ if(label=='air make-up'):
|
|
|
|
+ steps[0]=True
|
|
|
|
+ if(IoU(hand_box,[x1,y1,x2,y2])>0.1):steps[4]=True
|
|
|
|
+
|
|
|
|
+ if(label=='nose clip'):
|
|
|
|
+ if(IoU(head_box,[x1,y1,x2,y2])>0.1):steps[5]=True
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+ for i, step in enumerate(steps):
|
|
|
|
+ if step and redis_client.get(f"compressed_oxygen_step_{i+1}")=='False':
|
|
|
|
+ redis_client.rpush("compressed_oxygen_order", f"{i+1}")
|
|
|
|
+ redis_client.set(f"compressed_oxygen_step_{i+1}",'True')
|
|
|
|
+
|
|
start_event.set()
|
|
start_event.set()
|
|
|
|
|
|
|
|
|