add segmentation benchmark

Author: Zhang-Yang-Sustech

benchmark/config/image_segmentation_efficientsam.yaml

@@ -0,0 +1,16 @@
+Benchmark:
+  name: "Image Segmentation Benchmark"
+  type: "Segmentation"
+  data:
+    path: "data/image_segmentation"
+    files: ["messi5.jpg", "100040721_1.jpg"]
+    sizes: # [[w1, h1], ...], Omit to run at original scale
+      - [640, 640]
+  metric:
+    warmup: 30
+    repeat: 10
+  backend: "default"
+  target: "cpu"
+
+Model:
+  name: "EfficientSAM"

benchmark/utils/dataloaders/__init__.py

@@ -2,5 +2,6 @@
 from .classification import ClassificationImageLoader
 from .recognition import RecognitionImageLoader
 from .tracking import TrackingVideoLoader
+from .segmentation import SegmentationImageLoader
 
-__all__ = ['BaseImageLoader', 'BaseVideoLoader', 'ClassificationImageLoader', 'RecognitionImageLoader', 'TrackingVideoLoader']
+__all__ = ['BaseImageLoader', 'BaseVideoLoader', 'ClassificationImageLoader', 'RecognitionImageLoader', 'SegmentationImageLoader', 'TrackingVideoLoader']

benchmark/utils/dataloaders/segmentation.py

@@ -0,0 +1,66 @@
+import os
+
+import numpy as np
+import cv2 as cv
+
+from .base_dataloader import _BaseImageLoader
+from ..factory import DATALOADERS
+
+@DATALOADERS.register
+class SegmentationImageLoader(_BaseImageLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self._to_rgb = kwargs.pop('toRGB', False)
+        self._point_label= self._load_point_and_label()
+
+    def _load_point_and_label(self):
+        points_labels = dict.fromkeys(self._files, None)
+        for filename in self._files:
+            if os.path.exists(os.path.join(self._path, '{}.txt'.format(filename[:-4]))):
+                points_labels[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), ndmin=2)
+            else:
+                points_labels[filename] = None
+        # for filename in self._files:
+        #     label_file = os.path.join(self._path, '{}.txt'.format(filename[:-4]))
+        #     if os.path.exists(label_file):
+        #         # 假设标签文件的每一行格式为：x y label
+        #         # 其中 x, y 是点的坐标，label 是标签（0 或 1）
+        #         with open(label_file, 'r') as file:
+        #             lines = file.readlines()
+        #             current_point_label = []
+        #             for line in lines:
+        #                 parts = line.strip().split()
+        #                 if len(parts) == 3:
+        #                     x, y, label = map(int, parts)
+        #                     current_point_label.append((x, y, label))
+        #         points_labels[filename] = current_point_label
+        #     else:
+        #         points_labels[filename] = None
+        return points_labels
+
+
+    def _toRGB(self, image):
+        return cv.cvtColor(image, cv.COLOR_BGR2RGB)
+
+    def __iter__(self):
+        for filename in self._files:
+            image = cv.imread(os.path.join(self._path, filename))
+            
+            if self._to_rgb:
+                image = self._toRGB(image)
+                
+            if [0, 0] in self._sizes:
+                point_and_label = self._point_label.get(filename)
+                if point_and_label is not None:
+                    yield filename, image, point_and_label
+                else:
+                    yield filename, image, None
+            else:
+                for size in self._sizes:
+                    image_r = cv.resize(image, size)
+                    point_and_label = self._point_label.get(filename)
+                    if point_and_label is not None:
+                        yield filename, image_r, point_and_label
+                    else:
+                        yield filename, image_r, None

benchmark/utils/metrics/__init__.py

@@ -2,5 +2,6 @@
 from .detection import Detection
 from .recognition import Recognition
 from .tracking import Tracking
+from .segmentation import Segmentation
 
-__all__ = ['Base', 'Detection', 'Recognition', 'Tracking']
+__all__ = ['Base', 'Detection', 'Recognition', 'Segmentation', 'Tracking']

benchmark/utils/metrics/segmentation.py

@@ -0,0 +1,35 @@
+import cv2 as cv
+
+from .base_metric import BaseMetric
+from ..factory import METRICS
+
+@METRICS.register
+class Segmentation(BaseMetric):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def forward(self, model, *args, **kwargs):
+        img, point_and_label = args
+        size = [img.shape[1], img.shape[0]]
+        self._timer.reset()
+        if point_and_label is not None:
+            for idx, pl in enumerate(point_and_label):
+                point = [[pl[0], pl[1]]]
+                label = [[pl[2]]]
+                for _ in range(self._warmup):
+                    model.infer(img, point, label)
+                for _ in range(self._repeat):
+                    self._timer.start()
+                    model.infer(img, point, label)
+                    self._timer.stop()
+        else:
+            point = [[int(size[0]/2), int(size[1]/2)]]
+            label = [[1]]
+            for _ in range(self._warmup):
+                model.infer(img, point, label)
+            for _ in range(self._repeat):
+                self._timer.start()
+                model.infer(img, point, label)
+                self._timer.stop()
+
+        return self._timer.getRecords()