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)
