onnx script fix and added as main package

metaseg/sahi_predict.py

@@ -4,7 +4,7 @@
 import torch
 from PIL import Image
 
-from metaseg import SamPredictor, sam_model_registry
+from metaseg.generator import SamPredictor, sam_model_registry
 from metaseg.utils import (
  download_model,
  load_image,

metaseg/utils/__init__.py

@@ -8,6 +8,7 @@
 from .data_utils import load_image as load_image
 from .data_utils import load_mask as load_mask
 from .data_utils import load_server_image as load_server_image
+from .data_utils import load_video as load_video
 from .data_utils import multi_boxes as multi_boxes
 from .data_utils import plt_load_box as plt_load_box
 from .data_utils import plt_load_mask as plt_load_mask

pyproject.toml

@@ -34,6 +34,7 @@ pillow = "^9.5.0"
 pycocotools = "^2.0.6"
 onnx = "^1.14.0"
 onnxruntime = "^1.15.1"
+fal-serverless = "^0.6.35"
 
 
 

scripts/amg.py

@@ -9,7 +9,7 @@
 import os
 from typing import Any, Dict, List
 
-import cv2 # type: ignore
+import cv2
 
 from metaseg.generator import SamAutomaticMaskGenerator, sam_model_registry
 

scripts/export_onnx_model.py

@@ -8,17 +8,13 @@
 import warnings
 
 import torch
+from onnxruntime import InferenceSession
+from onnxruntime.quantization import QuantType
+from onnxruntime.quantization.quantize import quantize_dynamic
 
-from metaseg import build_sam, build_sam_vit_b, build_sam_vit_l
+from metaseg.generator import build_sam, build_sam_vit_b, build_sam_vit_l
 from metaseg.utils.onnx import SamOnnxModel
 
-try:
- import onnxruntime # type: ignore
-
- onnxruntime_exists = True
-except ImportError:
- onnxruntime_exists = False
-
 parser = argparse.ArgumentParser(
  description="Export the SAM prompt encoder and mask decoder to an ONNX model."
 )
@@ -169,11 +165,10 @@ def run_export(
  dynamic_axes=dynamic_axes,
  )
 
- if onnxruntime_exists:
- ort_inputs = {k: to_numpy(v) for k, v in dummy_inputs.items()}
- ort_session = onnxruntime.InferenceSession(output)
- _ = ort_session.run(None, ort_inputs)
- print("Model has successfully been run with ONNXRuntime.")
+ ort_inputs = {k: to_numpy(v) for k, v in dummy_inputs.items()}
+ ort_session = InferenceSession(output)
+ _ = ort_session.run(None, ort_inputs)
+ print("Model has successfully been run with ONNXRuntime.")
 
 
 def to_numpy(tensor):
@@ -193,18 +188,13 @@ def to_numpy(tensor):
  return_extra_metrics=args.return_extra_metrics,
  )
 
- if args.quantize_out is not None:
- assert onnxruntime_exists, "onnxruntime is required to quantize the model."
- from onnxruntime.quantization import QuantType # type: ignore
- from onnxruntime.quantization.quantize import quantize_dynamic # type: ignore
-
- print(f"Quantizing model and writing to {args.quantize_out}...")
- quantize_dynamic(
- model_input=args.output,
- model_output=args.quantize_out,
- optimize_model=True,
- per_channel=False,
- reduce_range=False,
- weight_type=QuantType.QUInt8,
- )
- print("Done!")
+ print(f"Quantizing model and writing to {args.quantize_out}...")
+ quantize_dynamic(
+ model_input=args.output,
+ model_output=args.quantize_out,
+ optimize_model=True,
+ per_channel=False,
+ reduce_range=False,
+ weight_type=QuantType.QUInt8,
+ )
+ print("Done!")