How to infer yolov6 model for single image

[mallapraveen]

i need a python script which take single image and do the prediction on it and return the results. can someone help on this

[rjshrivastav]

try using infer.py and pass your image path in source argument. for eg. python tools/infer.py --weights <saved_model_path> --source <img path>

[mallapraveen]

@Rjshrivastav that is one way of doing it, I need script which take a single image and return the bounding boxes and class labels.

[MTChengMeng]

@Rjshrivastav that is one way of doing it, I need script which take a single image and return the bounding boxes and class labels.

Hi, you can add --save-txt option to save predicted scores and labels of yolo format.

[mallapraveen]

@MTChengMeng do we have a predict function which takes image as input and gives the output as bboxes and labels. I have currently trained the model and want function to integrate with my codebase. A function would help instead of invoking like this 'python tools/infer.py --weights <saved_model_path> --source '

[rjshrivastav]

@mallapraveen you can useyolov6/core/inferer.py that contains all the necessary functions.

[sssssshf]

@Rjshrivastav 这是一种方法，我需要一个脚本来获取单个图像并返回边界框和类标签。

改写一下infer的代码。 def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf): ''' Model Inference and results visualization '''

    for img_path in tqdm(self.img_paths):
        img, img_src = self.precess_image(img_path, self.img_size, self.stride, self.half)
        img = img.to(self.device)
        if len(img.shape) == 3:
            img = img[None]
            # expand for batch dim
        pred_results = self.model(img)
        det = non_max_suppression(pred_results, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)[0]
        if len(det):
            det[:, :4] = self.rescale(img.shape[2:], det[:, :4], img_src.shape).round()
            return img_src, det
        else:
            # return img_src,torch.tensor(np.zeros((1,6)))
            return img_src,np.zeros((1,6))

[mallapraveen]

@SongHfei Thanks