Merge pull request #93 from roboflow/json-response-trt-fix

Fix for json_response["image"] in prediction.py

Author: Tyler-Odenthal

roboflow/core/project.py

@@ -534,10 +534,6 @@ def single_upload(
 
     def __str__(self):
         # String representation of project
-        json_str = {
-            "name": self.name,
-            "type": self.type,
-            "workspace": self.__workspace,
-        }
+        json_str = {"name": self.name, "type": self.type, "workspace": self.__workspace}
 
         return json.dumps(json_str, indent=2)

roboflow/core/version.py

@@ -98,15 +98,9 @@ def __init__(
                     local=local,
                 )
             elif self.type == TYPE_INSTANCE_SEGMENTATION:
-                self.model = InstanceSegmentationModel(
-                    self.__api_key,
-                    self.id,
-                )
+                self.model = InstanceSegmentationModel(self.__api_key, self.id)
             elif self.type == TYPE_SEMANTIC_SEGMENTATION:
-                self.model = SemanticSegmentationModel(
-                    self.__api_key,
-                    self.id,
-                )
+                self.model = SemanticSegmentationModel(self.__api_key, self.id)
             else:
                 self.model = None
 
@@ -485,10 +479,7 @@ def __get_format_identifier(self, format):
                 "You must pass a format argument to version.download() or define a model in your Roboflow object"
             )
 
-        friendly_formats = {
-            "yolov5": "yolov5pytorch",
-            "yolov7": "yolov7pytorch",
-        }
+        friendly_formats = {"yolov5": "yolov5pytorch", "yolov7": "yolov7pytorch"}
         return friendly_formats.get(format, format)
 
     def __reformat_yaml(self, location, format):

roboflow/models/inference.py

@@ -31,24 +31,25 @@ def __get_image_params(self, image_path):
         """
         validate_image_path(image_path)
 
-        hosted_image = urllib.parse.urlparse(image_path).scheme in (
-            "http",
-            "https",
-        )
+        hosted_image = urllib.parse.urlparse(image_path).scheme in ("http", "https")
 
         if hosted_image:
-            return {"image": image_path}, {}
+            image_dims = {"width": "Undefined", "height": "Undefined"}
+            return {"image": image_path}, {}, image_dims
 
         image = Image.open(image_path)
+        dimensions = image.size
+        image_dims = {"width": str(dimensions[0]), "height": str(dimensions[1])}
         buffered = io.BytesIO()
         image.save(buffered, quality=90, format="JPEG")
         data = MultipartEncoder(
             fields={"file": ("imageToUpload", buffered.getvalue(), "image/jpeg")}
         )
-        return {}, {
-            "data": data,
-            "headers": {"Content-Type": data.content_type},
-        }
+        return (
+            {},
+            {"data": data, "headers": {"Content-Type": data.content_type}},
+            image_dims,
+        )
 
     def predict(self, image_path, prediction_type=None, **kwargs):
         """
@@ -60,7 +61,7 @@ def predict(self, image_path, prediction_type=None, **kwargs):
         :return: PredictionGroup - a group of predictions based on Roboflow JSON response
         :raises Exception: Image path is not valid
         """
-        params, request_kwargs = self.__get_image_params(image_path)
+        params, request_kwargs, image_dims = self.__get_image_params(image_path)
 
         params["api_key"] = self.__api_key
 
@@ -74,4 +75,5 @@ def predict(self, image_path, prediction_type=None, **kwargs):
             response.json(),
             image_path=image_path,
             prediction_type=prediction_type,
+            image_dims=image_dims,
         )

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,22 +159,27 @@ 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)
+                # Currently cv2.imencode does not properly return shape
+                dimensions = buffer.shape
                 img_str = base64.b64encode(buffer)
-                # print(img_str)
                 img_str = img_str.decode("ascii")
                 resp = requests.post(
                     self.api_url,
                     data=img_str,
                     headers={"Content-Type": "application/x-www-form-urlencoded"},
                 )
+                # Replace with dimensions variable once cv2.imencode shape solution is found
+                image_dims = {"width": "0", "height": "0"}
             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 +190,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/image_utils.py

@@ -22,10 +22,7 @@ def check_image_url(url):
     :param url: URL of image
     :returns: Boolean
     """
-    if urllib.parse.urlparse(url).scheme not in (
-        "http",
-        "https",
-    ):
+    if urllib.parse.urlparse(url).scheme not in ("http", "https"):
         return False
 
     r = requests.head(url)

roboflow/util/prediction.py

@@ -80,10 +80,7 @@ def plot_annotation(axes, prediction=None, stroke=1, transparency=60):
     elif prediction["prediction_type"] == INSTANCE_SEGMENTATION_MODEL:
         points = [[p["x"], p["y"]] for p in prediction["points"]]
         polygon = patches.Polygon(
-            points,
-            linewidth=stroke,
-            edgecolor="r",
-            facecolor="none",
+            points, linewidth=stroke, edgecolor="r", facecolor="none"
         )
         axes.add_patch(polygon)
     elif prediction["prediction_type"] == SEMANTIC_SEGMENTATION_MODEL:
@@ -476,14 +473,15 @@ 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 = []
@@ -494,15 +492,15 @@ def create_prediction_group(json_response, image_path, prediction_type):
                     prediction, image_path, prediction_type=prediction_type
                 )
                 prediction_list.append(prediction)
-            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)
-            img_dims = {}
+            img_dims = image_dims
         elif prediction_type == SEMANTIC_SEGMENTATION_MODEL:
             prediction = Prediction(json_response, image_path, prediction_type)
             prediction_list.append(prediction)
-            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)