YOLOX保存像YOLOv5一样的推理结果
1.第一步,在tools/demo.py文件中替换两个函数
def visual(self, output, img_info, cls_conf=0.35, image_name= None, vis_folder= None, current_time= None):
ratio = img_info["ratio"]
img = img_info["raw_img"]
if output is None:
return img
output = output.cpu()
bboxes = output[:, 0:4]
# preprocessing: resize
bboxes /= ratio
cls = output[:, 6]
scores = output[:, 4] * output[:, 5]
vis_res = vis(img, bboxes, scores, cls, cls_conf, self.cls_names, image_name, vis_folder, current_time)
return vis_res
def image_demo(predictor, vis_folder, path, current_time, save_result):
if os.path.isdir(path):
files = get_image_list(path)
else:
files = [path]
files.sort()
for image_name in files:
outputs, img_info = predictor.inference(image_name)
# print(len(predictor.visual(outputs[0], img_info, predictor.confthre)))
result_image= predictor.visual(outputs[0], img_info, predictor.confthre, image_name = image_name, vis_folder = vis_folder, current_time = current_time)
if save_result:
save_folder = os.path.join(
vis_folder, time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
)
os.makedirs(save_folder, exist_ok=True)
save_file_name = os.path.join(save_folder, os.path.basename(image_name))
logger.info("Saving detection result in {}".format(save_file_name))
save_labels_path = os.path.join(save_folder, 'lables')
cv2.imwrite(save_file_name, result_image)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord("q") or ch == ord("Q"):
break
2.第二部,替换yolox/utils/visualize.py中的函数
def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None, image_name=None, vis_folder=None, current_time = None):
result_list = []
for i in range(len(boxes)):
box = boxes[i]
cls_id = int(cls_ids[i])
score = scores[i]
if score < conf:
continue
x0 = int(box[0])
y0 = int(box[1])
x1 = int(box[2])
y1 = int(box[3])
class_name = class_names[cls_id]
one_line = (str(x0), str(y0), str(x1), str(y1), class_name, str(float(score)))
#save_yolo
h, w,c = img.shape
x_c = float((x0 +x1)/2.0/w)
y_c = float((y0 +y1)/2.0/h)
b_w = float(abs(x0-x1)*1.0/w)
b_h = float(abs(y0-y1)*1.0/h)
save_str = (str(cls_id), str(x_c), str(y_c), str(b_w), str(b_h), str(float(score)))
str_one_line = " ".join(save_str)
result_list.append(str_one_line)
######################
color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
text = '{}:{:.1f}%'.format(class_name, score * 100)
txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
font = cv2.FONT_HERSHEY_SIMPLEX
txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()
cv2.rectangle(
img,
(x0, y0 + 1),
(x0 + txt_size[0] + 1, y0 + int(1.5 * txt_size[1])),
txt_bk_color,
-1
)
cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)
save_folder = os.path.join(
vis_folder, time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
)
os.makedirs(save_folder, exist_ok=True)
save_labels_path = os.path.join(save_folder, 'lables')
os.makedirs(save_labels_path, exist_ok=True)
save_label_name = os.path.join(save_labels_path, os.path.basename(image_name[:-3] +'txt'))
f =open(save_label_name, 'w')
for line in result_list:
f.write(str(line) +'\n')
f.close()
return img
觉得好的话希望大家给我点赞关注收藏哦。