fix(ml-backend): 推理分批,避免 Retrieve Predictions 撑爆显存;新增批量推送脚本

Retrieve Predictions 会把选中的任务一次发给后端,而后端把整个列表交给
model.predict(source=paths)——ultralytics 不理会 batch 参数,会当成一个 batch,
于是显存爆掉。这与之前 visualize_detections 踩的是同一个坑。后端现在按
YOLO_BATCH(默认 16)自己分批。

另加 scripts/push_predictions.py:重训完在本地分批推理,再用批量导入接口
POST /api/projects/<id>/import/predictions 提交,不经过后端服务,显存与请求体都可控。
支持 --replace 先删同版本旧预测(预测是追加的,重复跑会堆两份)、--limit 与 --dry-run。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-22 12:11:06 +08:00
co-authored by Claude Opus 5
parent 8fa7a9d142
commit dce2eef0a7
4 changed files with 374 additions and 3 deletions
+88
View File
@@ -0,0 +1,88 @@
"""批量给 Label Studio 推预测:URL 还原与分页拉取。"""
import sys
import pytest
pytest.importorskip("requests")
sys.path.insert(0, "scripts")
import push_predictions as pp # noqa: E402
def test_local_path_from_local_files_url(tmp_path):
url = "/data/local-files/?d=boxes_shap_label_studio/images/blackdot_0489.png"
assert pp.local_path(url, tmp_path) == (
tmp_path / "boxes_shap_label_studio/images/blackdot_0489.png"
)
def test_local_path_unquotes_escaped_names(tmp_path):
"""文件名里有中文或空格时 URL 是转义过的,得还原回来。"""
url = "/data/local-files/?d=proj/images/%E5%8F%8D%E9%9D%A2%20a.png"
assert pp.local_path(url, tmp_path) == tmp_path / "proj/images/反面 a.png"
def test_local_path_rejects_non_local_urls(tmp_path):
with pytest.raises(ValueError, match="本地存储"):
pp.local_path("https://example.com/a.png", tmp_path)
class _FakeResponse:
def __init__(self, payload, status=200):
self._payload = payload
self.status_code = status
def json(self):
return self._payload
def raise_for_status(self):
if self.status_code >= 400:
raise AssertionError(f"HTTP {self.status_code}")
class _FakeSession:
"""按页返回任务,记录请求参数。"""
def __init__(self, pages):
self.pages = pages
self.headers = {}
self.calls = []
def get(self, url, params=None, timeout=None):
self.calls.append(params)
page = params["page"]
return _FakeResponse(self.pages[page - 1] if page <= len(self.pages) else {"tasks": []})
@pytest.fixture
def client(monkeypatch):
monkeypatch.setattr(pp, "LabelStudio", pp.LabelStudio)
instance = pp.LabelStudio.__new__(pp.LabelStudio)
instance.url = "http://ls"
return instance
def test_tasks_pagination_stops_at_total(client):
"""分页要拉完再停,不能只拿第一页。"""
client.session = _FakeSession([
{"tasks": [{"id": 1}, {"id": 2}], "total": 3},
{"tasks": [{"id": 3}], "total": 3},
])
assert [t["id"] for t in client.tasks(39, page_size=2)] == [1, 2, 3]
def test_tasks_handles_drf_style_pagination(client):
"""有的版本返回 results/count 而不是 tasks/total。"""
client.session = _FakeSession([
{"results": [{"id": 1}], "count": 2},
{"results": [{"id": 2}], "count": 2},
])
assert [t["id"] for t in client.tasks(39, page_size=1)] == [1, 2]
def test_tasks_stops_on_empty_page(client):
"""没有 total 字段时,靠空页收尾,不能无限翻。"""
client.session = _FakeSession([{"tasks": [{"id": 1}]}, {"tasks": []}])
assert [t["id"] for t in client.tasks(39)] == [1]
assert len(client.session.calls) == 2