Merge pull request #46 from roboflow-ai/inferenceEnhancements

Inference Improvements

roboflow/util/prediction.py

@@ -218,16 +218,18 @@ def __getitem__(self, key):
 
 
 class PredictionGroup:
-    def __init__(self, *args):
+    def __init__(self, image_dims, image_path, *args):
         """
         :param args: The prediction(s) to be added to the prediction group
         """
         # List of predictions (core of the PredictionGroup)
         self.predictions = []
         # Base image path (path of image of first prediction in prediction group)
-        self.base_image_path = ""
+        self.base_image_path = image_path
         # Base prediction type (prediction type of image of first prediction in prediction group)
         self.base_prediction_type = ""
+
+        self.image_dims = image_dims
         # Iterate through the arguments
         for index, prediction in enumerate(args):
             # Set base image path based on first prediction
@@ -434,6 +436,7 @@ def json(self):
         for prediction in self.predictions:
             prediction_group_json["predictions"].append(prediction.json())
 
+        prediction_group_json["image"] = self.image_dims
         return prediction_group_json
 
     @staticmethod
@@ -459,12 +462,14 @@ def create_prediction_group(json_response, image_path, prediction_type):
                 )
                 # Add to prediction list
                 prediction_list.append(prediction)
+            img_dims = json_response["image"]
         # For classification model
         elif prediction_type == CLASSIFICATION_MODEL:
             # Create prediction for predicted class
             prediction = Prediction(json_response, image_path, prediction_type)
             # add to prediction list
             prediction_list.append(prediction)
 
+            img_dims = {}
         # Seperate list and return as a prediction group
-        return PredictionGroup(*prediction_list)
+        return PredictionGroup(img_dims, image_path, *prediction_list)

setup.py

@@ -6,7 +6,7 @@
 
 setuptools.setup(
     name="roboflow",  # Replace with your own username
-    version="0.2.14",
+    version="0.2.15",
     author="Roboflow",
     author_email="[email protected]",
     description="python client for the Roboflow application",