HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.50 AMP check and YOLOv5u YAMLs (#1263)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Troy <wudashuo@vip.qq.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Huijae Lee <46982469+ZeroAct@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-03-06 11:39:26 +01:00 • committed by GitHub
parent 3861e6c82a
commit f0d8e4718b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
29 changed files with 440 additions and 83 deletions
Download patch file Download diff file Expand all files Collapse all files

10
ultralytics/yolo/data/base.py
View file

@ -35,7 +35,8 @@ class BaseDataset(Dataset):
batch_size=None,
stride=32,
pad=0.5,
single_cls=False):
single_cls=False,
classes=None):
super().__init__()
self.img_path = img_path
self.imgsz = imgsz
@ -45,8 +46,7 @@ class BaseDataset(Dataset):
self.im_files = self.get_img_files(self.img_path)
self.labels = self.get_labels()
if self.single_cls:
self.update_labels(include_class=[])
self.update_labels(include_class=classes) # single_cls and include_class
self.ni = len(self.labels)
@ -96,7 +96,7 @@ class BaseDataset(Dataset):
"""include_class, filter labels to include only these classes (optional)"""
include_class_array = np.array(include_class).reshape(1, -1)
for i in range(len(self.labels)):
if include_class:
if include_class is not None:
cls = self.labels[i]['cls']
bboxes = self.labels[i]['bboxes']
segments = self.labels[i]['segments']
@ -104,7 +104,7 @@ class BaseDataset(Dataset):
self.labels[i]['cls'] = cls[j]
self.labels[i]['bboxes'] = bboxes[j]
if segments:
self.labels[i]['segments'] = segments[j]
self.labels[i]['segments'] = [segments[si] for si, idx in enumerate(j) if idx]
if self.single_cls:
self.labels[i]['cls'][:, 0] = 0