Move processing bounding boxes to utils

lib/ex_vision/instance_segmentation/maskrcnn_resnet50_fpn_v2.ex

@@ -8,6 +8,8 @@ defmodule ExVision.InstanceSegmentation.MaskRCNN_ResNet50_FPN_V2 do
 
   require Logger
 
+  import ExVision.Utils
+
   alias ExVision.Types.BBoxWithMask
 
   @type output_t() :: [BBoxWithMask.t()]
@@ -46,16 +48,10 @@ defmodule ExVision.InstanceSegmentation.MaskRCNN_ResNet50_FPN_V2 do
     scale_x = w / 224
     scale_y = h / 224
 
-    bboxes =
-      bboxes
-      |> Nx.squeeze(axes: [0])
-      |> Nx.multiply(Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
-      |> Nx.round()
-      |> Nx.as_type(:s64)
-      |> Nx.to_list()
+    bboxes = process_bbox(bboxes, Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
 
-    scores = scores |> Nx.squeeze(axes: [0]) |> Nx.to_list()
-    labels = labels |> Nx.squeeze(axes: [0]) |> Nx.to_list()
+    scores = unbatch(scores)
+    labels = unbatch(labels)
 
     masks =
       masks

lib/ex_vision/keypoint_detection/keypointrcnn_resnet50_fpn.ex

@@ -8,6 +8,8 @@ defmodule ExVision.KeypointDetection.KeypointRCNN_ResNet50_FPN do
 
   require Logger
 
+  import ExVision.Utils
+
   alias ExVision.Types.BBoxWithKeypoints
 
   @typep output_t() :: [BBoxWithKeypoints.t()]
@@ -67,26 +69,14 @@ defmodule ExVision.KeypointDetection.KeypointRCNN_ResNet50_FPN do
     scale_x = w / 224
     scale_y = h / 224
 
-    bboxes =
-      bboxes
-      |> Nx.squeeze(axes: [0])
-      |> Nx.multiply(Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
-      |> Nx.round()
-      |> Nx.as_type(:s64)
-      |> Nx.to_list()
-
-    scores = scores |> Nx.squeeze(axes: [0]) |> Nx.to_list()
-    labels = labels |> Nx.squeeze(axes: [0]) |> Nx.to_list()
-
-    keypoints_list =
-      keypoints_list
-      |> Nx.squeeze(axes: [0])
-      |> Nx.multiply(Nx.tensor([scale_x, scale_y, 1]))
-      |> Nx.round()
-      |> Nx.as_type(:s64)
-      |> Nx.to_list()
-
-    keypoints_scores_list = keypoints_scores_list |> Nx.squeeze(axes: [0]) |> Nx.to_list()
+    bboxes = process_bbox(bboxes, Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
+
+    scores = unbatch(scores)
+    labels = unbatch(labels)
+
+    keypoints_list = process_bbox(keypoints_list, Nx.tensor([scale_x, scale_y, 1]))
+
+    keypoints_scores_list = unbatch(keypoints_scores_list)
 
     [bboxes, scores, labels, keypoints_list, keypoints_scores_list]
     |> Enum.zip()

lib/ex_vision/object_detection/generic_detector.ex

@@ -6,6 +6,8 @@ defmodule ExVision.ObjectDetection.GenericDetector do
 
   require Logger
 
+  import ExVision.Utils
+
   alias ExVision.Types.{BBox, ImageMetadata}
 
   @typep output_t() :: [BBox.t()]
@@ -29,16 +31,10 @@ defmodule ExVision.ObjectDetection.GenericDetector do
     scale_x = w / 224
     scale_y = h / 224
 
-    bboxes =
-      bboxes
-      |> Nx.squeeze(axes: [0])
-      |> Nx.multiply(Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
-      |> Nx.round()
-      |> Nx.as_type(:s64)
-      |> Nx.to_list()
+    bboxes = process_bbox(bboxes, Nx.tensor([scale_x, scale_y, scale_x, scale_y]))
 
-    scores = scores |> Nx.squeeze(axes: [0]) |> Nx.to_list()
-    labels = labels |> Nx.squeeze(axes: [0]) |> Nx.to_list()
+    scores = unbatch(scores)
+    labels = unbatch(labels)
 
     [bboxes, scores, labels]
     |> Enum.zip()

lib/ex_vision/utils.ex

@@ -155,4 +155,17 @@ defmodule ExVision.Utils do
   def batched_run(process_name, input) do
     process_name |> batched_run([input]) |> hd()
   end
+
+  defp process_bbox(bbox, scales, axes \\ [0]) do
+    bbox
+    |> Nx.squeeze(axes: axes)
+    |> Nx.multiply(scales)
+    |> Nx.round()
+    |> Nx.as_type(:s64)
+    |> Nx.to_list()
+  end
+
+  defp unbatch(batched_value, axes \\ [0]) do
+    batched_value |> Nx.squeeze(axes: axes) |> Nx.to_list()
+  end
 end