[detect] 完成连通域的跟踪和检测，离开区域的时候会输出连通域信息

Author: eternal-echo

detect/parts_segment.py

@@ -33,10 +33,21 @@ def process(self, img):
         num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
         # 获取连通域面积
         areas = stats[:, 4]
-        # 对连通域按面积从小到大排序
-        sorted_areas_indices = np.argsort(areas)
+        # 对连通域按面积从大到小排序
+        sorted_areas_indices = np.argsort(areas)[::-1]
         # 选择面积在 self.min_area 到 self.max_area 像素之间的连通域
-        selected_areas_indices = sorted_areas_indices[(areas[sorted_areas_indices] > self.min_area) & (areas[sorted_areas_indices] < self.max_area)]
+        # selected_areas_indices = sorted_areas_indices[(areas[sorted_areas_indices] > self.min_area) & (areas[sorted_areas_indices] < self.max_area)]
+        selected_areas_indices = []
+        for i in range(len(sorted_areas_indices)):
+            if (areas[sorted_areas_indices[i]] < self.max_area
+                # 区域的边缘处不再选择
+                and stats[sorted_areas_indices[i], 1] + stats[sorted_areas_indices[i], 3] < img.shape[0] - 10
+                and stats[sorted_areas_indices[i], 1] > 10
+                and stats[sorted_areas_indices[i], 0] > 5
+                and stats[sorted_areas_indices[i], 0] + stats[sorted_areas_indices[i], 2] < img.shape[1] - 5):
+                if areas[sorted_areas_indices[i]] < self.min_area:
+                    break
+                selected_areas_indices.append(sorted_areas_indices[i])
 
         num_labels = len(selected_areas_indices)
         stats = stats[selected_areas_indices]

detect/parts_tracker.py

@@ -32,6 +32,12 @@ def update(self, target: TargetInfo):
         self.last = target
         self.tracker.correct(target.center)
 
+    def get_position(self):
+        return self.tracker.get_position()
+
+    def get_velocity(self):
+        return self.tracker.get_velocity()
+
 class Tracker:
     def __init__(self, center: Point):
         # 初始化 Kalman 滤波器
@@ -73,3 +79,9 @@ def get_position(self):
         y = int(self.state[1][0])
         return (x, y)
 
+    def get_velocity(self):
+        # 获取速度
+        vx = int(self.state[2][0])
+        vy = int(self.state[3][0])
+        return (vx, vy)
+

parts_sorting.py

@@ -129,6 +129,7 @@ def run(self):
         if self.is_save is not None:
             self.belt_video = cv2.VideoWriter(os.path.join(self.results_dir, 'belt.avi'), cv2.VideoWriter_fourcc(*'MJPG'), 30, (int(self.bbox_belt[2]), int(self.bbox_belt[3])))
             self.cc_video = cv2.VideoWriter(os.path.join(self.results_dir, 'cc.avi'), cv2.VideoWriter_fourcc(*'MJPG'), 30, (int(self.bbox_belt[2]), int(self.bbox_belt[3])))
+            self.track_video = cv2.VideoWriter(os.path.join(self.results_dir, 'track.avi'), cv2.VideoWriter_fourcc(*'MJPG'), 30, (int(self.bbox_belt[2]), int(self.bbox_belt[3])))
 
         while True:
             ret, frame = self.camera.read()
@@ -166,10 +167,6 @@ def run(self):
             #     self.components_centroids = components_centroids
 
             # TODO: 对零件特征点使用光流法来跟踪，目前直接对连通域的坐标进行卡尔曼滤波
-            for candidate in self.candidates:
-                # 处理物体消失
-                if frame_id - candidate.last.timestamp > 10 or candidate.last.center.y < 30:
-                    self.candidates.remove(candidate)
             # 遍历当前连通域，与候选目标队列的连通域进行匹配
             for i in range(components_num_labels):
                 x, y, w, h, area = components_stats[i]
@@ -192,14 +189,24 @@ def run(self):
                     if candidate.last.timestamp < frame_id and overlap > 0.8:
                         matched = True
                         # 当x坐标基本不变，y坐标在减小，并且面积基本不变时，认为跟踪成功
-                        if abs(dx) < 3 and dy < 0 and abs(candidate.last.area - area) < 0.1 * area:
+                        if abs(dx) < 3 and dy < 0:
                             candidate.update(cur_info)
                             break
 
-                # 如果不属于已有物体，则加入候选目标队列
-                if not matched:
+                # 如果不属于已有物体，并且在检测区域的下半部分，才加入候选目标队列
+                if (not matched) and y > belt.shape[0] / 3:
                     self.candidates.append(TargetTrack(cur_info))
 
+            # 过期检查
+            for candidate in self.candidates:
+                if candidate.last.rect.y < 30:
+                    self.__logger.info("目标已经离开传送带，目标信息：{}".format(candidate.last))
+                    self.candidates.remove(candidate)
+                    continue
+                # 处理过期物体
+                if frame_id - candidate.last.timestamp > 3:
+                    self.candidates.remove(candidate)
+
             # TODO: 卡尔曼滤波
             if self.is_show:
                 # 形态学处理的结果
@@ -227,14 +234,19 @@ def run(self):
                     # 绘制矩形
                     cv2.rectangle(candidates_belt, (candidate.last.rect.x, candidate.last.rect.y), (candidate.last.rect.x + candidate.last.rect.w, candidate.last.rect.y + candidate.last.rect.h), (0, 255, 0), 2)
                     # 显示连通域编号，中心点坐标，面积，时间戳
-                    # cv2.putText(candidates_belt, str(candidate.last.timestamp) + ": (" + str(candidate.last.center.x) + ", " + str(candidate.last.center.y) + ")", (candidate.last.rect.x, candidate.last.rect.y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
-                    cv2.putText(candidates_belt, "{0}: ({1:.2f}, {2:.2f})".format(candidate.last.timestamp, candidate.last.center.x, candidate.last.center.y), (candidate.last.rect.x, candidate.last.rect.y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+                    cv2.putText(candidates_belt, "{0}: ({1:.2f}, {2:.2f})".format(candidate.last.timestamp, candidate.last.center.x, candidate.last.center.y), (candidate.last.rect.x, candidate.last.rect.y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
+                    cv2.putText(candidates_belt, "{0}x{1}, {2}".format(candidate.last.rect.w, candidate.last.rect.h, candidate.last.area), (int(candidate.last.center.x), int(candidate.last.center.y)), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
+                    # 显示坐标和速度
+                    c_x, c_y = candidate.get_position()
+                    v_x, v_y = candidate.get_velocity()
+                    cv2.putText(candidates_belt, "({0:.2f}, {1:.2f}), ({2:.2f}, {3:.2f})".format(c_x, c_y, v_x, v_y), (candidate.last.rect.x, candidate.last.rect.y + candidate.last.rect.h), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
                 cv2.namedWindow("Candidates Belt", cv2.WINDOW_NORMAL)
                 cv2.imshow("Candidates Belt", candidates_belt)
 
             if self.is_save:
                 self.belt_video.write(belt)
                 self.cc_video.write(selected_belt)
+                self.track_video.write(candidates_belt)
 
 
             keyboard = cv2.waitKey(30)