Flatten the structure of the model directory

software-mansion-labs · May 23, 2024 · d3ee058 · d3ee058
1 parent 7c04591
commit d3ee058
Show file tree

Hide file tree

Showing 18 changed files with 37 additions and 17 deletions.
diff --git a/.gitattributes b/.gitattributes
@@ -1 +1,5 @@
 models/**/*.onnx filter=lfs diff=lfs merge=lfs -text
+models/deeplab_v3_mobilenetv3_segmentation.onnx filter=lfs diff=lfs merge=lfs -text
+models/fasterrcnn_resnet50_fpn_detector.onnx filter=lfs diff=lfs merge=lfs -text
+models/mobilenetv3small-classifier.onnx filter=lfs diff=lfs merge=lfs -text
+models/ssdlite320_mobilenetv3_detector.onnx filter=lfs diff=lfs merge=lfs -text
diff --git a/lib/ex_vision/classification/mobilenet_v3_small.ex b/lib/ex_vision/classification/mobilenet_v3_small.ex
@@ -5,8 +5,8 @@ defmodule ExVision.Classification.MobileNetV3Small do
   Weights from Imagenet 1k.
   """
   use ExVision.Model.Definition.Ortex,
-    model: "classification/mobilenetv3small/model.onnx",
-    categories: "classification/mobilenetv3small/categories.json"
+    model: "mobilenetv3small-classifier.onnx",
+    categories: "imagenet_v2_categories.json"
 
   require Bunch.Typespec
   alias ExVision.Utils

diff --git a/lib/ex_vision/detection/fasterrcnn_resnet50_fpn.ex b/lib/ex_vision/detection/fasterrcnn_resnet50_fpn.ex
@@ -3,8 +3,8 @@ defmodule ExVision.Detection.FasterRCNN_ResNet50_FPN do
   FasterRCNN object detector with ResNet50 backbone and FPN detection head, exported from torchvision.
   """
   use ExVision.Model.Definition.Ortex,
-    model: "detection/fasterrcnn_resnet50_fpn/model.onnx",
-    categories: "detection/fasterrcnn_resnet50_fpn/categories.json"
+    model: "fasterrcnn_resnet50_fpn_detector.onnx",
+    categories: "coco_categories.json"
 
   use ExVision.Detection.GenericDetector
 

diff --git a/lib/ex_vision/detection/generic_detector.ex b/lib/ex_vision/detection/generic_detector.ex
@@ -15,8 +15,7 @@ defmodule ExVision.Detection.GenericDetector do
     ExVision.Utils.resize(img, {224, 224})
   end
 
-  @spec postprocessing({Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()}, ImageMetadata.t(), [atom()]) ::
-          output_t()
+  @spec postprocessing(map(), ImageMetadata.t(), [atom()]) :: output_t()
   def postprocessing(
         %{"boxes" => bboxes, "scores" => scores, "labels" => labels},
         metadata,

diff --git a/lib/ex_vision/detection/ssdlite320_mobilenetv3.ex b/lib/ex_vision/detection/ssdlite320_mobilenetv3.ex
@@ -3,8 +3,8 @@ defmodule ExVision.Detection.Ssdlite320_MobileNetv3 do
   SSDLite320 object detector with MobileNetV3 Large architecture, exported from torchvision.
   """
   use ExVision.Model.Definition.Ortex,
-    model: "detection/ssdlite320_mobilenetv3/model.onnx",
-    categories: "detection/ssdlite320_mobilenetv3/categories.json"
+    model: "ssdlite320_mobilenetv3_detector.onnx",
+    categories: "coco_categories.json"
 
   use ExVision.Detection.GenericDetector
 

diff --git a/lib/ex_vision/model/definition/ortex.ex b/lib/ex_vision/model/definition/ortex.ex
@@ -21,7 +21,7 @@ defmodule ExVision.Model.Definition.Ortex do
 
   In this callback, you should transform the output to match your desired format.
   """
-  @callback postprocessing(tuple(), ImageMetadata.t()) :: ExVision.Model.output_t()
+  @callback postprocessing(map(), ImageMetadata.t()) :: ExVision.Model.output_t()
 
   @typedoc """
   A type describing ONNX provider that can be used with ExVision.

diff --git a/lib/ex_vision/segmentation/deep_lab_v3_mobilenet_v3.ex b/lib/ex_vision/segmentation/deep_lab_v3_mobilenet_v3.ex
@@ -3,8 +3,8 @@ defmodule ExVision.Segmentation.DeepLabV3_MobileNetV3 do
   A semantic segmentation model for MobileNetV3 Backbone. Exported from torchvision.
   """
   use ExVision.Model.Definition.Ortex,
-    model: "segmentation/deeplab_v3_mobilenetv3/model.onnx",
-    categories: "segmentation/deeplab_v3_mobilenetv3/categories.json"
+    model: "deeplab_v3_mobilenetv3_segmentation.onnx",
+    categories: "coco_with_voc_labels_categories.json"
 
   @type output_t() :: %{category_t() => Nx.Tensor.t()}
 

diff --git a/...n/fasterrcnn_resnet50_fpn/categories.json → models/coco_categories.json b/...n/fasterrcnn_resnet50_fpn/categories.json → models/coco_categories.json
diff --git a/models/coco_with_voc_labels_categories.json b/models/coco_with_voc_labels_categories.json
@@ -0,0 +1,23 @@
+[
+  "__background__",
+  "aeroplane",
+  "bicycle",
+  "bird",
+  "boat",
+  "bottle",
+  "bus",
+  "car",
+  "cat",
+  "chair",
+  "cow",
+  "diningtable",
+  "dog",
+  "horse",
+  "motorbike",
+  "person",
+  "pottedplant",
+  "sheep",
+  "sofa",
+  "train",
+  "tvmonitor"
+]
diff --git a/models/segmentation/deeplab_v3/model.onnx → .../deeplab_v3_mobilenetv3_segmentation.onnx b/models/segmentation/deeplab_v3/model.onnx → .../deeplab_v3_mobilenetv3_segmentation.onnx
diff --git a/models/detection/ssdlite320_mobilenetv3/categories.json b/models/detection/ssdlite320_mobilenetv3/categories.json
diff --git a/...ection/fasterrcnn_resnet50_fpn/model.onnx → models/fasterrcnn_resnet50_fpn_detector.onnx b/...ection/fasterrcnn_resnet50_fpn/model.onnx → models/fasterrcnn_resnet50_fpn_detector.onnx
diff --git a/...fication/mobilenetv3small/categories.json → models/imagenet_v2_categories.json b/...fication/mobilenetv3small/categories.json → models/imagenet_v2_categories.json
diff --git a/...lassification/mobilenetv3small/model.onnx → models/mobilenetv3small-classifier.onnx b/...lassification/mobilenetv3small/model.onnx → models/mobilenetv3small-classifier.onnx
diff --git a/models/segmentation/deeplab_v3/categories.json b/models/segmentation/deeplab_v3/categories.json
diff --git a/models/segmentation/deeplab_v3_mobilenetv3/categories.json b/models/segmentation/deeplab_v3_mobilenetv3/categories.json
diff --git a/models/segmentation/deeplab_v3_mobilenetv3/model.onnx b/models/segmentation/deeplab_v3_mobilenetv3/model.onnx
diff --git a/...tection/ssdlite320_mobilenetv3/model.onnx → models/ssdlite320_mobilenetv3_detector.onnx b/...tection/ssdlite320_mobilenetv3/model.onnx → models/ssdlite320_mobilenetv3_detector.onnx