| | from argparse import ArgumentParser |
| |
|
| | from mmdet.apis import inference_detector, init_detector, show_result_pyplot |
| | from mmdet.core.mask.utils import encode_mask_results |
| | import numpy as np |
| | import mmcv |
| | import torch |
| | from imantics import Polygons, Mask |
| | import json |
| | import os |
| | import cv2, glob |
| |
|
| | class detections(): |
| | def __init__(self, cfg_path, device, model_path = 'data/models/walt_vehicle.pth', threshold=0.85): |
| | self.model = init_detector(cfg_path, model_path, device=device) |
| | self.all_preds = [] |
| | self.all_scores = [] |
| | self.index = [] |
| | self.score_thr = threshold |
| | self.result = [] |
| | self.record_dict = {'model': cfg_path,'results': []} |
| | self.detect_count = [] |
| |
|
| |
|
| | def run_on_image(self, image): |
| | self.result = inference_detector(self.model, image) |
| | image_labelled = self.model.show_result(image, self.result, score_thr=self.score_thr) |
| | return image_labelled |
| |
|
| | def process_output(self, count): |
| | result = self.result |
| | infer_result = {'url': count, |
| | 'boxes': [], |
| | 'scores': [], |
| | 'keypoints': [], |
| | 'segmentation': [], |
| | 'label_ids': [], |
| | 'track': [], |
| | 'labels': []} |
| |
|
| | if isinstance(result, tuple): |
| | bbox_result, segm_result = result |
| | |
| | if isinstance(segm_result, tuple): |
| | segm_result = segm_result[0] |
| | bboxes = np.vstack(bbox_result) |
| | labels = [np.full(bbox.shape[0], i, dtype=np.int32) for i, bbox in enumerate(bbox_result)] |
| |
|
| | labels = np.concatenate(labels) |
| | segms = None |
| | if segm_result is not None and len(labels) > 0: |
| | segms = mmcv.concat_list(segm_result) |
| | if isinstance(segms[0], torch.Tensor): |
| | segms = torch.stack(segms, dim=0).detach().cpu().numpy() |
| | else: |
| | segms = np.stack(segms, axis=0) |
| |
|
| | for i, (bbox, label, segm) in enumerate(zip(bboxes, labels, segms)): |
| | if bbox[-1].item() <0.3: |
| | continue |
| | box = [bbox[0].item(), bbox[1].item(), bbox[2].item(), bbox[3].item()] |
| | polygons = Mask(segm).polygons() |
| | |
| | infer_result['boxes'].append(box) |
| | infer_result['segmentation'].append(polygons.segmentation) |
| | infer_result['scores'].append(bbox[-1].item()) |
| | infer_result['labels'].append(self.model.CLASSES[label]) |
| | infer_result['label_ids'].append(label) |
| | self.record_dict['results'].append(infer_result) |
| | self.detect_count = labels |
| | |
| | def write_json(self, filename): |
| | with open(filename + '.json', 'w') as f: |
| | json.dump(self.record_dict, f) |
| |
|
| | |
| | def main(): |
| | if torch.cuda.is_available() == False: |
| | device='cpu' |
| | else: |
| | device='cuda:0' |
| | detect_people = detections('configs/walt/walt_people.py', device, model_path='data/models/walt_people.pth') |
| | detect = detections('configs/walt/walt_vehicle.py', device, model_path='data/models/walt_vehicle.pth') |
| | filenames = sorted(glob.glob('demo/images/*')) |
| | count = 0 |
| | for filename in filenames: |
| | img=cv2.imread(filename) |
| | try: |
| | img = detect_people.run_on_image(img) |
| | img = detect.run_on_image(img) |
| | except: |
| | continue |
| | count=count+1 |
| |
|
| | try: |
| | import os |
| | os.makedirs(os.path.dirname(filename.replace('demo','demo/results/'))) |
| | os.mkdirs(os.path.dirname(filename)) |
| | except: |
| | print('done') |
| | cv2.imwrite(filename.replace('demo','demo/results/'),img) |
| | if count == 30000: |
| | break |
| | try: |
| | detect.process_output(count) |
| | except: |
| | continue |
| | ''' |
| | |
| | np.savez('FC', a= detect.record_dict) |
| | with open('check.json', 'w') as f: |
| | json.dump(detect.record_dict, f) |
| | detect.write_json('seq3') |
| | asas |
| | detect.process_output(0) |
| | ''' |
| | if __name__ == "__main__": |
| | main() |
| |
|
