import sys
import os
import numpy as np
os.environ['OPENCV_VIDEOIO_PRIORITY_LIST'] = 'FFMPEG,V4L2,MSMF,DSHOW'
os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = 'video_codec;h264_cuvid'
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QLabel, QComboBox
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QImage, QPixmap, QCursor
def import_cv2():
try:
import cv2
_ = cv2.__version__
print(f"OpenCV version: {cv2.__version__}")
return cv2
except Exception as e:
print(f"OpenCV import failed: {e}")
try:
import subprocess
subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python-headless==4.5.5.64", "-q"])
import cv2
print("Successfully installed opencv-python-headless")
return cv2
except:
print("Please install OpenCV manually: pip install opencv-python-headless")
return None
cv2 = None
try:
from ultralytics import YOLO
except ImportError:
print("Installing ultralytics...")
import subprocess
subprocess.check_call([sys.executable, "-m", "pip", "install", "ultralytics"])
from ultralytics import YOLO
try:
from deep_sort_realtime.deepsort_tracker import DeepSort
except ImportError:
print("Installing deep-sort-realtime...")
import subprocess
subprocess.check_call([sys.executable, "-m", "pip", "install", "deep-sort-realtime"])
from deep_sort_realtime.deepsort_tracker import DeepSort
try:
from PIL import ImageFont, ImageDraw, Image
except ImportError:
print("Installing Pillow...")
import subprocess
subprocess.check_call([sys.executable, "-m", "pip", "install", "Pillow"])
from PIL import ImageFont, ImageDraw, Image
class VideoPreviewThread(QThread):
frame_update_signal = pyqtSignal(np.ndarray)
preview_info_signal = pyqtSignal(str)
finished_signal = pyqtSignal()
def __init__(self, source_type, source_path=None):
super().__init__()
global cv2
if cv2 is None:
cv2 = import_cv2()
if cv2 is None:
raise ImportError("OpenCV is required but could not be imported")
self.source_type = source_type
self.source_path = source_path
self.is_running = False
self.cap = None
def run(self):
global cv2
self.is_running = True
try:
if self.source_type == 'camera':
self.cap = cv2.VideoCapture(0)
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
elif self.source_type == 'video' and self.source_path:
self.cap = cv2.VideoCapture(self.source_path)
if not self.cap or not self.cap.isOpened():
self.preview_info_signal.emit("无法打开视频源")
self.finished_signal.emit()
return
ret, test_frame = self.cap.read()
if not ret or test_frame is None:
self.preview_info_signal.emit("无法读取视频帧")
self.finished_signal.emit()
return
if self.source_type == 'video':
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
self.preview_info_signal.emit("视频源已打开，准备预览")
while self.is_running:
ret, frame = self.cap.read()
if not ret or frame is None or frame.size == 0:
if self.source_type == 'video':
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
continue
else:
break
display_frame = frame.copy()
h, w = frame.shape[:2]
cv2.putText(display_frame, "点击'选择目标'按钮，然后在画面中框选目标",
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
cv2.putText(display_frame, f"视频尺寸: {w}x{h}",
(10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
self.frame_update_signal.emit(display_frame)
self.preview_info_signal.emit(f"预览中... 尺寸: {w}x{h}")
except Exception as e:
self.preview_info_signal.emit(f"预览错误: {str(e)}")
finally:
if self.cap:
self.cap.release()
self.finished_signal.emit()
def stop(self):
self.is_running = False
class TargetTrackingThread(QThread):
frame_update_signal = pyqtSignal(np.ndarray)
tracking_info_signal = pyqtSignal(str)
finished_signal = pyqtSignal()
def __init__(self, source_type, source_path=None, selected_target=None):
super().__init__()
global cv2
if cv2 is None:
cv2 = import_cv2()
if cv2 is None:
raise ImportError("OpenCV is required but could not be imported")
self.source_type = source_type
self.source_path = source_path
self.selected_target = selected_target
self.is_running = False
self.is_paused = False
self.model = YOLO('yolov8n.pt')
self.deepsort = DeepSort(max_age=50, n_init=5, nn_budget=100) # 增加max_age和n_init以提高追踪稳定性
self.cap = None
self.current_frame_idx = 0
self.last_known_position = None
self.consecutive_missing_frames = 0
self.max_consecutive_missing = 10
self.tracking_id = None
self.target_features = None
def set_selected_target(self, bbox):
self.selected_target = bbox
def cv2_add_text(self, img, text, pos, text_color=(0, 255, 0), text_size=20):
global cv2
try:
font = ImageFont.truetype("simhei.ttf", text_size)
except IOError:
try:
font = ImageFont.truetype("arial.ttf", text_size)
except IOError:
return cv2.putText(img, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 0.7, text_color, 2)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
draw = ImageDraw.Draw(Image.fromarray(img_rgb))
draw.text(pos, text, fill=text_color, font=font)
return cv2.cvtColor(np.asarray(img_rgb), cv2.COLOR_RGB2BGR)
def extract_target_features(self, frame, bbox):
x1, y1, x2, y2 = bbox
x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(frame.shape[1], x2)
y2 = min(frame.shape[0], y2)
if x2 <= x1 or y2 <= y1:
return None
target_region = frame[y1:y2, x1:x2]
hist = cv2.calcHist([target_region], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
hist = cv2.normalize(hist, hist).flatten()
return hist
def compare_features(self, features1, features2):
if features1 is None or features2 is None:
return 0
return cv2.compareHist(features1, features2, cv2.HISTCMP_CORREL)
def run(self):
global cv2
self.is_running = True
try:
if self.source_type == 'camera':
self.cap = cv2.VideoCapture(0)
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
elif self.source_type == 'video' and self.source_path:
self.cap = cv2.VideoCapture(self.source_path)
if not self.cap or not self.cap.isOpened():
self.tracking_info_signal.emit("无法打开视频源")
self.finished_signal.emit()
return
if self.source_type == 'video':
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
self.current_frame_idx = 0
while self.is_running:
if not self.is_paused:
ret, frame = self.cap.read()
if not ret or frame is None or frame.size == 0:
break
self.current_frame_idx += 1
if self.selected_target is not None:
results = self.model(frame, conf=0.4, verbose=False)[0]
detections = []
for box in results.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = float(box.conf[0])
cls = int(box.cls[0])
detections.append(([x1, y1, x2 - x1, y2 - y1], conf, cls))
tracks = self.deepsort.update_tracks(detections, frame=frame)
display_frame = frame.copy()
target_tracked = False
if self.target_features is None and self.selected_target is not None:
self.target_features = self.extract_target_features(frame, self.selected_target)
for track in tracks:
if not track.is_confirmed():
continue
track_id = track.track_id
ltrb = track.to_ltrb()
x1, y1, x2, y2 = map(int, ltrb)
st_x1, st_y1, st_x2, st_y2 = self.selected_target
iou = self.calculate_iou((st_x1, st_y1, st_x2, st_y2), (x1, y1, x2, y2))
if iou > 0.3 or track_id == self.tracking_id:
if self.tracking_id is None:
self.tracking_id = track_id
if track_id == self.tracking_id or iou > 0.3:
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
display_frame = self.cv2_add_text(display_frame, f"追踪目标 ID: {track_id}", (x1, y1 - 10), (0, 0, 255), 16)
target_tracked = True
self.last_known_position = (x1, y1, x2, y2)
self.consecutive_missing_frames = 0 # 重置丢失计数
current_features = self.extract_target_features(frame, (x1, y1, x2, y2))
if current_features is not None:
self.target_features = current_features
if not target_tracked and self.target_features is not None:
best_match = None
best_similarity = 0.7 # 相似度阈值
for track in tracks:
if not track.is_confirmed():
continue
track_id = track.track_id
ltrb = track.to_ltrb()
x1, y1, x2, y2 = map(int, ltrb)
current_features = self.extract_target_features(frame, (x1, y1, x2, y2))
if current_features is not None:
similarity = self.compare_features(self.target_features, current_features)
if similarity > best_similarity:
best_similarity = similarity
best_match = track
if best_match is not None:
track_id = best_match.track_id
ltrb = best_match.to_ltrb()
x1, y1, x2, y2 = map(int, ltrb)
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 255), 3) # 使用黄色表示重新找到的目标
display_frame = self.cv2_add_text(display_frame, f"重新找到 ID: {track_id}", (x1, y1 - 10), (0, 255, 255), 16)
target_tracked = True
self.tracking_id = track_id
self.last_known_position = (x1, y1, x2, y2)
self.consecutive_missing_frames = 0
if not target_tracked:
self.consecutive_missing_frames += 1
if self.last_known_position is not None and self.consecutive_missing_frames <= self.max_consecutive_missing:
x1, y1, x2, y2 = self.last_known_position
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (255, 0, 0), 2) # 使用蓝色表示预测位置
display_frame = self.cv2_add_text(display_frame, "预测位置", (x1, y1 - 10), (255, 0, 0), 16)
display_frame = self.cv2_add_text(display_frame, f"目标遮挡/丢失 {self.consecutive_missing_frames}/{self.max_consecutive_missing}",
(10, 30), (255, 0, 0), 20)
else:
display_frame = self.cv2_add_text(display_frame, "目标丢失", (10, 30), (0, 0, 255), 20)
if self.selected_target is not None:
x1, y1, x2, y2 = self.selected_target
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
display_frame = self.cv2_add_text(display_frame, "初始目标", (x1, y1 - 30), (0, 255, 0), 16)
else:
display_frame = frame.copy()
display_frame = self.cv2_add_text(display_frame, "请先选择目标", (10, 30), (0, 255, 0), 20)
display_frame = self.cv2_add_text(display_frame, f"帧: {self.current_frame_idx}", (10, 60), (255, 255, 255), 16)
self.frame_update_signal.emit(display_frame)
status = "追踪中" if target_tracked else "目标丢失" if self.consecutive_missing_frames > 0 else "未追踪"
self.tracking_info_signal.emit(f"尺寸: {frame.shape[1]}x{frame.shape[0]} | 状态: {status}")
except Exception as e:
self.tracking_info_signal.emit(f"错误: {str(e)}")
finally:
if self.cap:
self.cap.release()
self.finished_signal.emit()
def calculate_iou(self, box1, box2):
x1, y1, x2, y2 = box1
x1_, y1_, x2_, y2_ = box2
xi1, yi1 = max(x1, x1_), max(y1, y1_)
xi2, yi2 = min(x2, x2_), min(y2, y2_)
inter_area = max(0, xi2 - xi1) * max(0, yi2 - yi1)
box1_area = (x2 - x1) * (y2 - y1)
box2_area = (x2_ - x1_) * (y2_ - y1_)
return inter_area / (box1_area + box2_area - inter_area + 1e-6)
def pause(self):
self.is_paused = True
def resume(self):
self.is_paused = False
def stop(self):
self.is_running = False
def restart(self):
if self.cap and self.source_type == 'video':
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
self.current_frame_idx = 0
self.last_known_position = None
self.consecutive_missing_frames = 0
self.tracking_id = None
self.target_features = None
class VideoLabel(QLabel):
mouse_press_signal = pyqtSignal(int, int)
mouse_move_signal = pyqtSignal(int, int)
mouse_release_signal = pyqtSignal(int, int, int, int)
def __init__(self, parent=None):
super().__init__(parent)
self.setMouseTracking(True)
self.select_start = (0, 0)
self.setCursor(QCursor(Qt.CrossCursor))
self.setMinimumSize(640, 480)
self.setText("请打开视频源")
self.setStyleSheet("border: 2px solid #000; background: #333; color: white; font-size: 16px;")
def mousePressEvent(self, event):
if event.button() == Qt.LeftButton:
self.select_start = (event.x(), event.y())
self.mouse_press_signal.emit(event.x(), event.y())
def mouseMoveEvent(self, event):
self.mouse_move_signal.emit(event.x(), event.y())
def mouseReleaseEvent(self, event):
if event.button() == Qt.LeftButton:
x1, y1 = self.select_start
x2, y2 = event.x(), event.y()
self.mouse_release_signal.emit(x1, y1, x2, y2)
class TargetTrackingGUI(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("目标追踪系统 - 改进版（处理遮挡）")
self.setGeometry(100, 100, 1200, 800)
self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)
self.preview_thread = None
self.tracking_thread = None
self.is_selecting_target = False
self.original_frame_size = (0, 0)
self.scale = 1.0
self.select_start = (0, 0)
self.temp_frame = None
self.selected_target = None
self.source_type = None
self.source_path = None
self.init_ui()
def init_ui(self):
central_widget = QWidget()
self.setCentralWidget(central_widget)
layout = QVBoxLayout(central_widget)
control_layout = QHBoxLayout()
self.source_combo = QComboBox()
self.source_combo.addItems(["摄像头", "本地视频"])
control_layout.addWidget(QLabel("视频源:"))
control_layout.addWidget(self.source_combo)
self.open_btn = QPushButton("打开视频源")
self.open_btn.clicked.connect(self.open_source)
control_layout.addWidget(self.open_btn)
self.select_target_btn = QPushButton("选择目标")
self.select_target_btn.clicked.connect(self.start_select_target)
self.select_target_btn.setEnabled(False)
control_layout.addWidget(self.select_target_btn)
self.start_tracking_btn = QPushButton("开始追踪")
self.start_tracking_btn.clicked.connect(self.start_tracking)
self.start_tracking_btn.setEnabled(False)
control_layout.addWidget(self.start_tracking_btn)
self.pause_btn = QPushButton("暂停")
self.pause_btn.clicked.connect(self.pause_tracking)
self.pause_btn.setEnabled(False)
control_layout.addWidget(self.pause_btn)
self.stop_btn = QPushButton("停止")
self.stop_btn.clicked.connect(self.stop_all)
self.stop_btn.setEnabled(False)
control_layout.addWidget(self.stop_btn)
layout.addLayout(control_layout)
self.info_label = QLabel("状态: 未开始 | 请选择视频源并打开")
layout.addWidget(self.info_label)
self.video_label = VideoLabel()
self.video_label.setAlignment(Qt.AlignCenter)
self.video_label.mouse_press_signal.connect(self.on_mouse_press)
self.video_label.mouse_move_signal.connect(self.on_mouse_move)
self.video_label.mouse_release_signal.connect(self.on_mouse_release)
layout.addWidget(self.video_label)
def open_source(self):
self.stop_all()
self.source_type = self.source_combo.currentText()
self.source_path = None
if self.source_type == "本地视频":
self.source_path, _ = QFileDialog.getOpenFileName(self, "选择视频", "", "Video (*.mp4 *.avi *.mov *.mkv)")
if not self.source_path:
return
global cv2
if cv2 is None:
cv2 = import_cv2()
if cv2 is None:
self.info_label.setText("错误: 无法导入OpenCV")
return
self.preview_thread = VideoPreviewThread("camera" if self.source_type == "摄像头" else "video", self.source_path)
self.preview_thread.frame_update_signal.connect(self.update_frame)
self.preview_thread.preview_info_signal.connect(self.update_info)
self.preview_thread.finished_signal.connect(self.preview_finished)
self.preview_thread.start()
self.open_btn.setEnabled(False)
self.select_target_btn.setEnabled(True)
self.stop_btn.setEnabled(True)
self.info_label.setText("状态: 预览中 | 请选择目标")
def start_select_target(self):
if self.preview_thread is None:
self.info_label.setText("错误: 请先打开视频源")
return
self.is_selecting_target = True
self.info_label.setText("状态: 框选目标中 | 在视频画面上拖动鼠标框选目标")
self.video_label.setCursor(QCursor(Qt.CrossCursor))
def start_tracking(self):
if self.selected_target is None:
self.info_label.setText("状态: 错误 | 请先选择目标")
return
if self.preview_thread:
self.preview_thread.stop()
self.preview_thread.wait()
self.preview_thread = None
self.tracking_thread = TargetTrackingThread(
"camera" if self.source_type == "摄像头" else "video",
self.source_path,
self.selected_target
)
self.tracking_thread.frame_update_signal.connect(self.update_frame)
self.tracking_thread.tracking_info_signal.connect(self.update_info)
self.tracking_thread.finished_signal.connect(self.tracking_finished)
self.tracking_thread.start()
self.pause_btn.setText("暂停")
self.start_tracking_btn.setEnabled(False)
self.select_target_btn.setEnabled(False)
self.pause_btn.setEnabled(True)
self.info_label.setText("状态: 追踪中 | 正在追踪选定目标")
def update_frame(self, frame):
self.temp_frame = frame.copy()
h, w, ch = frame.shape
self.original_frame_size = (w, h)
label_width = self.video_label.width()
label_height = self.video_label.height()
if label_width > 0 and label_height > 0:
self.scale = min(label_width / w, label_height / h)
new_w, new_h = int(w * self.scale), int(h * self.scale)
else:
new_w, new_h = w, h
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
rgb_frame = cv2.resize(rgb_frame, (new_w, new_h))
bytes_per_line = ch * new_w
qt_image = QImage(rgb_frame.data, new_w, new_h, bytes_per_line, QImage.Format_RGB888)
self.video_label.setPixmap(QPixmap.fromImage(qt_image))
def on_mouse_press(self, x, y):
if self.is_selecting_target and self.temp_frame is not None:
self.select_start = (x, y)
def on_mouse_move(self, x, y):
if self.is_selecting_target and self.temp_frame is not None:
global cv2
frame = self.temp_frame.copy()
x1, y1 = self.select_start
x2, y2 = x, y
scaled_x1 = int(x1 / self.scale)
scaled_y1 = int(y1 / self.scale)
scaled_x2 = int(x2 / self.scale)
scaled_y2 = int(y2 / self.scale)
cv2.rectangle(frame, (scaled_x1, scaled_y1), (scaled_x2, scaled_y2), (255, 0, 0), 2)
width = abs(scaled_x2 - scaled_x1)
height = abs(scaled_y2 - scaled_y1)
cv2.putText(frame, f"{width}x{height}", (scaled_x1, scaled_y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
new_w, new_h = int(self.original_frame_size[0] * self.scale), int(self.original_frame_size[1] * self.scale)
rgb_frame = cv2.resize(rgb_frame, (new_w, new_h))
bytes_per_line = 3 * new_w
qt_image = QImage(rgb_frame.data, new_w, new_h, bytes_per_line, QImage.Format_RGB888)
self.video_label.setPixmap(QPixmap.fromImage(qt_image))
def on_mouse_release(self, x1, y1, x2, y2):
if self.is_selecting_target and self.temp_frame is not None:
scaled_x1 = int(min(x1, x2) / self.scale)
scaled_y1 = int(min(y1, y2) / self.scale)
scaled_x2 = int(max(x1, x2) / self.scale)
scaled_y2 = int(max(y1, y2) / self.scale)
width = abs(scaled_x2 - scaled_x1)
height = abs(scaled_y2 - scaled_y1)
if width > 20 and height > 20:
self.selected_target = (scaled_x1, scaled_y1, scaled_x2, scaled_y2)
self.start_tracking_btn.setEnabled(True)
self.info_label.setText(f"状态: 目标已选择 | 坐标: {self.selected_target} | 尺寸: {width}x{height} | 点击'开始追踪'从头追踪")
else:
self.info_label.setText("状态: 框选过小 | 请选择更大的区域")
self.is_selecting_target = False
self.video_label.setCursor(QCursor(Qt.ArrowCursor))
def update_info(self, info):
current_text = self.info_label.text().split(" | ")[0]
self.info_label.setText(f"{current_text} | {info}")
def pause_tracking(self):
if self.tracking_thread:
if self.tracking_thread.is_paused:
self.tracking_thread.resume()
self.pause_btn.setText("暂停")
self.info_label.setText("状态: 追踪中")
else:
self.tracking_thread.pause()
self.pause_btn.setText("继续")
self.info_label.setText("状态: 已暂停")
def stop_all(self):
if self.preview_thread:
self.preview_thread.stop()
self.preview_thread.wait()
self.preview_thread = None
if self.tracking_thread:
self.tracking_thread.stop()
self.tracking_thread.wait()
self.tracking_thread = None
self.video_label.clear()
self.video_label.setText("请打开视频源")
self.info_label.setText("状态: 已停止 | 选择视频源并打开")
self.open_btn.setEnabled(True)
self.select_target_btn.setEnabled(False)
self.start_tracking_btn.setEnabled(False)
self.pause_btn.setEnabled(False)
self.stop_btn.setEnabled(False)
self.is_selecting_target = False
self.temp_frame = None
self.selected_target = None
def preview_finished(self):
self.info_label.setText("状态: 预览结束 | 重新选择视频源")
self.open_btn.setEnabled(True)
self.select_target_btn.setEnabled(False)
self.start_tracking_btn.setEnabled(False)
self.pause_btn.setEnabled(False)
self.stop_btn.setEnabled(False)
def tracking_finished(self):
self.info_label.setText("状态: 追踪结束 | 重新选择视频源")
self.open_btn.setEnabled(True)
self.select_target_btn.setEnabled(False)
self.start_tracking_btn.setEnabled(False)
self.pause_btn.setEnabled(False)
self.stop_btn.setEnabled(False)
def closeEvent(self, event):
self.stop_all()
event.accept()
if __name__ == "__main__":
app = QApplication(sys.argv)
window = TargetTrackingGUI()
window.show()
sys.exit(app.exec_())有什么问题吗
最新发布
本文介绍了一个简单的Windows消息处理示例，包括定义消息常量IDM_EXIT和IDM_ABOUT，以及实现窗口过程函数WndProc和关于框处理函数About。
扫一扫
373
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
