AndreyGermanov / yolov8_segmentation_python

YOLOv8 image segmentation through ONNX in Python
GNU General Public License v3.0
28 stars 3 forks source link

About a class that can support both PyTorch Model and ONNX model #3

Open curtis18 opened 1 week ago

curtis18 commented 1 week ago

Thank you for your tutorial. I would like to ask can the output of boxes and masks from ONNX model be the same as the output of PyTorch model? So that a classs can be created to support both ONNX and PyTorch model output?

                # YOLO
                results = self.model.predict(source=image, device=device)
                boxes = results[0].boxes.data.cpu().numpy()
                masks = results[0].masks.data.cpu().numpy() if results[0].masks else None
                # ONNX
                outputs = self.model.run(None, {input_name: image_np})
                boxes, masks = self.model_onnx_seg_convert(outputs)  # a new function