fix type checking for object detection utils

Signed-off-by: Farhan Ahmed <[email protected]>

Author: f4str

art/estimators/object_detection/utils.py

@@ -97,7 +97,7 @@ def convert_pt_to_tf(y: List[Dict[str, np.ndarray]], height: int, width: int) ->
 def cast_inputs_to_pt(
     x: Union[np.ndarray, "torch.Tensor"],
     y: Optional[List[Dict[str, Union[np.ndarray, "torch.Tensor"]]]] = None,
-) -> Tuple["torch.Tensor", List[Dict[str, "torch.Tensor"]]]:
+) -> Tuple["torch.Tensor", Optional[List[Dict[str, "torch.Tensor"]]]]:
     """
     Cast object detection inputs `(x, y)` to PyTorch tensors.
 
@@ -117,25 +117,43 @@ def cast_inputs_to_pt(
     else:
         x_tensor = x
 
+    y_tensor: Optional[List[Dict[str, torch.Tensor]]] = None
+
     # Convert labels into tensor
-    if y is not None and isinstance(y, list) and isinstance(y[0]["boxes"], np.ndarray):
+    if isinstance(y, list):
         y_tensor = []
         for y_i in y:
-            y_t = {
-                "boxes": torch.from_numpy(y_i["boxes"]).to(dtype=torch.float32),
-                "labels": torch.from_numpy(y_i["labels"]).to(dtype=torch.int64),
-            }
+            y_t = {}
+
+            if isinstance(y_i["boxes"], np.ndarray):
+                y_t["boxes"] = torch.from_numpy(y_i["boxes"]).to(dtype=torch.float32)
+            else:
+                y_t["boxes"] = y_i["boxes"]
+
+            if isinstance(y_i["labels"], np.ndarray):
+                y_t["labels"] = torch.from_numpy(y_i["labels"]).to(dtype=torch.int64)
+            else:
+                y_t["labels"] = y_i["labels"]
+
             if "masks" in y_i:
-                y_t["masks"] = torch.from_numpy(y_i["masks"]).to(dtype=torch.uint8)
+                if isinstance(y_i["masks"], np.ndarray):
+                    y_t["masks"] = torch.from_numpy(y_i["masks"]).to(dtype=torch.uint8)
+                else:
+                    y_t["masks"] = y_i["masks"]
+
             y_tensor.append(y_t)
-    elif y is not None and isinstance(y, dict):
+    elif isinstance(y, dict):
         y_tensor = []
-        for i in range(y["boxes"].shape[0]):
+        for i in range(len(y["boxes"])):
             y_t = {}
+
             y_t["boxes"] = y["boxes"][i]
             y_t["labels"] = y["labels"][i]
+            if "masks" in y:
+                y_t["masks"] = y["masks"][i]
+
             y_tensor.append(y_t)
     else:
-        y_tensor = y  # type: ignore
+        y_tensor = y
 
     return x_tensor, y_tensor