General refactoring and improvements (#373)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@ -143,7 +143,7 @@ class SegLoss(Loss):
|
||||
|
||||
@hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent), config_name=DEFAULT_CONFIG.name)
|
||||
def train(cfg):
|
||||
cfg.model = cfg.model or "yolov8n-seg.yaml"
|
||||
cfg.model = cfg.model or "yolov8n-seg.pt"
|
||||
cfg.data = cfg.data or "coco128-seg.yaml" # or yolo.ClassificationDataset("mnist")
|
||||
cfg.device = cfg.device if cfg.device is not None else ''
|
||||
# trainer = SegmentationTrainer(cfg)
|
||||
@ -154,11 +154,4 @@ def train(cfg):
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
CLI usage:
|
||||
python ultralytics/yolo/v8/segment/train.py model=yolov8n-seg.yaml data=coco128-segments epochs=100 imgsz=640
|
||||
|
||||
TODO:
|
||||
Direct cli support, i.e, yolov8 classify_train args.epochs 10
|
||||
"""
|
||||
train()
|
||||
|
@ -114,8 +114,9 @@ class SegmentationValidator(DetectionValidator):
|
||||
masks=True)
|
||||
if self.args.plots:
|
||||
self.confusion_matrix.process_batch(predn, labelsn)
|
||||
self.stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:,
|
||||
5], cls.squeeze(-1))) # conf, pcls, tcls
|
||||
|
||||
# Append correct_masks, correct_boxes, pconf, pcls, tcls
|
||||
self.stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:, 5], cls.squeeze(-1)))
|
||||
|
||||
pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8)
|
||||
if self.args.plots and self.batch_i < 3:
|
||||
|
Reference in New Issue
Block a user