ultralytics 8.0.84 results JSON outputs (#2171)

Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Talia Bender <85292283+taliabender@users.noreply.github.com>

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = '8.0.83'
+__version__ = '8.0.84'
 
 from ultralytics.hub import start
 from ultralytics.yolo.engine.model import YOLO

ultralytics/yolo/data/base.py

@@ -120,11 +120,14 @@ class BaseDataset(Dataset):
                 cls = self.labels[i]['cls']
                 bboxes = self.labels[i]['bboxes']
                 segments = self.labels[i]['segments']
+                keypoints = self.labels[i]['keypoints']
                 j = (cls == include_class_array).any(1)
                 self.labels[i]['cls'] = cls[j]
                 self.labels[i]['bboxes'] = bboxes[j]
                 if segments:
                     self.labels[i]['segments'] = [segments[si] for si, idx in enumerate(j) if idx]
+                if keypoints is not None:
+                    self.labels[i]['keypoints'] = keypoints[j]
             if self.single_cls:
                 self.labels[i]['cls'][:, 0] = 0
 

ultralytics/yolo/engine/results.py

@@ -97,11 +97,6 @@ class Results(SimpleClass):
         self.path = path
         self._keys = ('boxes', 'masks', 'probs', 'keypoints')
 
-    def pandas(self):
-        """Convert the results to a pandas DataFrame."""
-        pass
-        # TODO masks.pandas + boxes.pandas + cls.pandas
-
     def __getitem__(self, idx):
         """Return a Results object for the specified index."""
         r = self.new()
@@ -315,6 +310,35 @@ class Results(SimpleClass):
                          file=save_dir / self.names[int(d.cls)] / f'{file_name.stem}.jpg',
                          BGR=True)
 
+    def pandas(self):
+        """Convert the object to a pandas DataFrame (not yet implemented)."""
+        LOGGER.warning("WARNING ⚠️ 'Results.pandas' method is not yet implemented.")
+
+    def tojson(self, normalize=False):
+        """Convert the object to JSON format."""
+        import json
+
+        # Create list of detection dictionaries
+        results = []
+        data = self.boxes.data.cpu().tolist()
+        h, w = self.orig_shape if normalize else (1, 1)
+        for i, row in enumerate(data):
+            box = {'x1': row[0] / w, 'y1': row[1] / h, 'x2': row[2] / w, 'y2': row[3] / h}
+            conf = row[4]
+            id = int(row[5])
+            name = self.names[id]
+            result = {'name': name, 'class': id, 'confidence': conf, 'box': box}
+            if self.masks:
+                x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1]  # numpy array
+                result['segments'] = {'x': (x / w).tolist(), 'y': (y / h).tolist()}
+            if self.keypoints is not None:
+                x, y, visible = self.keypoints[i].cpu().unbind(dim=1)  # torch Tensor
+                result['keypoints'] = {'x': (x / w).tolist(), 'y': (y / h).tolist(), 'visible': visible.tolist()}
+            results.append(result)
+
+        # Convert detections to JSON
+        return json.dumps(results, indent=2)
+
 
 class Boxes(BaseTensor):
     """
@@ -397,10 +421,6 @@ class Boxes(BaseTensor):
         """Return the boxes in xywh format normalized by original image size."""
         return self.xywh / self.orig_shape[[1, 0, 1, 0]]
 
-    def pandas(self):
-        """Convert the object to a pandas DataFrame (not yet implemented)."""
-        LOGGER.info('results.pandas() method not yet implemented')
-
     @property
     def boxes(self):
         """Return the raw bboxes tensor (deprecated)."""
@@ -466,3 +486,7 @@ class Masks(BaseTensor):
         """Return the raw masks tensor (deprecated)."""
         LOGGER.warning("WARNING ⚠️ 'Masks.masks' is deprecated. Use 'Masks.data' instead.")
         return self.data
+
+    def pandas(self):
+        """Convert the object to a pandas DataFrame (not yet implemented)."""
+        LOGGER.warning("WARNING ⚠️ 'Masks.pandas' method is not yet implemented.")

ultralytics/yolo/utils/instance.py

@@ -34,7 +34,7 @@ class Bboxes:
     """Now only numpy is supported."""
 
     def __init__(self, bboxes, format='xyxy') -> None:
-        assert format in _formats
+        assert format in _formats, f'Invalid bounding box format: {format}, format must be one of {_formats}'
         bboxes = bboxes[None, :] if bboxes.ndim == 1 else bboxes
         assert bboxes.ndim == 2
         assert bboxes.shape[1] == 4
@@ -66,7 +66,7 @@ class Bboxes:
 
     def convert(self, format):
         """Converts bounding box format from one type to another."""
-        assert format in _formats
+        assert format in _formats, f'Invalid bounding box format: {format}, format must be one of {_formats}'
         if self.format == format:
             return
         elif self.format == 'xyxy':