Solved image_dims issue in object_detection.py

Was able to pass variables from object_detection.py instead of generating in prediction.py

Author: Tyler Odenthal

roboflow/models/object_detection.py

@@ -146,6 +146,7 @@ def predict(
         if not hosted:
             if type(image_path) is str:
                 image = Image.open(image_path).convert("RGB")
+                dimensions = image.size
                 # Create buffer
                 buffered = io.BytesIO()
                 image.save(buffered, format="PNG")
@@ -158,9 +159,11 @@ def predict(
                     data=img_str,
                     headers={"Content-Type": "application/x-www-form-urlencoded"},
                 )
+                image_dims = {"width": str(dimensions[0]), "height": str(dimensions[1])}
             elif isinstance(image_path, np.ndarray):
                 # Performing inference on a OpenCV2 frame
                 retval, buffer = cv2.imencode(".jpg", image_path)
+                dimensions = buffer.shape
                 img_str = base64.b64encode(buffer)
                 # print(img_str)
                 img_str = img_str.decode("ascii")
@@ -169,11 +172,13 @@ def predict(
                     data=img_str,
                     headers={"Content-Type": "application/x-www-form-urlencoded"},
                 )
+                image_dims = {"width": str(dimensions[0]), "height": str(dimensions[1])}
             else:
                 raise ValueError("image_path must be a string or a numpy array.")
         else:
             # Create API URL for hosted image (slightly different)
             self.api_url += "&image=" + urllib.parse.quote_plus(image_path)
+            image_dims = {"width": "0", "height": "0"}
             # POST to the API
             resp = requests.get(self.api_url)
 
@@ -184,6 +189,7 @@ def predict(
                 resp.json(),
                 image_path=image_path,
                 prediction_type=OBJECT_DETECTION_MODEL,
+                image_dims=image_dims,
             )
         # Returns base64 encoded Data
         elif self.format == "image":

roboflow/util/prediction.py

@@ -473,53 +473,34 @@ def json(self):
         return prediction_group_json
 
     @staticmethod
-    def create_prediction_group(json_response, image_path, prediction_type):
+    def create_prediction_group(json_response, image_path, prediction_type, image_dims):
         """
         Method to create a prediction group based on the JSON Response
 
         :param prediction_type:
         :param json_response: Based on Roboflow JSON Response from Inference API
         :param model:
         :param image_path:
+        :param image_dims:
         :return:
         """
         prediction_list = []
 
-        # LOADING PHOTO TAKES TIME - PASS VARIABLES IF POSSIBLE
-        image_loaded = Image.open(image_path)
-        dimensions = image_loaded.size
-        # LOADING PHOTO TAKES TIME - PASS VARIABLES IF POSSIBLE
-
         if prediction_type in [OBJECT_DETECTION_MODEL, INSTANCE_SEGMENTATION_MODEL]:
             for prediction in json_response["predictions"]:
                 prediction = Prediction(
                     prediction, image_path, prediction_type=prediction_type
                 )
                 prediction_list.append(prediction)
-                if "image" not in json_response:
-                    json_response["image"] = {
-                        "width": dimensions[0],
-                        "height": dimensions[1],
-                    }
-            img_dims = json_response["image"]
+            img_dims = image_dims
         elif prediction_type == CLASSIFICATION_MODEL:
             prediction = Prediction(json_response, image_path, prediction_type)
             prediction_list.append(prediction)
-            if "image" not in json_response:
-                json_response["image"] = {
-                    "width": dimensions[0],
-                    "height": dimensions[1],
-                }
-            img_dims = json_response["image"]
+            img_dims = image_dims
         elif prediction_type == SEMANTIC_SEGMENTATION_MODEL:
             prediction = Prediction(json_response, image_path, prediction_type)
             prediction_list.append(prediction)
-            if "image" not in json_response:
-                json_response["image"] = {
-                    "width": dimensions[0],
-                    "height": dimensions[1],
-                }
-            img_dims = json_response["image"]
+            img_dims = image_dims
 
         # Seperate list and return as a prediction group
         return PredictionGroup(img_dims, image_path, *prediction_list)