Upload folder using huggingface_hub

utils/inference.py

@@ -1,4 +1,5 @@
 import torch
+import argparse
 import numpy as np
 from PIL import Image
 from skimage import io
@@ -6,6 +7,31 @@ from ormbg import ORMBG
 import torch.nn.functional as F
 
 
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Remove background from images using ORMBG model."
+    )
+    parser.add_argument(
+        "--input",
+        type=str,
+        default="example.png",
+        help="Path to the input image file.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="no-background.png",
+        help="Path to the output image file.",
+    )
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default="models/ormbg.pth",
+        help="Path to the model file.",
+    )
+    return parser.parse_args()
+
+
 def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
     if len(im.shape) < 3:
         im = im[:, :, np.newaxis]
@@ -27,19 +53,19 @@ def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
     return im_array
 
 
-def example_inference():
-
-    image_path = "example.png"
-    result_name = "no-background.png"
+def inference(args):
+    image_path = args.input
+    result_name = args.output
+    model_path = args.model_path
 
     net = ORMBG()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     if torch.cuda.is_available():
-        net.load_state_dict(torch.load("models/ormbg.pth"))
+        net.load_state_dict(torch.load(model_path))
         net = net.cuda()
     else:
-        net.load_state_dict(torch.load("models/ormbg.pth", map_location="cpu"))
+        net.load_state_dict(torch.load(model_path, map_location="cpu"))
     net.eval()
 
     model_input_size = [1024, 1024]
@@ -61,4 +87,4 @@ def example_inference():
 
 
 if __name__ == "__main__":
-    example_inference()
+    inference(parse_args())