| 285 | |
| 286 | |
| 287 | class Detector(object): |
| 288 | def __init__(self, args, model=None, seq_num=2): |
| 289 | |
| 290 | self.args = args |
| 291 | self.detr = model |
| 292 | |
| 293 | self.seq_num = seq_num |
| 294 | img_list = os.listdir(os.path.join(self.args.mot_path, self.seq_num, 'img1')) |
| 295 | img_list = [os.path.join(self.args.mot_path, self.seq_num, 'img1', _) for _ in img_list if |
| 296 | ('jpg' in _) or ('png' in _)] |
| 297 | |
| 298 | self.img_list = sorted(img_list) |
| 299 | self.img_len = len(self.img_list) |
| 300 | self.tr_tracker = MOTR() |
| 301 | |
| 302 | ''' |
| 303 | common settings |
| 304 | ''' |
| 305 | self.img_height = 800 |
| 306 | self.img_width = 1536 |
| 307 | self.mean = [0.485, 0.456, 0.406] |
| 308 | self.std = [0.229, 0.224, 0.225] |
| 309 | |
| 310 | self.save_path = os.path.join(self.args.output_dir, 'results/{}'.format(seq_num)) |
| 311 | os.makedirs(self.save_path, exist_ok=True) |
| 312 | |
| 313 | self.predict_path = os.path.join(self.args.output_dir, 'preds', self.seq_num) |
| 314 | os.makedirs(self.predict_path, exist_ok=True) |
| 315 | if os.path.exists(os.path.join(self.predict_path, 'gt.txt')): |
| 316 | os.remove(os.path.join(self.predict_path, 'gt.txt')) |
| 317 | |
| 318 | def load_img_from_file(self,f_path): |
| 319 | label_path = f_path.replace('images', 'labels_with_ids').replace('.png', '.txt').replace('.jpg', '.txt') |
| 320 | cur_img = cv2.imread(f_path) |
| 321 | cur_img = cv2.cvtColor(cur_img, cv2.COLOR_BGR2RGB) |
| 322 | targets = load_label(label_path, cur_img.shape[:2]) if os.path.exists(label_path) else None |
| 323 | return cur_img, targets |
| 324 | |
| 325 | def init_img(self, img): |
| 326 | ori_img = img.copy() |
| 327 | self.seq_h, self.seq_w = img.shape[:2] |
| 328 | scale = self.img_height / min(self.seq_h, self.seq_w) |
| 329 | if max(self.seq_h, self.seq_w) * scale > self.img_width: |
| 330 | scale = self.img_width / max(self.seq_h, self.seq_w) |
| 331 | target_h = int(self.seq_h * scale) |
| 332 | target_w = int(self.seq_w * scale) |
| 333 | img = cv2.resize(img, (target_w, target_h)) |
| 334 | img = F.normalize(F.to_tensor(img), self.mean, self.std) |
| 335 | img = img.unsqueeze(0) |
| 336 | return img, ori_img |
| 337 | |
| 338 | @staticmethod |
| 339 | def filter_dt_by_score(dt_instances: Instances, prob_threshold: float) -> Instances: |
| 340 | keep = dt_instances.scores > prob_threshold |
| 341 | return dt_instances[keep] |
| 342 | |
| 343 | @staticmethod |
| 344 | def filter_dt_by_area(dt_instances: Instances, area_threshold: float) -> Instances: |