From bc7fd12d8998fcf926c6560d126c303c0416db64 Mon Sep 17 00:00:00 2001 From: Iuri de Silvio Date: Mon, 19 Oct 2026 10:08:49 +0200 Subject: [PATCH] Fix the crashes in Prediction.save Prediction.save() raised before writing the file for most predictions: - Object detection and keypoint: the label background got float corners (x - width / 2, ...), which cv2.rectangle rejects on every supported OpenCV. The corners are now ints, and the label text sits on that background at the top of the box, where PredictionGroup.save puts it. - Classification: "Confidence: " + confidence failed on the float the API returns, so the confidence goes through str() first. - Hosted images: only http:// URLs were downloaded, so an https:// image went to cv2.imread and came back empty. Both are downloaded now. - Semantic segmentation: mask_image decoded the mask with np.fromstring, whose binary mode NumPy 2.3 removed. It uses np.frombuffer now, which fixes PredictionGroup.save for semantic segmentation too. Co-Authored-By: Claude Opus 5.5 --- roboflow/util/image_utils.py | 2 +- roboflow/util/prediction.py | 12 ++--- tests/util/test_prediction.py | 89 +++++++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 tests/util/test_prediction.py diff --git a/roboflow/util/image_utils.py b/roboflow/util/image_utils.py index 6ed7b539..9efb2204 100644 --- a/roboflow/util/image_utils.py +++ b/roboflow/util/image_utils.py @@ -53,7 +53,7 @@ def mask_image(image, encoded_mask, transparency=60): import cv2 import numpy as np - np_data = np.fromstring(base64.b64decode(encoded_mask), np.uint8) # type: ignore[no-overload] + np_data = np.frombuffer(base64.b64decode(encoded_mask), np.uint8) mask = cv2.imdecode(np_data, cv2.IMREAD_UNCHANGED) # Fallback in case the API returns an incorrectly sized mask diff --git a/roboflow/util/prediction.py b/roboflow/util/prediction.py index 77d4cd73..347db804 100644 --- a/roboflow/util/prediction.py +++ b/roboflow/util/prediction.py @@ -126,7 +126,7 @@ def __load_image(self): import cv2 import numpy as np - if "http://" in self.image_path: + if "http://" in self.image_path or "https://" in self.image_path: req = urllib.request.urlopen(self.image_path) arr = np.asarray(bytearray(req.read()), dtype=np.uint8) image = cv2.imdecode(arr, -1) # 'Load it as it is' @@ -182,10 +182,10 @@ def save(self, output_path="predictions.jpg", stroke=2, transparency=60): # Draw background rectangle for text cv2.rectangle( image, - (x - width / 2, y - height / 2 + 1), + (int(x - width / 2), int(y - height / 2 + 1)), ( - x - width / 2 + text_size[0] + 1, - y - height / 2 + int(1.5 * text_size[1]), + int(x - width / 2 + text_size[0] + 1), + int(y - height / 2 + int(1.5 * text_size[1])), ), stroke_color, -1, @@ -194,7 +194,7 @@ def save(self, output_path="predictions.jpg", stroke=2, transparency=60): cv2.putText( image, class_name, - (int(x - width / 2), y + text_size[1]), + (int(x - width / 2), int(y - height / 2 + text_size[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), @@ -224,7 +224,7 @@ def save(self, output_path="predictions.jpg", stroke=2, transparency=60): # Add text and relax cv2.putText( image, - (self["top"] + " | " + "Confidence: " + self["confidence"]), + (self["top"] + " | " + "Confidence: " + str(self["confidence"])), (int(width / 2), 5), cv2.FONT_HERSHEY_DUPLEX, 0.5, diff --git a/tests/util/test_prediction.py b/tests/util/test_prediction.py new file mode 100644 index 00000000..d6a8d74a --- /dev/null +++ b/tests/util/test_prediction.py @@ -0,0 +1,89 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from roboflow.config import ( + CLASSIFICATION_MODEL, + KEYPOINT_DETECTION_MODEL, + OBJECT_DETECTION_MODEL, + SEMANTIC_SEGMENTATION_MODEL, +) +from roboflow.util.prediction import Prediction + +IMAGE_PATH = "tests/images/rabbit.JPG" +BOX = {"x": 300.5, "y": 200.5, "width": 120.0, "height": 90.0, "class": "rabbit", "confidence": 0.91} + + +class TestPredictionSave(unittest.TestCase): + def _save(self, prediction): + """Save the prediction as a PNG, so pixels aren't altered by compression, and read it back.""" + import cv2 + + with tempfile.TemporaryDirectory() as tmp: + output_path = os.path.join(tmp, "prediction.png") + prediction.save(output_path) + image = cv2.imread(output_path) + self.assertIsNotNone(image) + return image + + def test_object_detection_draws_the_label_on_its_background(self): + import cv2 + import numpy as np + + image = self._save(Prediction(dict(BOX), IMAGE_PATH, prediction_type=OBJECT_DETECTION_MODEL)) + + # The label background is filled with the stroke color, so white pixels in it can only be the label text. + text_width, text_height = cv2.getTextSize("rabbit", cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)[0] + left = int(BOX["x"] - BOX["width"] / 2) + top = int(BOX["y"] - BOX["height"] / 2) + label = image[top + 1 : top + int(1.5 * text_height), left : left + text_width + 1] + self.assertTrue(np.all(label[0, -1] == (255, 0, 0))) + self.assertTrue(np.any(np.all(label == 255, axis=-1))) + + def test_keypoint_detection(self): + image = self._save(Prediction(dict(BOX, keypoints=[]), IMAGE_PATH, prediction_type=KEYPOINT_DETECTION_MODEL)) + self.assertEqual(image.shape, (992, 1024, 3)) + + def test_classification(self): + prediction = Prediction( + {"top": "rabbit", "confidence": 0.91, "predictions": [{"class": "rabbit", "confidence": 0.91}]}, + IMAGE_PATH, + prediction_type=CLASSIFICATION_MODEL, + ) + image = self._save(prediction) + # The label goes on a border around the image. + self.assertEqual(image.shape, (992 + 200, 1024 + 200, 3)) + + def test_semantic_segmentation(self): + import base64 + + import cv2 + import numpy as np + + mask = np.zeros((992, 1024), np.uint8) + mask[100:300, 100:300] = 255 + encoded_mask = base64.b64encode(cv2.imencode(".png", mask)[1].tobytes()).decode() + prediction = Prediction( + {"segmentation_mask": encoded_mask, "class_map": {"1": "rabbit"}}, + IMAGE_PATH, + prediction_type=SEMANTIC_SEGMENTATION_MODEL, + ) + + image = self._save(prediction) + + self.assertEqual(image.shape, (992, 1024, 3)) + + @mock.patch("urllib.request.urlopen") + def test_hosted_https_image(self, urlopen): + urlopen.return_value.read.return_value = Path(IMAGE_PATH).read_bytes() + + image = self._save(Prediction(dict(BOX), "https://example.com/rabbit.jpg")) + + urlopen.assert_called_once() + self.assertEqual(image.shape, (992, 1024, 3)) + + +if __name__ == "__main__": + unittest.main()