ultralytics 8.2.59 use Results.save_txt for validation (#14496)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-07-18 06:17:25 +08:00 committed by GitHub
parent ebf7dcf5a8
commit bfcd85323d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 83 additions and 32 deletions

View file

@ -48,9 +48,8 @@ class SegmentationValidator(DetectionValidator):
self.plot_masks = []
if self.args.save_json:
check_requirements("pycocotools>=2.0.6")
self.process = ops.process_mask_upsample # more accurate
else:
self.process = ops.process_mask # faster
# more accurate vs faster
self.process = ops.process_mask_upsample if self.args.save_json or self.args.save_txt else ops.process_mask
self.stats = dict(tp_m=[], tp=[], conf=[], pred_cls=[], target_cls=[], target_img=[])
def get_desc(self):
@ -148,14 +147,23 @@ class SegmentationValidator(DetectionValidator):
# Save
if self.args.save_json:
pred_masks = ops.scale_image(
pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(),
pbatch["ori_shape"],
ratio_pad=batch["ratio_pad"][si],
self.pred_to_json(
predn,
batch["im_file"][si],
ops.scale_image(
pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(),
pbatch["ori_shape"],
ratio_pad=batch["ratio_pad"][si],
),
)
if self.args.save_txt:
self.save_one_txt(
predn,
pred_masks,
self.args.save_conf,
pbatch["ori_shape"],
self.save_dir / "labels" / f'{Path(batch["im_file"][si]).stem}.txt',
)
self.pred_to_json(predn, batch["im_file"][si], pred_masks)
# if self.args.save_txt:
# save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
def finalize_metrics(self, *args, **kwargs):
"""Sets speed and confusion matrix for evaluation metrics."""
@ -235,6 +243,18 @@ class SegmentationValidator(DetectionValidator):
) # pred
self.plot_masks.clear()
def save_one_txt(self, predn, pred_masks, save_conf, shape, file):
"""Save YOLO detections to a txt file in normalized coordinates in a specific format."""
from ultralytics.engine.results import Results
Results(
np.zeros((shape[0], shape[1]), dtype=np.uint8),
path=None,
names=self.names,
boxes=predn[:, :6],
masks=pred_masks,
).save_txt(file, save_conf=save_conf)
def pred_to_json(self, predn, filename, pred_masks):
"""
Save one JSON result.