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:
co-authored by
Claude Opus 5
parent
8fa7a9d142
commit
dce2eef0a7
@@ -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
|
||||
Reference in New Issue
Block a user