def run_demo(cfg, weights_file, iou_threshold, score_threshold, images_dir, output_dir, dataset_type):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
else:
raise NotImplementedError('Not implemented now.')
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg, is_test=True)
model.load(weights_file)
print('Loaded weights from {}.'.format(weights_file))
model = model.to(device)
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
output = predictor.predict(image)
boxes, labels, scores = [o.to(cpu_device).numpy() for o in output]
drawn_image = draw_bounding_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
image_name = os.path.basename(image_path)
Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
def _log_images_tensorboard(cfg, global_step, images, orig_boxes, orig_labels, summary_writer, j_images=None):
from random import randrange
size = images.shape[0]
i = randrange(size)
imagei = images[i]
imagei = imagei.cpu()
labels_i = orig_labels[i]
labels_i = labels_i.numpy()
nz_indices_i = (np.nonzero(labels_i))[0]
boxes_i = orig_boxes[i]
boxes_i = boxes_i.numpy()
boxes_i = boxes_i[nz_indices_i]
labels_i = labels_i[nz_indices_i]
labels_i = labels_i.astype(int)
boxes_i = boxes_i * cfg.INPUT.IMAGE_SIZE
boxes_i = boxes_i.astype(int)
rescaled = _scale_back_image(cfg, imagei)
drawn_image = draw_bounding_boxes(rescaled, boxes_i, labels=labels_i,
class_name_map=VOCDataset.class_names).astype(np.uint8)
drawn_image = np.transpose(drawn_image, (2, 0, 1))
summary_writer.add_image('imaugmall', drawn_image, global_step=global_step)
if cfg.MODEL.SELF_SUPERVISED:
j_imagei = j_images[i]
j_imagei = j_imagei.cpu()
rescaled2 = _scale_back_image(cfg, j_imagei)
drawn_image2 = np.transpose(rescaled2, (2, 0, 1))
summary_writer.add_image('aux_task_image', drawn_image2, global_step=global_step)
def run_demo(cfg, weights_file, iou_threshold, score_threshold, video_dir, output_dir, dataset_type):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
elif dataset_type == 'cla':
class_names = CLADataset.class_names
else:
raise NotImplementedError('Not implemented now.')
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
model.load(weights_file)
print('Loaded weights from {}.'.format(weights_file))
model = model.to(device)
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
stream = cv2.VideoCapture(video_dir)
# image_paths = glob.glob(os.path.join(video_dir, '*.jpg'))
# 获得输出视频大小，与原视频大小相同
shape=(int(stream.get(cv2.CAP_PROP_FRAME_WIDTH)),int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT)))
# 获取输出视频的帧率
_fps = stream.get(cv2.CAP_PROP_FPS)
# 指定视频编码
fourcc = cv2.VideoWriter_fourcc(*'XVID')
# 如果输出目录不存在 创建输出目录
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 创建视频输出对象
output_name = os.path.basename(video_dir)
output_name = output_name.split('.')[0] + ".avi"
writer = cv2.VideoWriter(os.path.join(output_dir, output_name), fourcc, _fps, shape)
while True:
ret, image = stream.read()
if ret is False :
break
output = predictor.predict(image)
boxes, labels, scores = [o.to(cpu_device).numpy() for o in output]
drawn_image = draw_bounding_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
writer.write(drawn_image)
writer.release()
stream.release()
def run_demo(cfg, checkpoint_file, iou_threshold, score_threshold, images_dir,
output_dir):
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['state_dict'])
print('Loaded weights from {}.'.format(checkpoint_file))
model = model.to(device)
model.eval()
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
add_count = 0
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
# image_mirror = image[:, ::-1]
output = predictor.predict(image)
boxes, scores, seg_map = [o.to(cpu_device).numpy() for o in output]
seg_map = cv2.resize(seg_map, (512, 512)) * 255
seg_map = seg_map.astype(np.uint8)
# seg_map = cv2.applyColorMap(seg_map, cv2.COLORMAP_JET)
seg_map = cv2.resize(seg_map, (1280, 720),
interpolation=cv2.INTER_CUBIC)
drawn_image = draw_bounding_boxes(image, boxes).astype(np.uint8)
image_name = os.path.basename(image_path)
txt_path = os.path.join(output_dir, 'txtes')
if not os.path.exists(txt_path):
os.makedirs(txt_path)
txt_path = os.path.join(txt_path,
'res_' + image_name.replace('jpg', 'txt'))
#multi-output merge
merge_output = False
if merge_output:
ret, binary = cv2.threshold(seg_map, 75, 255, cv2.THRESH_BINARY)
# cv2.imshow('binary:',binary)
# cv2.waitKey()
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
w, h = np.shape(binary)
for contour in contours:
# 获取最小包围矩形
rect = cv2.minAreaRect(contour)
# 中心坐标
x, y = rect[0]
# cv2.circle(img, (int(x), int(y)), 3, (0, 255, 0), 5)
# 长宽,总有 width>=height
width, height = rect[1]
if width < 10 or height < 10:
continue
# 角度:[-90,0)
angle = rect[2]
box = cv2.boxPoints(rect)
box = np.int0(box)
box[:, 0] = np.clip(box[:, 0], 0, h)
box[:, 1] = np.clip(box[:, 1], 0, w)
poly1 = Polygon(box).convex_hull
intersect = False
for item in boxes:
print('item:', item)
poly2 = Polygon(item.reshape(4, 2)).convex_hull
if poly1.intersects(poly2): # 如果两四边形相交
intersect = True
break
if not intersect:
print('boxes.shape:', np.shape(boxes))
box = box.reshape((1, 8))
print('box.shape:', np.shape(box))
num, _ = np.shape(boxes)
if num == 0:
print('num == 0')
boxes = box
else:
boxes = np.concatenate((boxes, box))
print('boxes.shape:', np.shape(boxes))
print('add one box')
add_count += 1
# cv2.line(image, (box[0][0], box[0][1]), (box[0][2], box[0][3]), (0, 0, 255), thickness=4)
# cv2.line(image,(box[0][2], box[0][3]), (box[0][4], box[0][5]), (0, 0, 255), thickness=4)
# cv2.line(image,(box[0][4], box[0][5]), (box[0][6], box[0][7]), (0, 0, 255), thickness=4)
# cv2.line(image, (box[0][6], box[0][7]), (box[0][0], box[0][1]), (0, 0, 255), thickness=4)
# cv2.imshow('img',image)
# cv2.waitKey()
# print('txt_path:',txt_path)
with open(txt_path, 'w+') as f:
for box in boxes:
box_temp = np.reshape(box, (4, 2))
box = order_points_quadrangle(box_temp)
box = np.reshape(box, (1, 8)).squeeze(0)
is_valid = validate_clockwise_points(box)
if not is_valid:
continue
# print('box:',box)
line = ''
for item in box:
if item < 0:
item = 0
line += str(int(item)) + ','
line = line[:-1] + '\n'
f.write(line)
path = os.path.join(output_dir, image_name)
print('path:', path)
Image.fromarray(drawn_image).save(path)
path = os.path.join(
output_dir,
image_name.split('.')[0] + '_segmap.' + image_name.split('.')[1])
# print(path)
# 存储score_map
cv2.imwrite(path, seg_map)
print('add count:', add_count)
def run_demo(cfg, checkpoint_file, iou_threshold, score_threshold, images_dir,
output_dir):
basename = os.path.basename(checkpoint_file).split('.')[0]
epoch = basename[21:]
epoch = int(epoch)
if epoch < min_epoch:
return
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['state_dict'])
print('Loaded weights from {}.'.format(checkpoint_file))
model = model.to(device)
model.eval()
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
# image = image[:, ::-1]
#image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
output = predictor.predict(image)
boxes, scores, score_map = [o.to(cpu_device).numpy() for o in output]
# score_map = cv2.resize(score_map, (512, 512)) * 255
# score_map = score_map.astype(np.uint8)
# score_map = cv2.applyColorMap(score_map, cv2.COLORMAP_JET)
# score_map = cv2.resize(score_map,(1280,720),interpolation=cv2.INTER_CUBIC)
drawn_image = draw_bounding_boxes(image, boxes).astype(np.uint8)
image_name = os.path.basename(image_path)
txt_path = os.path.join(output_dir, 'txtes')
if not os.path.exists(txt_path):
os.makedirs(txt_path)
txt_path = os.path.join(txt_path,
'res_' + image_name.replace('jpg', 'txt'))
# print('txt_path:',txt_path)
with open(txt_path, 'w+') as f:
for box in boxes:
box_temp = np.reshape(box, (4, 2))
box = order_points_quadrangle(box_temp)
box = np.reshape(box, (1, 8)).squeeze(0)
is_valid = validate_clockwise_points(box)
if not is_valid:
continue
# print('box:',box)
line = ''
for item in box:
if item < 0:
item = 0
line += str(int(item)) + ','
line = line[:-1] + '\n'
f.write(line)
# path = os.path.join(output_dir, image_name)
# Image.fromarray(drawn_image).save(path)
# path = os.path.join(output_dir, image_name.split('.')[0]+'_score_map.'+image_name.split('.')[1])
# print(path)
# cv2.imwrite(path,score_map)
submit_path = MyZip(
'/home/binchengxiong/ssd_fcn_multitask_text_detection_pytorch1.0/demo/result223'
+ str(min_epoch) + '/txtes', epoch)
hmean_this_epoch = compute_hmean(submit_path, epoch)