Improved method

README.md

@@ -165,6 +165,12 @@ def prediction(model, processor, img_batch: List[Image.Image], class_to_predict:
     """
     Predict if images meet or exceed a specific NSFW threshold
     """
+    if class_to_predict not in ["low", "medium", "high"]:
+        raise ValueError("class_to_predict must be one of: low, medium, high")
+    
+    if not 0 <= threshold <= 1:
+        raise ValueError("threshold must be between 0 and 1")
+
     output = predict_batch_values(model, processor, img_batch)
     return [output[i][class_to_predict] >= threshold for i in range(len(output))]