def predict(self, path, is_output_polygon=False, short_size=1024, output_folder='../output'): img = cv2.imread(path, 1 if self.img_mode != 'GRAY' else 0) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w = img.shape[:2] img = resize_image(img, short_size) tensor = self.transform(img) tensor = tensor.unsqueeze_(0) tensor = tensor.to(self.device) batch = {'shape': [(h, w)]} with torch.no_grad(): if str(self.device).__contains__('cuda'): torch.cuda.synchronize(self.device) preds = self.model(tensor) if str(self.device).__contains__('cuda'): torch.cuda.synchronize(self.device) box_list, _ = self.post_process(batch, preds, is_output_polygon=is_output_polygon) box_list = box_list[0] if len(box_list) > 0: if is_output_polygon: idx = [x.sum() > 0 for x in box_list] box_list = [box_list[i] for i, v in enumerate(idx) if v] else: idx = box_list.reshape(box_list.shape[0], -1).sum(axis=1) > 0 box_list = box_list[idx] else: box_list, score_list = [], [] img = draw_bbox(cv2.imread(path)[:, :, ::-1], box_list) os.makedirs(output_folder, exist_ok=True) img_path = pathlib.Path(path) output_path = os.path.join(output_folder, img_path.stem + '_result.jpg') cv2.imwrite(output_path, img[:, :, ::-1]) return output_path, box_list
def main_debug(): os.environ['CUDA_VISIBLE_DEVICES'] = '7' import pathlib from tqdm import tqdm import matplotlib.pyplot as plt from utils.util import show_img, draw_bbox, save_result, get_file_list model_path = '/red_detection/DBNet/DBNet_fzh/phone_code_model/model_0.88_depoly.pth' path = './测试图片' output_path = './测试图片_条形码检测' debug = True # 保存结果到路径 os.makedirs(output_path, exist_ok=True) imgs_list_path = [os.path.join(path, i) for i in os.listdir(path)] # 初始化网络 model = Code_model(model_path) TIMes = [] for i, img_list_path in enumerate(imgs_list_path): # if i<1: print('==img_list_path:', img_list_path) img = cv2.imread(img_list_path) st = time.time() preds, boxes_list, score_list = model.predict(img, is_output_polygon=False) print('每张图片分割耗时:{}'.format(time.time() - st)) TIMes.append(time.time() - st) if debug: img = draw_bbox(img, boxes_list) pred_path = os.path.join( output_path, img_list_path.split('/')[-1].split('.')[0] + '_pred.jpg') cv2.imwrite( os.path.join( output_path, img_list_path.split('/')[-1].split('.')[0] + '.jpg'), img) cv2.imwrite(pred_path, preds * 255)
def detect_on_dataset_ui(model): sg.theme(cfg.UI.THEME) layout = [ [sg.Text('Dataset Folder Path:')], [ sg.Input(key='-INPUT-PATH-', enable_events=True), sg.FolderBrowse(initial_folder='./') ], [sg.Text('Prediction Output Path:')], [sg.Input(key='-OUTPUT-PATH-'), sg.FolderBrowse(initial_folder='./')], [sg.Text('Annotation Output Path')], [sg.Input(key='-ANNO-PATH-'), sg.FileSaveAs(initial_folder='./')], [sg.Text('Input File Type:')], [ sg.Combo(('.jpg', '.png', '.bmp', '.tif', '.tiff', '.gif'), readonly=True, default_value='.jpg', size=(15, 1), key='-FILE-TYPE-') ], [ sg.Text('Confidence Threshold'), sg.Slider((0, 1), 0.4, 0.05, orientation='h', key='-CONF-SLIDER-'), sg.Text('NMS IOU Threshold'), sg.Slider((0, 1), 0.45, 0.05, orientation='h', key='-NMS-SLIDER-') ], [sg.B('Detect'), sg.B('Exit')] ] window = Window('Detect Dataset', layout, font=(cfg.UI.FONT, 12)) while True: event, value = window.read() if event in ['Exit', None]: break # Update output path and annotation file path according to input path if event in ['-INPUT-PATH-']: window['-OUTPUT-PATH-'].update( os.path.join(value['-INPUT-PATH-'], 'Prediction')) window['-ANNO-PATH-'].update( os.path.join(value['-INPUT-PATH-'], 'annotation.txt')) # Start detection if event in ['Detect']: input_path = value['-INPUT-PATH-'] output_path = value['-OUTPUT-PATH-'] anno_path = value['-ANNO-PATH-'] file_type = value['-FILE-TYPE-'] nms_iou = value['-NMS-SLIDER-'] conf_thresh = value['-CONF-SLIDER-'] anno_file = None try: anno_file = open(anno_path, 'w') images = os.listdir(input_path) if not os.path.exists(output_path): os.mkdir(output_path) # Count file conform the given file type max_process = [x.endswith(file_type) for x in images].count(True) cur_process = 1 window.disable() for image in images: if image.endswith(file_type): # If user push cancel button on the progress popup if not sg.one_line_progress_meter( 'Generating Prediction', cur_process, max_process, key='-PROGRESS-'): break img_RGB = io.imread(os.path.join(input_path, image)) img = cv2.cvtColor(img_RGB, cv2.COLOR_RGB2BGR) # Predict bboxes, _ = predict(img_RGB, model, nms_iou, conf_thresh) pred_img = util.draw_bbox(img.copy(), bboxes) cv2.imwrite(os.path.join(output_path, image), pred_img) anno_line = util.encode_annotation(bboxes, image) anno_file.write(anno_line.strip() + '\n') cur_process += 1 except IOError as e: sg.PopupError("Invalid input file,\n{}".format(e)) sg.one_line_progress_meter('Generating Prediction', 1, 1, key='-PROGRESS-') finally: if anno_file is not None: anno_file.close() window.enable() sg.popup_ok("Detection finished", title="Finish") window.close() del window del layout
# 初始化网络 model = DBNet(args.model_path, post_p_thre=args.thre, gpu_id=0, imageH=960, imageW=480) img_folder = pathlib.Path(args.input_folder) for img_path in tqdm(get_file_list(args.input_folder, p_postfix=['.jpg'])): img = cv2.imread(img_path) # if img.shape[0] /img.shape[1] > 2 or img.shape[1]/img.shape[0] > 2: # continue preds, boxes_list, score_list, t = model.predict( img_path, is_output_polygon=args.polygon, runtime='trt') print('time cost: {}s'.format(t)) crops = crop_bbox(img[:, :, ::-1], boxes_list) img = draw_bbox(img[:, :, ::-1], boxes_list) if args.show: show_img(preds) show_img(img, title=os.path.basename(img_path)) plt.show() # 保存结果到路径 os.makedirs(args.output_folder, exist_ok=True) img_path = pathlib.Path(img_path) output_path = os.path.join(args.output_folder, img_path.stem + '_result.jpg') pred_path = os.path.join(args.output_folder, img_path.stem + '_pred.jpg') cv2.imwrite(output_path, img[:, :, ::-1]) cv2.imwrite(pred_path, preds * 255) for i, crop in enumerate(crops): cv2.imwrite(
start = time.time() preds = self.net(tensor)[0] print('net time:', time.time() - start) preds, boxes_list = decode(preds) scale = (preds.shape[1] / w, preds.shape[0] / h) if len(boxes_list): boxes_list = boxes_list / scale t = time.time() - start return preds, boxes_list, t if __name__ == '__main__': import matplotlib.pyplot as plt from utils.util import show_img, draw_bbox os.environ['CUDA_VISIBLE_DEVICES'] = str('2') model_path = 'output/PAN_gt_mask_resnet50/checkpoint/model_best.pth' img_id = 10 img_path = '/data1/zj/ocr/icdar2015/test/img/img_{}.jpg'.format(img_id) # 初始化网络 model = Pytorch_model(model_path, gpu_id=None) preds, boxes_list, t = model.predict(img_path) show_img(cv2.imread(img_path)[:, :, ::-1], color=True) show_img(preds) img = draw_bbox(img_path, boxes_list) show_img(img, color=True) plt.show()
torch.cuda.synchronize(self.device) box_list, score_list = self.post_process(batch, preds) box_list, score_list = box_list[0], score_list[0] if len(box_list) > 0: idx = box_list.reshape(box_list.shape[0], -1).sum(axis=1) > 0 box_list, score_list = box_list[idx], score_list[idx] else: box_list, score_list = [], [] t = time.time() - start return preds[0, 0, :, :].detach().cpu().numpy(), box_list, t if __name__ == '__main__': import matplotlib.pyplot as plt from utils.util import show_img, draw_bbox os.environ['CUDA_VISIBLE_DEVICES'] = str('0') model_path = '../output/DBNet_resnet18_FPN_DBHead/checkpoint/DBNet_latest.pth' img_id = 10 img_path = 'E:/zj/dataset/icdar2015/test/img/img_{}.jpg'.format(img_id) # 初始化网络 model = Pytorch_model(model_path, gpu_id=0) preds, boxes_list, t = model.predict(img_path) show_img(preds) img = draw_bbox(cv2.imread(img_path)[:, :, ::-1], boxes_list) show_img(img) plt.show()
def annotation_ui(): global amended sg.theme(cfg.UI.THEME) col1 = [ [sg.Button('Open', size=(10, 1)), sg.Button('Save', size=(10, 1), disabled=True), sg.Button('Save as..', size=(10, 1), disabled=True), sg.Button('Exit', size=(10, 1))], [sg.Image(data=placeholder, key='-SCREEN-')] ] col2 = [ [sg.Text('Images')], [sg.Listbox([], size=(20, 15), key='-IMAGE-', enable_events=True)], [sg.Text('Objects')], [sg.Listbox([], size=(20, 5), key='-OBJECT-', enable_events=True)], [sg.Button('Add', size=(10, 1), disabled=True), sg.Button('Delete', size=(10, 1), disabled=True)] ] layout = [ [sg.Column(col1), sg.Column(col2, pad=(0, 10))] ] window = Window('Annotation', layout=layout, font=(cfg.UI.FONT, 12)) anno_data = None anno_path = None def on_close(): global amended if amended: confirm = sg.popup_yes_no('Do you want to exit without saving your amendment?', title="Exit..") if confirm in ['No', None]: return if window.CurrentlyRunningMainloop: amended = False window.TKroot.quit() window.TKroot.destroy() window.RootNeedsDestroying = True window.TKrootDestroyed = True window.finalize() window.TKroot.protocol('WM_DELETE_WINDOW', on_close) while True: event, value = window.Read() if event in ['Exit', None]: if amended: confirm = sg.popup_yes_no('Do you want to exit without saving your amendment?', title="Exit..") if confirm in ['No', None]: continue break # Load annotation file if event in ['Open']: annotation_file = None anno_path = sg.popup_get_file('Annotation file', 'Open the annotation file:', file_types=(('Annotation file', 'annotation.txt'), ('All files', '*.*')), initial_folder='./') if anno_path is None: continue dir_path = os.path.dirname(anno_path) try: annotation_file = open(anno_path, 'r+') annotations = annotation_file.readlines() anno_data = [] for anno in annotations: anno = anno.strip() bboxes, img_path = util.parse_annotation(anno) img_path = os.path.join(dir_path, img_path) anno_data.append([img_path, bboxes]) if len(anno_data) == 0: raise IOError() anno_data = np.array(anno_data) window['-IMAGE-'].update(values=[os.path.basename(x) for x in anno_data[:, 0]], set_to_index=0) window['-OBJECT-'].update(values=range(1, len(anno_data[0][1]) + 1)) if len(anno_data[0][1]) > 0: window['-OBJECT-'].update(set_to_index=0) amended = False except IOError: sg.PopupError('Invalid Annotation') anno_data = None except ValueError: sg.PopupError('Invalid Annotation format') anno_data = None finally: if annotation_file is not None: annotation_file.close() if anno_data is not None: if len(window['-IMAGE-'].GetIndexes()) > 0: image_idx = window['-IMAGE-'].GetIndexes()[0] else: continue if event in ['-IMAGE-']: window['-OBJECT-'].update(values=range(1, len(anno_data[image_idx][1]) + 1)) if len(anno_data[image_idx][1]) > 0: window['-OBJECT-'].update(set_to_index=0) # Check number of objects in this image if len(window['-OBJECT-'].GetIndexes()) > 0: object_idx = window['-OBJECT-'].GetIndexes()[0] else: object_idx = -1 window['Delete'].update(disabled=True) # Load image if anno_data[image_idx][0].endswith('.tif'): img_RGB = io.imread(anno_data[image_idx][0]) img = cv2.cvtColor(img_RGB, cv2.COLOR_RGB2BGR) else: img = cv2.imread(anno_data[image_idx][0]) # Add button if event in ['Add']: img_to_draw = util.draw_bbox(img.copy(), anno_data[image_idx][1]) input_coor = util.get_bbox_from_UI(img_to_draw) if input_coor is not None: anno_data[image_idx][1].append([input_coor[0], input_coor[1], input_coor[2], input_coor[3], 1.0, 0]) window['-OBJECT-'].update(values=range(1, len(anno_data[image_idx][1]) + 1)) # Delete button if event in ['Delete']: anno_data[image_idx][1].pop(object_idx) window['-OBJECT-'].update(values=range(1, len(anno_data[image_idx][1]) + 1)) # Check number of objects in this image again after amending if event in ['Add', 'Delete']: amended = True if len(window['-OBJECT-'].GetIndexes()) > 0: object_idx = window['-OBJECT-'].GetIndexes()[0] else: object_idx = -1 window['Delete'].update(disabled=True) # Save and Save as if event in ['Save']: util.save_anno(anno_data, anno_path) amended = False sg.popup_ok('Annotation saved', title='Annotation saved') if event in ['Save as..']: save_as_path = sg.popup_get_file('Select file path to save annotation', 'Save as..', default_extension='txt', save_as=True, file_types=( ('Annotation file', 'annotation.txt'), ('All files', '*.*'))) if save_as_path is not None: util.save_anno(anno_data, save_as_path, copy_image=True) anno_path = save_as_path amended = False sg.popup_ok('Annotation saved at ' + save_as_path, title='Annotation saved') # Display image and bounding boxes draw_img = util.draw_bbox(img.copy(), anno_data[image_idx][1], highlight=object_idx) img_reshaped = draw_img if draw_img.shape[1] > 800: img_reshaped = cv2.resize(draw_img, (800, 800 * img.shape[0] // img.shape[1],)) img_bytes = cv2.imencode('.png', img_reshaped)[1].tobytes() window['-SCREEN-'].update(data=img_bytes) window['Save'].update(disabled=not amended) window['Save as..'].update(disabled=anno_data is None) window['Add'].update(disabled=anno_data is None) window['Delete'].update(disabled=anno_data is None) window.close() del window del layout del col1 del col2
def realtime_detect_ui(model): sg.theme(cfg.UI.THEME) layout = [[sg.Image('', key='-SCREEN-')], [sg.Text('', size=(30, 1), key='-PROCESS-TIME-')], [ sg.Text('Confidence Threshold'), sg.Slider((0, 1), 0.4, 0.05, orientation='h', key='-CONF-SLIDER-'), sg.Text('NMS IOU Threshold'), sg.Slider((0, 1), 0.45, 0.05, orientation='h', key='-NMS-SLIDER-') ], [ sg.Button('Record', size=(10, 1), key='-REC-'), sg.Button('Collect', size=(10, 1), key='-COL-'), sg.Button('Close', size=(10, 1)) ]] cam_no = sg.popup_get_text('Detection device number:', default_text='0') if cam_no is None: return # Initialise resource variables cam = cv2.VideoCapture() window = Window('Realtime Detect', layout=layout, finalize=True, font=cfg.UI.FONT) record = False collect = False anno_file = None collect_path = None out = None anno_index = 1 try: cam_no = int(cam_no) cam.open(cam_no) while True: event, value = window.read(timeout=0) flag, img = cam.read() if event in [None, 'Close']: break # Record button pushed if event in ['-REC-']: record = not record window['-REC-'].update(text=('Stop' if record else 'Record')) # Collect button pushed if event in ['-COL-']: if collect: anno_file.flush() collect = not collect window['-COL-'].update(text=('Stop' if collect else 'Collect')) # Read from Webcam successfully if flag: # Predict bboxes img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) bboxes, exe_time = predict(img_RGB, model, value['-NMS-SLIDER-'], value['-CONF-SLIDER-']) # Draw bbox to img pred_img = util.draw_bbox(img.copy(), bboxes) # Record video to the file if record: if out is None: out = cv2.VideoWriter( time.strftime('%Y%m%d-%H%M', time.localtime(time.time())) + '.mp4', -1, 20.0, (640, 480)) out.write(pred_img) # Collect image and annotation if collect: if collect_path is None: collect_path = './anno_data/collected_data_' + time.strftime( '%Y%m%d-%H%M', time.localtime(time.time())) os.mkdir(collect_path) os.mkdir(os.path.join(collect_path, 'img')) if anno_file is None: anno_file = open( os.path.join(collect_path, 'annotation.txt'), 'a') anno_index = 1 img_path = os.path.join('img', str(anno_index) + '.jpg') cv2.imwrite(os.path.join(collect_path, img_path), img) anno_line = util.encode_annotation(bboxes, img_path) anno_file.write(anno_line.strip() + '\n') anno_index += 1 # Display predicted image img_bytes = cv2.imencode('.png', pred_img)[1].tobytes() window['-SCREEN-'].update(data=img_bytes) window['-PROCESS-TIME-'].update( value='Process time: {}ms'.format(int(exe_time * 1000))) else: raise ValueError except ValueError: sg.popup_error("Selected Device is unavailable") return finally: cam.release() if out is not None: out.release() if anno_file is not None: anno_file.close() window.close() del cam del window del layout
boxes_list.append(box) t = time.time() - start + t boxes_list = np.array(boxes_list) return dilated_polys, boxes_list, t if __name__ == '__main__': from utils.util import show_img, draw_bbox os.environ["CUDA_VISIBLE_DEVICES"] = "2" model_path = './output/DB_shufflenetv2_FPN/checkpoint/DBNet_best_loss.pth' ## icdar 2013 / 2015 img_id = 130 img_path = '/home1/surfzjy/data/ic13/test_images/img_{}.jpg'.format(img_id) # 初始化网络 model = Pytorch_model(model_path, gpu_id = 0) ## set GPU id contours, boxes_list, t = model.predict(img_path) print('Time: %.4f' %t) img = cv2.imread(img_path)[:, :, ::-1] imgc = img.copy() cv2.drawContours(imgc, contours, -1, (22,222,22), 2, cv2.LINE_AA) cv2.imwrite('contour.png', imgc) img = draw_bbox(img, boxes_list) cv2.imwrite('predict.jpg', img)