Fix Issue #345: KeypointDetectionModel returns correct prediction type

Author: kadhirash

roboflow/config.py

@@ -71,8 +71,11 @@ def get_conditional_configuration_variable(key, default):
 TYPE_OBJECT_DETECTION = "object-detection"
 TYPE_INSTANCE_SEGMENTATION = "instance-segmentation"
 TYPE_SEMANTIC_SEGMENTATION = "semantic-segmentation"
+TYPE_SEMANTIC_SEGMENTATION = "semantic-segmentation"
 TYPE_KEYPOINT_DETECTION = "keypoint-detection"
 
+KEYPOINT_DETECTION_MODEL = "KeypointDetectionModel"
+
 DEFAULT_BATCH_NAME = "Pip Package Upload"
 DEFAULT_JOB_NAME = "Annotated via API"
 

roboflow/models/keypoint_detection.py

@@ -8,7 +8,7 @@
 import requests
 from PIL import Image
 
-from roboflow.config import CLASSIFICATION_MODEL
+from roboflow.config import KEYPOINT_DETECTION_MODEL
 from roboflow.models.inference import InferenceModel
 from roboflow.util.image_utils import check_image_url
 from roboflow.util.prediction import PredictionGroup
@@ -119,7 +119,7 @@ def predict(self, image_path, hosted=False, confidence=None):  # type: ignore[ov
             resp.json(),
             image_dims=img_dims,
             image_path=image_path,
-            prediction_type=CLASSIFICATION_MODEL,
+            prediction_type=KEYPOINT_DETECTION_MODEL,
             colors=self.colors,
         )
 

roboflow/util/image_utils.py

@@ -5,8 +5,8 @@
 import urllib
 
 # Third-party imports
-import pi_heif  # type: ignore[import-untyped]
-import pillow_avif  # type: ignore[import-untyped]
+import pi_heif  # type: ignore[import-untyped, import-not-found]
+import pillow_avif  # type: ignore[import-untyped, import-not-found]
 import requests
 import yaml
 from PIL import Image

roboflow/util/prediction.py

@@ -10,6 +10,7 @@
 from roboflow.config import (
     CLASSIFICATION_MODEL,
     INSTANCE_SEGMENTATION_MODEL,
+    KEYPOINT_DETECTION_MODEL,
     OBJECT_DETECTION_MODEL,
     PREDICTION_OBJECT,
     SEMANTIC_SEGMENTATION_MODEL,
@@ -509,7 +510,7 @@ def create_prediction_group(json_response, image_path, prediction_type, image_di
         colors = {} if colors is None else colors
         prediction_list = []
 
-        if prediction_type in [OBJECT_DETECTION_MODEL, INSTANCE_SEGMENTATION_MODEL]:
+        if prediction_type in [OBJECT_DETECTION_MODEL, INSTANCE_SEGMENTATION_MODEL, KEYPOINT_DETECTION_MODEL]:
             for prediction in json_response["predictions"]:
                 prediction = Prediction(
                     prediction,