Use global model

bhky · Dec 1, 2023 · 6db3c35 · 6db3c35
1 parent 2776f5f
commit 6db3c35
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/opennsfw2/_inference.py b/opennsfw2/_inference.py
@@ -16,11 +16,12 @@
 from ._model import make_open_nsfw_model
 from ._typing import NDFloat32Array
 
+model: Optional[Model] = None
+
 
 def predict_image(
         image_path: str,
         preprocessing: Preprocessing = Preprocessing.YAHOO,
-        model: Optional[Model] = None,
         weights_path: Optional[str] = get_default_weights_path(),
         grad_cam_path: Optional[str] = None,
         alpha: float = 0.8
@@ -31,6 +32,7 @@ def predict_image(
     """
     pil_image = Image.open(image_path)
     image = preprocess_image(pil_image, preprocessing)
+    global model
     if model is None:
         model = make_open_nsfw_model(weights_path=weights_path)
     nsfw_probability = float(model(np.expand_dims(image, 0))[0][1])
@@ -50,7 +52,6 @@ def predict_images(
         image_paths: Sequence[str],
         batch_size: int = 8,
         preprocessing: Preprocessing = Preprocessing.YAHOO,
-        model: Optional[Model] = None,
         weights_path: Optional[str] = get_default_weights_path(),
         grad_cam_paths: Optional[Sequence[str]] = None,
         alpha: float = 0.8
@@ -63,6 +64,7 @@ def predict_images(
         preprocess_image(Image.open(image_path), preprocessing)
         for image_path in image_paths
     ])
+    global model
     if model is None:
         model = make_open_nsfw_model(weights_path=weights_path)
     predictions = model.predict(images, batch_size=batch_size, verbose=0)
@@ -112,7 +114,6 @@ def predict_video_frames(
         batch_size: int = 8,
         output_video_path: Optional[str] = None,
         preprocessing: Preprocessing = Preprocessing.YAHOO,
-        model: Optional[Model] = None,
         weights_path: Optional[str] = get_default_weights_path(),
         progress_bar: bool = True
 ) -> Tuple[List[float], List[float]]:
@@ -122,6 +123,7 @@ def predict_video_frames(
     cap = cv2.VideoCapture(video_path)  # pylint: disable=no-member
     fps = cap.get(cv2.CAP_PROP_FPS)  # pylint: disable=no-member
 
+    global model
     if model is None:
         model = make_open_nsfw_model(weights_path=weights_path)