fix(ml): batch size setting (#26524)

This commit is contained in:
Mert
2026-03-05 12:01:47 -05:00
committed by GitHub
parent 09fabb36b6
commit 35a521c6ec
6 changed files with 89 additions and 8 deletions
@@ -22,7 +22,7 @@ class TextDetector(InferenceModel):
depends = []
identity = (ModelType.DETECTION, ModelTask.OCR)
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
def __init__(self, model_name: str, min_score: float = 0.5, **model_kwargs: Any) -> None:
super().__init__(model_name.split("__")[-1], **model_kwargs, model_format=ModelFormat.ONNX)
self.max_resolution = 736
self.mean = np.array([0.5, 0.5, 0.5], dtype=np.float32)
@@ -33,7 +33,7 @@ class TextDetector(InferenceModel):
}
self.postprocess = DBPostProcess(
thresh=0.3,
box_thresh=model_kwargs.get("minScore", 0.5),
box_thresh=model_kwargs.get("minScore", min_score),
max_candidates=1000,
unclip_ratio=1.6,
use_dilation=True,