YOLOX保存txt推理结果

本文讲述了如何在YOLOX中复制YOLOv5的推理过程,并指导用户修改`demo.py`和`visualize.py`中的关键函数以获取相似输出。
摘要由CSDN通过智能技术生成

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

觉得好的话希望大家给我点赞关注收藏哦。

  • 1
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值