ultralytics 8.0.239 Ultralytics Actions and hub-sdk adoption (#7431)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com> Co-authored-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
parent
e795277391
commit
fe27db2f6e
139 changed files with 6870 additions and 5125 deletions
|
|
@ -21,18 +21,21 @@ class YOLOv8Seg:
|
|||
"""
|
||||
|
||||
# Build Ort session
|
||||
self.session = ort.InferenceSession(onnx_model,
|
||||
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
|
||||
if ort.get_device() == 'GPU' else ['CPUExecutionProvider'])
|
||||
self.session = ort.InferenceSession(
|
||||
onnx_model,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
if ort.get_device() == "GPU"
|
||||
else ["CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
# Numpy dtype: support both FP32 and FP16 onnx model
|
||||
self.ndtype = np.half if self.session.get_inputs()[0].type == 'tensor(float16)' else np.single
|
||||
self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single
|
||||
|
||||
# Get model width and height(YOLOv8-seg only has one input)
|
||||
self.model_height, self.model_width = [x.shape for x in self.session.get_inputs()][0][-2:]
|
||||
|
||||
# Load COCO class names
|
||||
self.classes = yaml_load(check_yaml('coco128.yaml'))['names']
|
||||
self.classes = yaml_load(check_yaml("coco128.yaml"))["names"]
|
||||
|
||||
# Create color palette
|
||||
self.color_palette = Colors()
|
||||
|
|
@ -60,14 +63,16 @@ class YOLOv8Seg:
|
|||
preds = self.session.run(None, {self.session.get_inputs()[0].name: im})
|
||||
|
||||
# Post-process
|
||||
boxes, segments, masks = self.postprocess(preds,
|
||||
im0=im0,
|
||||
ratio=ratio,
|
||||
pad_w=pad_w,
|
||||
pad_h=pad_h,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
nm=nm)
|
||||
boxes, segments, masks = self.postprocess(
|
||||
preds,
|
||||
im0=im0,
|
||||
ratio=ratio,
|
||||
pad_w=pad_w,
|
||||
pad_h=pad_h,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
nm=nm,
|
||||
)
|
||||
return boxes, segments, masks
|
||||
|
||||
def preprocess(self, img):
|
||||
|
|
@ -98,7 +103,7 @@ class YOLOv8Seg:
|
|||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
|
||||
|
||||
# Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)
|
||||
img = np.ascontiguousarray(np.einsum('HWC->CHW', img)[::-1], dtype=self.ndtype) / 255.0
|
||||
img = np.ascontiguousarray(np.einsum("HWC->CHW", img)[::-1], dtype=self.ndtype) / 255.0
|
||||
img_process = img[None] if len(img.shape) == 3 else img
|
||||
return img_process, ratio, (pad_w, pad_h)
|
||||
|
||||
|
|
@ -124,7 +129,7 @@ class YOLOv8Seg:
|
|||
x, protos = preds[0], preds[1] # Two outputs: predictions and protos
|
||||
|
||||
# Transpose the first output: (Batch_size, xywh_conf_cls_nm, Num_anchors) -> (Batch_size, Num_anchors, xywh_conf_cls_nm)
|
||||
x = np.einsum('bcn->bnc', x)
|
||||
x = np.einsum("bcn->bnc", x)
|
||||
|
||||
# Predictions filtering by conf-threshold
|
||||
x = x[np.amax(x[..., 4:-nm], axis=-1) > conf_threshold]
|
||||
|
|
@ -138,7 +143,6 @@ class YOLOv8Seg:
|
|||
|
||||
# Decode and return
|
||||
if len(x) > 0:
|
||||
|
||||
# Bounding boxes format change: cxcywh -> xyxy
|
||||
x[..., [0, 1]] -= x[..., [2, 3]] / 2
|
||||
x[..., [2, 3]] += x[..., [0, 1]]
|
||||
|
|
@ -173,13 +177,13 @@ class YOLOv8Seg:
|
|||
segments (List): list of segment masks.
|
||||
"""
|
||||
segments = []
|
||||
for x in masks.astype('uint8'):
|
||||
for x in masks.astype("uint8"):
|
||||
c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[0] # CHAIN_APPROX_SIMPLE
|
||||
if c:
|
||||
c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
|
||||
else:
|
||||
c = np.zeros((0, 2)) # no segments found
|
||||
segments.append(c.astype('float32'))
|
||||
segments.append(c.astype("float32"))
|
||||
return segments
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -219,7 +223,7 @@ class YOLOv8Seg:
|
|||
masks = np.matmul(masks_in, protos.reshape((c, -1))).reshape((-1, mh, mw)).transpose(1, 2, 0) # HWN
|
||||
masks = np.ascontiguousarray(masks)
|
||||
masks = self.scale_mask(masks, im0_shape) # re-scale mask from P3 shape to original input image shape
|
||||
masks = np.einsum('HWN -> NHW', masks) # HWN -> NHW
|
||||
masks = np.einsum("HWN -> NHW", masks) # HWN -> NHW
|
||||
masks = self.crop_mask(masks, bboxes)
|
||||
return np.greater(masks, 0.5)
|
||||
|
||||
|
|
@ -250,8 +254,9 @@ class YOLOv8Seg:
|
|||
if len(masks.shape) < 2:
|
||||
raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
|
||||
masks = masks[top:bottom, left:right]
|
||||
masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]),
|
||||
interpolation=cv2.INTER_LINEAR) # INTER_CUBIC would be better
|
||||
masks = cv2.resize(
|
||||
masks, (im0_shape[1], im0_shape[0]), interpolation=cv2.INTER_LINEAR
|
||||
) # INTER_CUBIC would be better
|
||||
if len(masks.shape) == 2:
|
||||
masks = masks[:, :, None]
|
||||
return masks
|
||||
|
|
@ -279,32 +284,46 @@ class YOLOv8Seg:
|
|||
cv2.fillPoly(im_canvas, np.int32([segment]), self.color_palette(int(cls_), bgr=True))
|
||||
|
||||
# draw bbox rectangle
|
||||
cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),
|
||||
self.color_palette(int(cls_), bgr=True), 1, cv2.LINE_AA)
|
||||
cv2.putText(im, f'{self.classes[cls_]}: {conf:.3f}', (int(box[0]), int(box[1] - 9)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.7, self.color_palette(int(cls_), bgr=True), 2, cv2.LINE_AA)
|
||||
cv2.rectangle(
|
||||
im,
|
||||
(int(box[0]), int(box[1])),
|
||||
(int(box[2]), int(box[3])),
|
||||
self.color_palette(int(cls_), bgr=True),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
cv2.putText(
|
||||
im,
|
||||
f"{self.classes[cls_]}: {conf:.3f}",
|
||||
(int(box[0]), int(box[1] - 9)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.7,
|
||||
self.color_palette(int(cls_), bgr=True),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Mix image
|
||||
im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)
|
||||
|
||||
# Show image
|
||||
if vis:
|
||||
cv2.imshow('demo', im)
|
||||
cv2.imshow("demo", im)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
# Save image
|
||||
if save:
|
||||
cv2.imwrite('demo.jpg', im)
|
||||
cv2.imwrite("demo.jpg", im)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
# Create an argument parser to handle command-line arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--model', type=str, required=True, help='Path to ONNX model')
|
||||
parser.add_argument('--source', type=str, default=str(ASSETS / 'bus.jpg'), help='Path to input image')
|
||||
parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
|
||||
parser.add_argument('--iou', type=float, default=0.45, help='NMS IoU threshold')
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to ONNX model")
|
||||
parser.add_argument("--source", type=str, default=str(ASSETS / "bus.jpg"), help="Path to input image")
|
||||
parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
|
||||
parser.add_argument("--iou", type=float, default=0.45, help="NMS IoU threshold")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Build model
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue