ultralytics 8.0.41 TF SavedModel and EdgeTPU export (#1034)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Noobtoss <96134731+Noobtoss@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>

ultralytics/yolo/v8/classify/predict.py

@@ -22,7 +22,7 @@ class ClassificationPredictor(BasePredictor):
         results = []
         for i, pred in enumerate(preds):
             orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img
-            results.append(Results(probs=pred.softmax(0), orig_img=orig_img, names=self.model.names))
+            results.append(Results(probs=pred, orig_img=orig_img, names=self.model.names))
 
         return results
 

ultralytics/yolo/v8/classify/val.py

@@ -30,6 +30,9 @@ class ClassificationValidator(BaseValidator):
         self.pred.append(preds.argsort(1, descending=True)[:, :5])
         self.targets.append(batch['cls'])
 
+    def finalize_metrics(self, *args, **kwargs):
+        self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
+
     def get_stats(self):
         self.metrics.process(self.targets, self.pred)
         return self.metrics.results_dict

ultralytics/yolo/v8/detect/val.py

@@ -111,6 +111,9 @@ class DetectionValidator(BaseValidator):
             # 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):
+        self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
+
     def get_stats(self):
         stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)]  # to numpy
         if len(stats) and stats[0].any():

ultralytics/yolo/v8/segment/val.py

@@ -1,6 +1,5 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
-import os
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
 
@@ -10,7 +9,7 @@ import torch.nn.functional as F
 
 from ultralytics.yolo.utils import DEFAULT_CFG, NUM_THREADS, ops
 from ultralytics.yolo.utils.checks import check_requirements
-from ultralytics.yolo.utils.metrics import ConfusionMatrix, SegmentMetrics, box_iou, mask_iou
+from ultralytics.yolo.utils.metrics import SegmentMetrics, box_iou, mask_iou
 from ultralytics.yolo.utils.plotting import output_to_target, plot_images
 from ultralytics.yolo.v8.detect import DetectionValidator
 
@@ -120,6 +119,9 @@ class SegmentationValidator(DetectionValidator):
             # 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):
+        self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
+
     def _process_batch(self, detections, labels, pred_masks=None, gt_masks=None, overlap=False, masks=False):
         """
         Return correct prediction matrix