xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ultralytics 8.0.42 DDP fix and Docs updates (#1065)

Browse Source 
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: Noobtoss <96134731+Noobtoss@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Glenn Jocher
2023-02-20 12:56:20 +01:00
committed by GitHub
parent f6e393c1d2
commit f2a7a29e53
33 changed files with 196 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
ultralytics/yolo/engine/exporter.py
View File

@ -75,7 +75,6 @@ from ultralytics.yolo.utils.files import file_size
from ultralytics.yolo.utils.ops import Profile
from ultralytics.yolo.utils.torch_utils import get_latest_opset, select_device, smart_inference_mode
CUDA = torch.cuda.is_available()
ARM64 = platform.machine() in ('arm64', 'aarch64')
@ -324,7 +323,7 @@ class Exporter:
# Simplify
if self.args.simplify:
try:
check_requirements(('onnxsim', 'onnxruntime-gpu' if CUDA else 'onnxruntime'))
check_requirements(('onnxsim', 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'))
import onnxsim
LOGGER.info(f'{prefix} simplifying with onnxsim {onnxsim.__version__}...')
@ -506,10 +505,12 @@ class Exporter:
try:
import tensorflow as tf # noqa
except ImportError:
check_requirements(f"tensorflow{'-macos' if MACOS else '-aarch64' if ARM64 else '' if CUDA else '-cpu'}")
check_requirements(
f"tensorflow{'-macos' if MACOS else '-aarch64' if ARM64 else '' if torch.cuda.is_available() else '-cpu'}"
)
import tensorflow as tf # noqa
check_requirements(('onnx', 'onnx2tf', 'sng4onnx', 'onnxsim', 'onnx_graphsurgeon', 'tflite_support',
'onnxruntime-gpu' if CUDA else 'onnxruntime'),
'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'),
cmds='--extra-index-url https://pypi.ngc.nvidia.com')
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')

7
ultralytics/yolo/engine/model.py
View File

@ -32,7 +32,7 @@ class YOLO:
YOLO (You Only Look Once) object detection model.
Args:
model (str or Path): Path to the model file to load or create.
model (str, Path): Path to the model file to load or create.
type (str): Type/version of models to use. Defaults to "v8".
Attributes:
@ -62,7 +62,7 @@ class YOLO:
predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model.
Returns:
List[ultralytics.yolo.engine.results.Results]: The prediction results.
list(ultralytics.yolo.engine.results.Results): The prediction results.
"""
def __init__(self, model='yolov8n.pt', type='v8') -> None:
@ -114,6 +114,7 @@ class YOLO:
self.task = guess_model_task(cfg_dict)
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
self.model = self.ModelClass(cfg_dict, verbose=verbose and RANK == -1) # initialize
self.overrides['model'] = self.cfg
def _load(self, weights: str):
"""
@ -204,7 +205,7 @@ class YOLO:
def track(self, source=None, stream=False, **kwargs):
from ultralytics.tracker.track import register_tracker
register_tracker(self)
# bytetrack-based method needs low confidence predictions as input
# ByteTrack-based method needs low confidence predictions as input
conf = kwargs.get('conf') or 0.1
kwargs['conf'] = conf
kwargs['mode'] = 'track'

1
ultralytics/yolo/engine/predictor.py
View File

@ -92,6 +92,7 @@ class BasePredictor:
self.annotator = None
self.data_path = None
self.source_type = None
self.batch = None
self.callbacks = defaultdict(list, callbacks.default_callbacks) # add callbacks
callbacks.add_integration_callbacks(self)

42
ultralytics/yolo/engine/results.py
View File

@ -28,13 +28,14 @@ class Results:
"""
def __init__(self, boxes=None, masks=None, probs=None, orig_img=None, names=None) -> None:
def __init__(self, orig_img, path, names, boxes=None, masks=None, probs=None) -> None:
self.orig_img = orig_img
self.orig_shape = orig_img.shape[:2]
self.boxes = Boxes(boxes, self.orig_shape) if boxes is not None else None # native size boxes
self.masks = Masks(masks, self.orig_shape) if masks is not None else None # native size or imgsz masks
self.probs = probs if probs is not None else None
self.names = names
self.path = path
self.comp = ['boxes', 'masks', 'probs']
def pandas(self):
@ -42,7 +43,7 @@ class Results:
# TODO masks.pandas + boxes.pandas + cls.pandas
def __getitem__(self, idx):
r = Results(orig_img=self.orig_img)
r = Results(orig_img=self.orig_img, path=self.path, names=self.names)
for item in self.comp:
if getattr(self, item) is None:
continue
@ -58,7 +59,7 @@ class Results:
self.probs = probs
def cpu(self):
r = Results(orig_img=self.orig_img)
r = Results(orig_img=self.orig_img, path=self.path, names=self.names)
for item in self.comp:
if getattr(self, item) is None:
continue
@ -66,7 +67,7 @@ class Results:
return r
def numpy(self):
r = Results(orig_img=self.orig_img)
r = Results(orig_img=self.orig_img, path=self.path, names=self.names)
for item in self.comp:
if getattr(self, item) is None:
continue
@ -74,7 +75,7 @@ class Results:
return r
def cuda(self):
r = Results(orig_img=self.orig_img)
r = Results(orig_img=self.orig_img, path=self.path, names=self.names)
for item in self.comp:
if getattr(self, item) is None:
continue
@ -82,7 +83,7 @@ class Results:
return r
def to(self, *args, **kwargs):
r = Results(orig_img=self.orig_img)
r = Results(orig_img=self.orig_img, path=self.path, names=self.names)
for item in self.comp:
if getattr(self, item) is None:
continue
@ -123,7 +124,7 @@ class Results:
orig_shape (tuple, optional): Original image size.
""")
def visualize(self, show_conf=True, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
def plot(self, show_conf=True, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
"""
Plots the given result on an input RGB image. Accepts cv2(numpy) or PIL Image
@ -146,9 +147,9 @@ class Results:
annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True))
if masks is not None:
im_gpu = torch.as_tensor(img, dtype=torch.float16).permute(2, 0, 1).flip(0).contiguous()
im_gpu = F.resize(im_gpu, masks.data.shape[1:]) / 255
annotator.masks(masks.data, colors=[colors(x, True) for x in boxes.cls], im_gpu=im_gpu)
im = torch.as_tensor(img, dtype=torch.float16, device=masks.data.device).permute(2, 0, 1).flip(0)
im = F.resize(im.contiguous(), masks.data.shape[1:]) / 255
annotator.masks(masks.data, colors=[colors(x, True) for x in boxes.cls], im_gpu=im)
if logits is not None:
top5i = logits.argsort(0, descending=True)[:5].tolist() # top 5 indices
@ -371,24 +372,3 @@ class Masks:
Properties:
segments (list): A list of segments which includes x,y,w,h,label,confidence, and mask of each detection masks.
""")
if __name__ == '__main__':
# test examples
results = Results(boxes=torch.randn((2, 6)), masks=torch.randn((2, 160, 160)), orig_shape=[640, 640])
results = results.cuda()
print('--cuda--pass--')
results = results.cpu()
print('--cpu--pass--')
results = results.to('cuda:0')
print('--to-cuda--pass--')
results = results.to('cpu')
print('--to-cpu--pass--')
results = results.numpy()
print('--numpy--pass--')
# box = Boxes(boxes=torch.randn((2, 6)), orig_shape=[5, 5])
# box = box.cuda()
# box = box.cpu()
# box = box.numpy()
# for b in box:
# print(b)