Compare commits
11 commits
7cd7b01f76
...
537ef2e489
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
537ef2e489 | ||
| 27d012256e | |||
| 7e71102f86 | |||
| 48efb1f4e7 | |||
|
|
b3308e6e64 | ||
| 8115224846 | |||
| a71f6dc731 | |||
| c2de3e306f | |||
| 640344b0f3 | |||
| b9427ca574 | |||
| 32441f36fb |
29 changed files with 855 additions and 60 deletions
4
.agents/PROJECT_CONTEXT.md
Normal file → Executable file
4
.agents/PROJECT_CONTEXT.md
Normal file → Executable file
|
|
@ -352,8 +352,8 @@ auto_augment=randaugment
|
|||
|
||||
Вероятности и доли валидируются в диапазоне `0…1`; `degrees`, `shear` и
|
||||
`close_mosaic` не могут быть отрицательными. Режимы copy-paste: `flip`, `mixup`.
|
||||
Политики AutoAugment: `randaugment`, `autoaugment`, `augmix`. Если augmentation
|
||||
выключена, эти kwargs вообще не передаются в Ultralytics.
|
||||
Политики AutoAugment: `none` (отключено), `randaugment`, `autoaugment`, `augmix`. Если augmentation
|
||||
выключена, эти kwargs вообще не передаются в Ultralytics. При `auto_augment="none"` параметр транслируется в `None` для Ultralytics.
|
||||
|
||||
## 10. Работа с датасетами
|
||||
|
||||
|
|
|
|||
0
.agents/PROJECT_ISSUES.md
Normal file → Executable file
0
.agents/PROJECT_ISSUES.md
Normal file → Executable file
0
.gitignore
vendored
Normal file → Executable file
0
.gitignore
vendored
Normal file → Executable file
35
Dockerfile
Normal file → Executable file
35
Dockerfile
Normal file → Executable file
|
|
@ -1,37 +1,22 @@
|
|||
FROM python:3.11-slim
|
||||
FROM ultralytics/ultralytics:latest
|
||||
|
||||
# Install system dependencies needed for OpenCV, PyTorch, and Ultralytics
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libgomp1 \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
# The ultralytics image already has python, system dependencies, PyTorch, TensorRT, and ultralytics installed.
|
||||
# We just need to install uv and our project dependencies.
|
||||
|
||||
# Pin the installer as well as application dependencies.
|
||||
# Pin the installer
|
||||
RUN pip install --no-cache-dir uv==0.10.6
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /workspace
|
||||
|
||||
ENV UV_COMPILE_BYTECODE=1 \
|
||||
UV_LINK_MODE=copy \
|
||||
UV_PROJECT_ENVIRONMENT=/opt/venv \
|
||||
ULTRALYTICS_SAFE_LOAD=1
|
||||
# Install our dependencies into the system python without resolving/syncing
|
||||
# which would remove packages installed by ultralytics base image (like tensorrt).
|
||||
COPY pyproject.toml README.md ./
|
||||
RUN uv pip install --system fastapi mlflow uvicorn websockets
|
||||
|
||||
# Install the exact dependency set recorded in uv.lock. Keeping the project out of
|
||||
# this layer allows dependency caching while source files change.
|
||||
COPY pyproject.toml uv.lock README.md ./
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --locked --no-dev --no-install-project
|
||||
|
||||
# Copy source code and install the project without re-resolving dependencies.
|
||||
# Copy source code and install the project
|
||||
COPY src ./src
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --locked --no-dev
|
||||
|
||||
ENV PATH="/opt/venv/bin:$PATH"
|
||||
RUN uv pip install --system -e .
|
||||
|
||||
# Expose Web UI port and MLflow port
|
||||
EXPOSE 8000
|
||||
|
|
|
|||
0
README.md
Normal file → Executable file
0
README.md
Normal file → Executable file
17
docker-compose.yml
Normal file → Executable file
17
docker-compose.yml
Normal file → Executable file
|
|
@ -5,18 +5,19 @@ services:
|
|||
image: yolo-train-webui:latest
|
||||
ports:
|
||||
# The training API has no built-in user accounts, so expose it locally only.
|
||||
- "127.0.0.1:8000:8000"
|
||||
- "0.0.0.0:8000:8000"
|
||||
volumes:
|
||||
- ./datasets:/workspace/datasets
|
||||
- ./runs:/workspace/runs
|
||||
- ./models:/workspace/models
|
||||
- ./models/.config:/root/.config/Ultralytics
|
||||
# Uncomment the block below on Linux with NVIDIA GPU to pass the graphics card into the container:
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
ipc: host
|
||||
restart: unless-stopped
|
||||
|
|
|
|||
2
pyproject.toml
Normal file → Executable file
2
pyproject.toml
Normal file → Executable file
|
|
@ -3,7 +3,7 @@ name = "yolo-train-webui"
|
|||
version = "0.1.0"
|
||||
description = "Web UI for training Ultralytics YOLO models with MLflow tracking"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"fastapi>=0.110.0",
|
||||
"mlflow>=3.0",
|
||||
|
|
|
|||
0
scripts/create_yolo26_smoke_datasets.py
Normal file → Executable file
0
scripts/create_yolo26_smoke_datasets.py
Normal file → Executable file
0
scripts/run_yolo26_smoke_training.py
Normal file → Executable file
0
scripts/run_yolo26_smoke_training.py
Normal file → Executable file
0
scripts/verify_mlflow_smoke.py
Normal file → Executable file
0
scripts/verify_mlflow_smoke.py
Normal file → Executable file
0
src/yolo_webui/__init__.py
Normal file → Executable file
0
src/yolo_webui/__init__.py
Normal file → Executable file
0
src/yolo_webui/__main__.py
Normal file → Executable file
0
src/yolo_webui/__main__.py
Normal file → Executable file
293
src/yolo_webui/app.py
Normal file → Executable file
293
src/yolo_webui/app.py
Normal file → Executable file
|
|
@ -410,20 +410,40 @@ async def list_datasets():
|
|||
|
||||
@app.get("/api/models")
|
||||
async def list_models():
|
||||
models_dir = Path("models")
|
||||
if not models_dir.exists():
|
||||
return []
|
||||
|
||||
items = []
|
||||
try:
|
||||
for path in models_dir.iterdir():
|
||||
if path.is_file() and path.suffix.lower() in (".pt", ".pth", ".yaml", ".yml"):
|
||||
|
||||
# Check models directory
|
||||
models_dir = Path("models")
|
||||
if models_dir.exists():
|
||||
try:
|
||||
for path in models_dir.iterdir():
|
||||
if path.is_file() and path.suffix.lower() in (".pt", ".pth", ".yaml", ".yml"):
|
||||
items.append({
|
||||
"name": path.name,
|
||||
"path": str(path.absolute()),
|
||||
"source": "models"
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list models in models/: {e}")
|
||||
|
||||
# Check runs directory for .pt and .pth files
|
||||
runs_dir = Path("runs")
|
||||
if runs_dir.exists():
|
||||
try:
|
||||
for path in runs_dir.rglob("*.pt"):
|
||||
items.append({
|
||||
"name": path.name,
|
||||
"name": f"{path.parent.parent.parent.name}/{path.parent.parent.name}/{path.name}",
|
||||
"path": str(path.absolute()),
|
||||
"source": "runs"
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list models: {e}")
|
||||
for path in runs_dir.rglob("*.pth"):
|
||||
items.append({
|
||||
"name": f"{path.parent.parent.parent.name}/{path.parent.parent.name}/{path.name}",
|
||||
"path": str(path.absolute()),
|
||||
"source": "runs"
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list models in runs/: {e}")
|
||||
|
||||
return sorted(items, key=lambda x: x["name"])
|
||||
|
||||
|
|
@ -508,6 +528,259 @@ async def stop_training():
|
|||
return {"message": "Запрос на остановку отправлен."}
|
||||
|
||||
|
||||
class ExportManager:
|
||||
def __init__(self) -> None:
|
||||
self.state = LiveState()
|
||||
self.active_websockets: set[WebSocket] = set()
|
||||
self._lock = threading.Lock()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||||
self._broadcast_lock: asyncio.Lock | None = None
|
||||
self._process: subprocess.Popen | None = None
|
||||
|
||||
def add_websocket(self, websocket: WebSocket) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
with self._lock:
|
||||
self.active_websockets.add(websocket)
|
||||
if self._event_loop is not loop:
|
||||
self._event_loop = loop
|
||||
self._broadcast_lock = asyncio.Lock()
|
||||
|
||||
def remove_websocket(self, websocket: WebSocket) -> None:
|
||||
with self._lock:
|
||||
self.active_websockets.discard(websocket)
|
||||
|
||||
def broadcast(self, data: dict[str, Any]) -> None:
|
||||
payload = json.dumps(data)
|
||||
with self._lock:
|
||||
loop = self._event_loop
|
||||
has_sockets = bool(self.active_websockets)
|
||||
|
||||
if not has_sockets:
|
||||
return
|
||||
if loop is None or loop.is_closed():
|
||||
return
|
||||
|
||||
coroutine = self._send_payload(payload)
|
||||
try:
|
||||
try:
|
||||
running_loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
running_loop = None
|
||||
|
||||
if running_loop is loop:
|
||||
future = loop.create_task(coroutine)
|
||||
else:
|
||||
future = asyncio.run_coroutine_threadsafe(coroutine, loop)
|
||||
except Exception:
|
||||
coroutine.close()
|
||||
|
||||
async def _send_payload(self, payload: str) -> None:
|
||||
broadcast_lock = self._broadcast_lock
|
||||
if broadcast_lock is None:
|
||||
return
|
||||
|
||||
async with broadcast_lock:
|
||||
with self._lock:
|
||||
sockets = list(self.active_websockets)
|
||||
|
||||
failed: list[WebSocket] = []
|
||||
for websocket in sockets:
|
||||
try:
|
||||
await websocket.send_text(payload)
|
||||
except Exception:
|
||||
failed.append(websocket)
|
||||
|
||||
if failed:
|
||||
with self._lock:
|
||||
for websocket in failed:
|
||||
self.active_websockets.discard(websocket)
|
||||
|
||||
def add_log(self, text: str, level: str = "info") -> None:
|
||||
log_entry = f"__LOG_LEVEL_{level.upper()}__:{text}"
|
||||
with self._lock:
|
||||
self.state.logs.append(log_entry)
|
||||
self.broadcast({"type": "log", "message": text, "level": level})
|
||||
|
||||
def start_export(self, config_data: dict[str, Any]) -> None:
|
||||
with self._lock:
|
||||
if self.state.status in ("preparing", "exporting", "stopping") or (
|
||||
self._thread is not None and self._thread.is_alive()
|
||||
):
|
||||
raise ValueError("Экспорт уже выполняется.")
|
||||
|
||||
self.state.reset()
|
||||
self.state.status = "preparing"
|
||||
self._thread = threading.Thread(target=self._run_subprocess, args=(config_data,), daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
self.broadcast({"type": "status", "status": self.state.status})
|
||||
self.add_log(f"Запуск экспорта: {config_data.get('model')} -> {config_data.get('format')}", "started")
|
||||
|
||||
def stop_export(self) -> None:
|
||||
with self._lock:
|
||||
if self.state.status not in ("preparing", "exporting"):
|
||||
return
|
||||
self.state.status = "stopping"
|
||||
self.state.stop_requested = True
|
||||
if self._process is not None:
|
||||
try:
|
||||
self._process.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.broadcast({"type": "status", "status": self.state.status})
|
||||
self.add_log("Запрошена остановка экспорта...", "warning")
|
||||
|
||||
def _run_subprocess(self, config_data: dict[str, Any]) -> None:
|
||||
temp_config_path = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(config_data, f)
|
||||
temp_config_path = f.name
|
||||
|
||||
cmd = [sys.executable, "-u", "-m", "yolo_webui.export_runner", temp_config_path]
|
||||
with self._lock:
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
|
||||
buffer = ""
|
||||
while True:
|
||||
char = self._process.stdout.read(1)
|
||||
if not char:
|
||||
if buffer:
|
||||
self._handle_subprocess_line(buffer)
|
||||
break
|
||||
|
||||
if char in ("\r", "\n"):
|
||||
if buffer:
|
||||
self._handle_subprocess_line(buffer)
|
||||
buffer = ""
|
||||
else:
|
||||
buffer += char
|
||||
|
||||
self._process.wait()
|
||||
rc = self._process.returncode
|
||||
|
||||
with self._lock:
|
||||
self._process = None
|
||||
|
||||
self._finalize_process_result(rc)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("Error in export process thread:")
|
||||
if self._process is not None:
|
||||
try:
|
||||
if self._process.poll() is None:
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with self._lock:
|
||||
self._process = None
|
||||
self.state.status = "failed"
|
||||
self.add_log(f"Внутренняя ошибка менеджера: {exc}", "error")
|
||||
finally:
|
||||
if temp_config_path and os.path.exists(temp_config_path):
|
||||
try:
|
||||
os.unlink(temp_config_path)
|
||||
except Exception:
|
||||
pass
|
||||
with self._lock:
|
||||
final_status = self.state.status
|
||||
output_dir = self.state.output_dir
|
||||
self.broadcast({"type": "status", "status": final_status, "output_dir": output_dir})
|
||||
|
||||
def _handle_subprocess_line(self, line_str: str) -> None:
|
||||
line_str = line_str.strip()
|
||||
if not line_str:
|
||||
return
|
||||
|
||||
if line_str == "__YOLO_WEBUI_READY__":
|
||||
with self._lock:
|
||||
self.state.status = "exporting"
|
||||
self.broadcast({"type": "status", "status": "exporting"})
|
||||
elif line_str.startswith("__YOLO_WEBUI_RESULT__:"):
|
||||
with self._lock:
|
||||
self.state.output_dir = line_str[len("__YOLO_WEBUI_RESULT__:") :]
|
||||
else:
|
||||
self.add_log(line_str, "info")
|
||||
|
||||
def _finalize_process_result(self, return_code: int) -> None:
|
||||
with self._lock:
|
||||
stopped = self.state.stop_requested
|
||||
|
||||
if stopped:
|
||||
status = "cancelled"
|
||||
message = "Экспорт остановлен пользователем."
|
||||
level = "warning"
|
||||
elif return_code == 0:
|
||||
status = "succeeded"
|
||||
message = "Экспорт успешно завершен."
|
||||
level = "success"
|
||||
else:
|
||||
status = "failed"
|
||||
message = "Процесс экспорта завершился с ошибкой."
|
||||
level = "error"
|
||||
|
||||
with self._lock:
|
||||
self.state.status = status
|
||||
self.add_log(message, level)
|
||||
|
||||
export_manager = ExportManager()
|
||||
|
||||
@app.get("/api/export/status")
|
||||
async def get_export_status():
|
||||
with export_manager._lock:
|
||||
return {
|
||||
"status": export_manager.state.status,
|
||||
"output_dir": export_manager.state.output_dir,
|
||||
"logs": export_manager.state.logs,
|
||||
}
|
||||
|
||||
@app.post("/api/export/start")
|
||||
async def start_export(config_data: dict[str, Any]):
|
||||
try:
|
||||
export_manager.start_export(config_data)
|
||||
return {"message": "Экспорт запущен."}
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
|
||||
@app.post("/api/export/stop")
|
||||
async def stop_export():
|
||||
export_manager.stop_export()
|
||||
return {"message": "Запрос на остановку экспорта отправлен."}
|
||||
|
||||
@app.websocket("/api/export/ws")
|
||||
async def export_websocket_endpoint(websocket: WebSocket):
|
||||
await websocket.accept()
|
||||
export_manager.add_websocket(websocket)
|
||||
|
||||
with export_manager._lock:
|
||||
state_dict = {
|
||||
"type": "init",
|
||||
"status": export_manager.state.status,
|
||||
"output_dir": export_manager.state.output_dir,
|
||||
"logs": [log.split(":", 1) for log in export_manager.state.logs if ":" in log],
|
||||
}
|
||||
try:
|
||||
await websocket.send_text(json.dumps(state_dict))
|
||||
while True:
|
||||
await websocket.receive_text()
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
export_manager.remove_websocket(websocket)
|
||||
|
||||
|
||||
@app.websocket("/api/ws")
|
||||
async def websocket_endpoint(websocket: WebSocket):
|
||||
await websocket.accept()
|
||||
|
|
|
|||
5
src/yolo_webui/config.py
Normal file → Executable file
5
src/yolo_webui/config.py
Normal file → Executable file
|
|
@ -8,7 +8,7 @@ from typing import Any, Literal
|
|||
|
||||
|
||||
YoloTask = Literal["detect", "segment", "classify", "pose", "obb"]
|
||||
AutoAugmentPolicy = Literal["randaugment", "autoaugment", "augmix"]
|
||||
AutoAugmentPolicy = Literal["none", "randaugment", "autoaugment", "augmix"]
|
||||
CopyPasteMode = Literal["flip", "mixup"]
|
||||
SUPPORTED_TASKS: tuple[YoloTask, ...] = (
|
||||
"detect",
|
||||
|
|
@ -18,6 +18,7 @@ SUPPORTED_TASKS: tuple[YoloTask, ...] = (
|
|||
"obb",
|
||||
)
|
||||
SUPPORTED_AUTO_AUGMENT_POLICIES: tuple[AutoAugmentPolicy, ...] = (
|
||||
"none",
|
||||
"randaugment",
|
||||
"autoaugment",
|
||||
"augmix",
|
||||
|
|
@ -161,7 +162,7 @@ class AugmentationConfig:
|
|||
"cutmix": self.cutmix,
|
||||
"copy_paste": self.copy_paste,
|
||||
"copy_paste_mode": self.copy_paste_mode,
|
||||
"auto_augment": self.auto_augment,
|
||||
"auto_augment": None if self.auto_augment == "none" else self.auto_augment,
|
||||
"erasing": self.erasing,
|
||||
"close_mosaic": self.close_mosaic,
|
||||
}
|
||||
|
|
|
|||
0
src/yolo_webui/dataset_splitter.py
Normal file → Executable file
0
src/yolo_webui/dataset_splitter.py
Normal file → Executable file
78
src/yolo_webui/export_runner.py
Executable file
78
src/yolo_webui/export_runner.py
Executable file
|
|
@ -0,0 +1,78 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
# Force headless Matplotlib to avoid any thread/process GUI issues
|
||||
os.environ["MPLBACKEND"] = "Agg"
|
||||
|
||||
def main() -> int:
|
||||
args = sys.argv[1:]
|
||||
if not args:
|
||||
print(
|
||||
"Usage: python -m yolo_webui.export_runner <config_json_path>",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
config_path = Path(args[0])
|
||||
|
||||
# The parent waits for this marker
|
||||
print("__YOLO_WEBUI_READY__", flush=True)
|
||||
|
||||
try:
|
||||
with config_path.open("r", encoding="utf-8") as config_file:
|
||||
config_dict = json.load(config_file)
|
||||
except Exception as exc:
|
||||
print(f"Error loading config: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
model_path = config_dict.get("model")
|
||||
export_format = config_dict.get("format", "onnx")
|
||||
imgsz = config_dict.get("imgsz", 640)
|
||||
half = config_dict.get("half", False)
|
||||
int8 = config_dict.get("int8", False)
|
||||
dynamic = config_dict.get("dynamic", False)
|
||||
simplify = config_dict.get("simplify", False)
|
||||
batch = config_dict.get("batch", 1)
|
||||
workspace = config_dict.get("workspace", 4)
|
||||
|
||||
if not model_path or not Path(model_path).exists():
|
||||
print(f"Model file not found: {model_path}")
|
||||
return 1
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
|
||||
print(f"Загрузка модели {model_path}...", flush=True)
|
||||
model = YOLO(model_path)
|
||||
|
||||
print(f"Запуск экспорта в формат {export_format}...", flush=True)
|
||||
print(f"Параметры: imgsz={imgsz}, half={half}, int8={int8}, dynamic={dynamic}, simplify={simplify}, batch={batch}, workspace={workspace}", flush=True)
|
||||
|
||||
exported_path = model.export(
|
||||
format=export_format,
|
||||
imgsz=imgsz,
|
||||
half=half,
|
||||
int8=int8,
|
||||
dynamic=dynamic,
|
||||
simplify=simplify,
|
||||
batch=batch,
|
||||
workspace=workspace
|
||||
)
|
||||
|
||||
print(f"Экспорт завершен успешно.", flush=True)
|
||||
if exported_path:
|
||||
# ultralytics returns either a single path (string) or list of paths depending on the format.
|
||||
if isinstance(exported_path, list):
|
||||
exported_path = exported_path[0]
|
||||
print(f"__YOLO_WEBUI_RESULT__:{exported_path}", flush=True)
|
||||
return 0
|
||||
except Exception as e:
|
||||
print(f"Ошибка при экспорте модели: {e}", flush=True)
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
254
src/yolo_webui/static/app.js
Normal file → Executable file
254
src/yolo_webui/static/app.js
Normal file → Executable file
|
|
@ -131,6 +131,32 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
let socket = null;
|
||||
let isTrainingActive = false;
|
||||
|
||||
// --- View Tab Switching ---
|
||||
const viewBtns = document.querySelectorAll('.view-btn');
|
||||
const views = document.querySelectorAll('.workspace-view');
|
||||
|
||||
viewBtns.forEach(btn => {
|
||||
btn.addEventListener('click', () => {
|
||||
viewBtns.forEach(b => b.classList.remove('active'));
|
||||
views.forEach(v => {
|
||||
v.classList.remove('active');
|
||||
});
|
||||
|
||||
btn.classList.add('active');
|
||||
const viewId = btn.dataset.view;
|
||||
const view = document.getElementById(viewId);
|
||||
view.classList.add('active');
|
||||
|
||||
localStorage.setItem('active_workspace_view', viewId);
|
||||
});
|
||||
});
|
||||
|
||||
const savedView = localStorage.getItem('active_workspace_view');
|
||||
if (savedView) {
|
||||
const viewBtn = Array.from(viewBtns).find(b => b.dataset.view === savedView);
|
||||
if (viewBtn) viewBtn.click();
|
||||
}
|
||||
|
||||
// --- Tab Switching ---
|
||||
tabs.forEach(tab => {
|
||||
tab.addEventListener('click', () => {
|
||||
|
|
@ -138,10 +164,13 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
tabContents.forEach(c => c.classList.remove('active'));
|
||||
|
||||
tab.classList.add('active');
|
||||
const contentId = `tab-${tab.dataset.tab}`;
|
||||
document.getElementById(contentId).classList.add('active');
|
||||
|
||||
localStorage.setItem('active_tab', tab.dataset.tab);
|
||||
// Check if it's a sub-tab (has data-tab)
|
||||
if (tab.dataset.tab) {
|
||||
const contentId = `tab-${tab.dataset.tab}`;
|
||||
const content = document.getElementById(contentId);
|
||||
if (content) content.classList.add('active');
|
||||
localStorage.setItem('active_tab', tab.dataset.tab);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -730,6 +759,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
const res = await fetch('/api/models');
|
||||
if (!res.ok) throw new Error();
|
||||
discoveredModels = await res.json();
|
||||
updateExportModelOptions();
|
||||
} catch (e) {
|
||||
console.error("Failed to load models list:", e);
|
||||
}
|
||||
|
|
@ -737,6 +767,46 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
|
||||
modelSelect.addEventListener('change', updateModelFieldsState);
|
||||
|
||||
// --- Export Model Selection ---
|
||||
const exportModelSelect = document.getElementById('export-model-select');
|
||||
const exportModelCustomWrapper = document.getElementById('export-model-custom-wrapper');
|
||||
const exportModelInput = document.getElementById('export-model');
|
||||
|
||||
function updateExportModelOptions() {
|
||||
exportModelSelect.innerHTML = '';
|
||||
|
||||
if (discoveredModels.length > 0) {
|
||||
const localGroup = document.createElement('optgroup');
|
||||
localGroup.label = 'Локальные/обученные модели';
|
||||
discoveredModels.forEach(m => {
|
||||
const opt = document.createElement('option');
|
||||
opt.value = m.path;
|
||||
opt.textContent = m.source === 'runs' ? `[RUNS] ${m.name}` : m.name;
|
||||
localGroup.appendChild(opt);
|
||||
});
|
||||
exportModelSelect.appendChild(localGroup);
|
||||
}
|
||||
|
||||
const customOpt = document.createElement('option');
|
||||
customOpt.value = '__custom__';
|
||||
customOpt.textContent = 'Указать путь вручную...';
|
||||
exportModelSelect.appendChild(customOpt);
|
||||
|
||||
updateExportModelFieldsState();
|
||||
}
|
||||
|
||||
function updateExportModelFieldsState() {
|
||||
const val = exportModelSelect.value;
|
||||
if (val === '__custom__') {
|
||||
exportModelCustomWrapper.style.display = 'block';
|
||||
} else {
|
||||
exportModelCustomWrapper.style.display = 'none';
|
||||
exportModelInput.value = val;
|
||||
}
|
||||
}
|
||||
|
||||
exportModelSelect.addEventListener('change', updateExportModelFieldsState);
|
||||
|
||||
sessionSelect.addEventListener('change', async () => {
|
||||
const name = sessionSelect.value;
|
||||
localStorage.setItem('selected_profile', name);
|
||||
|
|
@ -921,6 +991,181 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
`;
|
||||
document.head.appendChild(styleSheet);
|
||||
|
||||
// --- Export Logic ---
|
||||
const exportStartBtn = document.getElementById('export-start-btn');
|
||||
const exportStopBtn = document.getElementById('export-stop-btn');
|
||||
const exportStatusCard = document.getElementById('export-status-card');
|
||||
const exportStatusTitle = document.getElementById('export-status-title');
|
||||
const exportStatusText = document.getElementById('export-status-text');
|
||||
const exportLogContainer = document.getElementById('export-log-container');
|
||||
const exportClearLogBtn = document.getElementById('export-clear-log-btn');
|
||||
const exportAutoscrollCheck = document.getElementById('export-autoscroll');
|
||||
|
||||
let isExportActive = false;
|
||||
let exportSocket = null;
|
||||
|
||||
function addExportLogLine(message, level = 'info') {
|
||||
const line = document.createElement('div');
|
||||
line.className = `log-line log-level-${level.toLowerCase()}`;
|
||||
line.textContent = message;
|
||||
exportLogContainer.appendChild(line);
|
||||
|
||||
if (exportAutoscrollCheck.checked) {
|
||||
exportLogContainer.scrollTop = exportLogContainer.scrollHeight;
|
||||
}
|
||||
}
|
||||
|
||||
exportClearLogBtn.addEventListener('click', () => {
|
||||
exportLogContainer.innerHTML = '';
|
||||
});
|
||||
|
||||
function connectExportWebSocket() {
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const wsUrl = `${protocol}//${window.location.host}/api/export/ws`;
|
||||
|
||||
exportSocket = new WebSocket(wsUrl);
|
||||
|
||||
exportSocket.onopen = () => {
|
||||
addExportLogLine('Соединение с сервером установлено.', 'info');
|
||||
};
|
||||
|
||||
exportSocket.onclose = () => {
|
||||
addExportLogLine('Соединение потеряно. Повторная попытка через 5 секунд...', 'warning');
|
||||
setTimeout(connectExportWebSocket, 5000);
|
||||
};
|
||||
|
||||
exportSocket.onerror = (err) => {
|
||||
console.error('Export WS Error:', err);
|
||||
};
|
||||
|
||||
exportSocket.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data);
|
||||
|
||||
if (data.type === 'init') {
|
||||
updateExportUIStatus(data.status);
|
||||
exportLogContainer.innerHTML = '';
|
||||
data.logs.forEach(([levelCode, msg]) => {
|
||||
const level = levelCode.replace('__LOG_LEVEL_', '').replace('__', '').toLowerCase();
|
||||
addExportLogLine(msg, level);
|
||||
});
|
||||
} else if (data.type === 'status') {
|
||||
updateExportUIStatus(data.status);
|
||||
if (data.output_dir) {
|
||||
addExportLogLine(`Результаты сохранены: ${data.output_dir}`, 'success');
|
||||
}
|
||||
} else if (data.type === 'log') {
|
||||
addExportLogLine(data.message, data.level);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
function updateExportUIStatus(status) {
|
||||
exportStatusCard.className = `status-${status}`;
|
||||
|
||||
switch (status) {
|
||||
case 'idle':
|
||||
exportStatusTitle.textContent = 'ГОТОВО К ЭКСПОРТУ';
|
||||
exportStatusText.textContent = 'Выберите модель и формат для начала.';
|
||||
isExportActive = false;
|
||||
exportStartBtn.disabled = false;
|
||||
exportStopBtn.disabled = true;
|
||||
break;
|
||||
case 'preparing':
|
||||
case 'exporting':
|
||||
exportStatusTitle.textContent = 'ЭКСПОРТ';
|
||||
exportStatusText.textContent = 'Выполняется экспорт модели...';
|
||||
isExportActive = true;
|
||||
exportStartBtn.disabled = true;
|
||||
exportStopBtn.disabled = false;
|
||||
break;
|
||||
case 'stopping':
|
||||
exportStatusTitle.textContent = 'ОСТАНОВКА';
|
||||
exportStatusText.textContent = 'Остановка процесса экспорта...';
|
||||
isExportActive = true;
|
||||
exportStartBtn.disabled = true;
|
||||
exportStopBtn.disabled = true;
|
||||
break;
|
||||
case 'succeeded':
|
||||
exportStatusTitle.textContent = 'ГОТОВО';
|
||||
exportStatusText.textContent = 'Экспорт успешно завершен.';
|
||||
isExportActive = false;
|
||||
exportStartBtn.disabled = false;
|
||||
exportStopBtn.disabled = true;
|
||||
break;
|
||||
case 'cancelled':
|
||||
exportStatusTitle.textContent = 'ОСТАНОВЛЕНО';
|
||||
exportStatusText.textContent = 'Экспорт остановлен пользователем.';
|
||||
isExportActive = false;
|
||||
exportStartBtn.disabled = false;
|
||||
exportStopBtn.disabled = true;
|
||||
break;
|
||||
case 'failed':
|
||||
exportStatusTitle.textContent = 'ОШИБКА';
|
||||
exportStatusText.textContent = 'Экспорт завершился с ошибкой.';
|
||||
isExportActive = false;
|
||||
exportStartBtn.disabled = false;
|
||||
exportStopBtn.disabled = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
async function startExport() {
|
||||
if (isExportActive) return;
|
||||
|
||||
const modelVal = document.getElementById('export-model').value.trim();
|
||||
if (!modelVal) {
|
||||
showNotification('Пожалуйста, выберите или укажите модель для экспорта.', 'warning');
|
||||
return;
|
||||
}
|
||||
|
||||
const config = {
|
||||
model: modelVal,
|
||||
format: document.getElementById('export-format').value,
|
||||
imgsz: readNumber('export-imgsz', 640, true),
|
||||
half: document.getElementById('export-half').checked,
|
||||
int8: document.getElementById('export-int8').checked,
|
||||
dynamic: document.getElementById('export-dynamic').checked,
|
||||
simplify: document.getElementById('export-simplify').checked,
|
||||
batch: readNumber('export-batch', 1, true),
|
||||
workspace: readNumber('export-workspace', 4, true)
|
||||
};
|
||||
|
||||
try {
|
||||
const res = await fetch('/api/export/start', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(config)
|
||||
});
|
||||
|
||||
const data = await res.json();
|
||||
if (!res.ok) {
|
||||
throw new Error(data.detail || 'Failed to start export');
|
||||
}
|
||||
showNotification('Экспорт успешно запущен!', 'success');
|
||||
} catch (err) {
|
||||
console.error('Export start error:', err);
|
||||
showNotification(err.message, 'error');
|
||||
}
|
||||
}
|
||||
|
||||
async function stopExport() {
|
||||
if (!isExportActive) return;
|
||||
try {
|
||||
const res = await fetch('/api/export/stop', { method: 'POST' });
|
||||
if (!res.ok) {
|
||||
const data = await res.json();
|
||||
throw new Error(data.detail || 'Failed to stop export');
|
||||
}
|
||||
showNotification('Запрос на остановку отправлен.', 'info');
|
||||
} catch (err) {
|
||||
console.error('Export stop error:', err);
|
||||
showNotification(err.message, 'error');
|
||||
}
|
||||
}
|
||||
|
||||
exportStartBtn.addEventListener('click', startExport);
|
||||
exportStopBtn.addEventListener('click', stopExport);
|
||||
|
||||
// Initial load sequence
|
||||
loadDatasetsList().then(() => {
|
||||
return loadModelsList();
|
||||
|
|
@ -929,6 +1174,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}).then(() => {
|
||||
loadSessionsList();
|
||||
connectWebSocket();
|
||||
connectExportWebSocket();
|
||||
initChart();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
152
src/yolo_webui/static/index.html
Normal file → Executable file
152
src/yolo_webui/static/index.html
Normal file → Executable file
|
|
@ -20,9 +20,13 @@
|
|||
<polyline points="3.27 6.96 12 12.01 20.73 6.96"></polyline>
|
||||
<line x1="12" y1="22.08" x2="12" y2="12"></line>
|
||||
</svg>
|
||||
<div class="logo-text">
|
||||
<div class="logo-text" style="margin-right: 2rem;">
|
||||
<h1>YOLO Train Studio</h1>
|
||||
<span>Интерфейс обучения моделей</span>
|
||||
<span>Интерфейс обучения модели</span>
|
||||
</div>
|
||||
<div class="view-tabs" style="display: flex; gap: 10px;">
|
||||
<button class="view-btn active" data-view="train-view">Обучение</button>
|
||||
<button class="view-btn" data-view="export-view">Экспорт</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="header-actions">
|
||||
|
|
@ -38,6 +42,7 @@
|
|||
</header>
|
||||
|
||||
<main class="app-workspace">
|
||||
<div id="train-view" class="workspace-view active">
|
||||
<!-- Configuration Card -->
|
||||
<section class="pane" id="config-pane">
|
||||
<!-- Session Controls -->
|
||||
|
|
@ -291,6 +296,7 @@
|
|||
<div class="field">
|
||||
<label for="auto-augment">AutoAugment политика (classify)</label>
|
||||
<select id="auto-augment" name="auto-augment">
|
||||
<option value="none">Отключено</option>
|
||||
<option value="randaugment">RandAugment</option>
|
||||
<option value="autoaugment">AutoAugment</option>
|
||||
<option value="augmix">AugMix</option>
|
||||
|
|
@ -400,6 +406,148 @@
|
|||
</div>
|
||||
|
||||
</section>
|
||||
</div> <!-- End train-view -->
|
||||
|
||||
<div id="export-view" class="workspace-view">
|
||||
<!-- Export Configuration Card -->
|
||||
<section class="pane" id="export-config-pane">
|
||||
<div class="tabs">
|
||||
<button type="button" class="export-tab-btn active" style="flex: 1; background: var(--accent-gradient); border: none; border-radius: 6px; color: #fff; font-family: var(--font-sans); font-size: 0.9rem; font-weight: 600; padding: 0.6rem; cursor: default; box-shadow: 0 4px 10px rgba(0, 0, 0, 0.25);">Настройки экспорта</button>
|
||||
</div>
|
||||
<form id="export-form" class="form-container">
|
||||
<div style="display: block;">
|
||||
<div class="form-section-title">Параметры модели</div>
|
||||
<div class="field">
|
||||
<label for="export-model-select">Модель для экспорта</label>
|
||||
<select id="export-model-select">
|
||||
<option value="">Поиск моделей...</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="field" id="export-model-custom-wrapper" style="display: none;">
|
||||
<label for="export-model">Имя весов или путь к модели вручную</label>
|
||||
<input type="text" id="export-model" name="export-model" placeholder="например, /workspace/runs/.../best.pt">
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="export-format">Формат экспорта</label>
|
||||
<select id="export-format" name="export-format">
|
||||
<option value="onnx">ONNX (.onnx)</option>
|
||||
<option value="engine">TensorRT (.engine)</option>
|
||||
<option value="openvino">OpenVINO (_openvino_model/)</option>
|
||||
<option value="triton">Triton Inference Server (triton/)</option>
|
||||
<option value="torchscript">TorchScript (.torchscript)</option>
|
||||
<option value="coreml">CoreML (.mlpackage)</option>
|
||||
<option value="tflite">TFLite (.tflite)</option>
|
||||
<option value="pb">TensorFlow SavedModel (_saved_model/)</option>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
<div class="form-section-title">Опции</div>
|
||||
<div class="row">
|
||||
<div class="field">
|
||||
<label for="export-imgsz">Размер изображения (imgsz)</label>
|
||||
<input type="number" id="export-imgsz" name="export-imgsz" value="640">
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="export-batch" title="Максимальный размер батча. Важен при Dynamic=True для TensorRT и Triton.">Макс. батч (batch)</label>
|
||||
<input type="number" id="export-batch" name="export-batch" value="1" min="1">
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="export-workspace" title="Размер памяти для сборки TensorRT (ГБ).">Workspace (ГБ)</label>
|
||||
<input type="number" id="export-workspace" name="export-workspace" value="4" min="1" max="64">
|
||||
</div>
|
||||
</div>
|
||||
<div class="row">
|
||||
<div class="toggle-row" style="padding: 10px 0;">
|
||||
<div class="toggle-label">
|
||||
<h3>FP16 (Half)</h3>
|
||||
<p>Использовать полуточность</p>
|
||||
</div>
|
||||
<label class="switch-container">
|
||||
<input type="checkbox" id="export-half">
|
||||
<span class="switch-slider"></span>
|
||||
</label>
|
||||
</div>
|
||||
<div class="toggle-row" style="padding: 10px 0;">
|
||||
<div class="toggle-label">
|
||||
<h3>INT8</h3>
|
||||
<p>Квантование до INT8</p>
|
||||
</div>
|
||||
<label class="switch-container">
|
||||
<input type="checkbox" id="export-int8">
|
||||
<span class="switch-slider"></span>
|
||||
</label>
|
||||
</div>
|
||||
<div class="toggle-row" style="padding: 10px 0;">
|
||||
<div class="toggle-label">
|
||||
<h3>Dynamic</h3>
|
||||
<p>Динамические оси</p>
|
||||
</div>
|
||||
<label class="switch-container">
|
||||
<input type="checkbox" id="export-dynamic">
|
||||
<span class="switch-slider"></span>
|
||||
</label>
|
||||
</div>
|
||||
<div class="toggle-row" style="padding: 10px 0;">
|
||||
<div class="toggle-label" title="Только для ONNX. Оптимизирует и упрощает граф модели.">
|
||||
<h3>Simplify</h3>
|
||||
<p>Упростить граф (ONNX)</p>
|
||||
</div>
|
||||
<label class="switch-container">
|
||||
<input type="checkbox" id="export-simplify">
|
||||
<span class="switch-slider"></span>
|
||||
</label>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</form>
|
||||
</section>
|
||||
|
||||
<!-- Export Monitoring Card -->
|
||||
<section class="pane" id="export-run-pane">
|
||||
<div id="export-status-card" class="status-idle">
|
||||
<div class="status-header">
|
||||
<div class="status-indicator">
|
||||
<span class="status-dot"></span>
|
||||
<h2 id="export-status-title">ГОТОВО К ЭКСПОРТУ</h2>
|
||||
</div>
|
||||
</div>
|
||||
<p id="export-status-text">Выберите модель и формат для начала.</p>
|
||||
</div>
|
||||
|
||||
<!-- Action Buttons -->
|
||||
<div class="action-buttons">
|
||||
<button type="button" id="export-start-btn" class="btn btn-primary">
|
||||
<svg viewBox="0 0 24 24" fill="currentColor" class="btn-icon">
|
||||
<path d="M8 5v14l11-7z"></path>
|
||||
</svg>
|
||||
Начать экспорт
|
||||
</button>
|
||||
<button type="button" id="export-stop-btn" class="btn btn-danger" disabled>
|
||||
<svg viewBox="0 0 24 24" fill="currentColor" class="btn-icon">
|
||||
<path d="M6 19h4V5H6v14zm8-14v14h4V5h-4z"></path>
|
||||
</svg>
|
||||
Остановить
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<!-- Logs Box -->
|
||||
<div class="log-card" style="flex: 1;">
|
||||
<div class="log-header">
|
||||
<h3>Журнал экспорта</h3>
|
||||
<div class="log-actions">
|
||||
<label class="checkbox-inline">
|
||||
<input type="checkbox" id="export-autoscroll" checked>
|
||||
Автопрокрутка
|
||||
</label>
|
||||
<button type="button" id="export-clear-log-btn" class="btn-text">Очистить</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="log-body" id="export-log-container">
|
||||
<div class="log-line log-level-info">Ожидание запуска экспорта...</div>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
</div> <!-- End export-view -->
|
||||
</main>
|
||||
<script src="/static/app.js"></script>
|
||||
</body>
|
||||
|
|
|
|||
47
src/yolo_webui/static/style.css
Normal file → Executable file
47
src/yolo_webui/static/style.css
Normal file → Executable file
|
|
@ -129,9 +129,8 @@ body {
|
|||
/* App Layout Workspace */
|
||||
.app-workspace {
|
||||
flex: 1;
|
||||
display: grid;
|
||||
grid-template-columns: 46% 1fr;
|
||||
gap: 1.5rem;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 1.5rem 2rem;
|
||||
max-width: 1800px;
|
||||
width: 100%;
|
||||
|
|
@ -139,8 +138,21 @@ body {
|
|||
height: calc(100vh - 57px);
|
||||
}
|
||||
|
||||
.workspace-view {
|
||||
display: none;
|
||||
grid-template-columns: 46% 1fr;
|
||||
gap: 1.5rem;
|
||||
flex: 1;
|
||||
height: 100%;
|
||||
}
|
||||
|
||||
.workspace-view.active {
|
||||
display: grid;
|
||||
animation: fadeIn var(--transition-normal);
|
||||
}
|
||||
|
||||
@media (max-width: 1100px) {
|
||||
.app-workspace {
|
||||
.workspace-view {
|
||||
grid-template-columns: 1fr;
|
||||
height: auto;
|
||||
overflow-y: auto;
|
||||
|
|
@ -349,6 +361,33 @@ body {
|
|||
box-shadow: 0 4px 10px rgba(0, 0, 0, 0.25);
|
||||
}
|
||||
|
||||
/* View Button */
|
||||
.view-btn {
|
||||
flex: 0 1 auto;
|
||||
background: var(--bg-card);
|
||||
border: 1px solid var(--border-color);
|
||||
border-radius: 6px;
|
||||
color: var(--text-muted);
|
||||
font-family: var(--font-sans);
|
||||
font-size: 1rem;
|
||||
font-weight: 600;
|
||||
padding: 0.6rem 1.2rem;
|
||||
cursor: pointer;
|
||||
transition: var(--transition-fast);
|
||||
}
|
||||
|
||||
.view-btn:hover {
|
||||
color: var(--text-main);
|
||||
border-color: var(--accent);
|
||||
}
|
||||
|
||||
.view-btn.active {
|
||||
color: #fff;
|
||||
background: var(--bg-tertiary);
|
||||
border-color: var(--accent);
|
||||
box-shadow: 0 4px 10px rgba(0, 0, 0, 0.25);
|
||||
}
|
||||
|
||||
/* Configuration Form Layout */
|
||||
.form-container {
|
||||
flex: 1;
|
||||
|
|
|
|||
0
src/yolo_webui/subprocess_runner.py
Normal file → Executable file
0
src/yolo_webui/subprocess_runner.py
Normal file → Executable file
0
src/yolo_webui/trainer.py
Normal file → Executable file
0
src/yolo_webui/trainer.py
Normal file → Executable file
4
tests/frontend_smoke.js
Normal file → Executable file
4
tests/frontend_smoke.js
Normal file → Executable file
|
|
@ -141,13 +141,15 @@ async function flushPromises() {
|
|||
element('workers').value = '0';
|
||||
element('patience').value = '0';
|
||||
element('close-mosaic').value = '0';
|
||||
element('auto-augment').value = 'none';
|
||||
element('config-form').listeners.input();
|
||||
const savedConfig = JSON.parse(storage.get('draft_config'));
|
||||
assert.equal(savedConfig.workers, 0);
|
||||
assert.equal(savedConfig.patience, 0);
|
||||
assert.equal(savedConfig.augmentation.close_mosaic, 0);
|
||||
assert.equal(savedConfig.augmentation.auto_augment, 'none');
|
||||
|
||||
assert.equal(FakeWebSocket.instances.length, 1);
|
||||
assert.equal(FakeWebSocket.instances.length, 2);
|
||||
const socket = FakeWebSocket.instances[0];
|
||||
socket.onmessage({
|
||||
data: JSON.stringify({
|
||||
|
|
|
|||
0
tests/test_app.py
Normal file → Executable file
0
tests/test_app.py
Normal file → Executable file
22
tests/test_config.py
Normal file → Executable file
22
tests/test_config.py
Normal file → Executable file
|
|
@ -42,6 +42,28 @@ def test_augmentation_kwargs_are_passed_to_ultralytics() -> None:
|
|||
assert kwargs["auto_augment"] == "randaugment"
|
||||
|
||||
|
||||
def test_auto_augment_none_policy() -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
augmentation=AugmentationConfig(auto_augment="none"),
|
||||
)
|
||||
|
||||
kwargs = config.train_kwargs()
|
||||
assert kwargs["auto_augment"] is None
|
||||
|
||||
|
||||
def test_invalid_auto_augment_policy() -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
augmentation=AugmentationConfig(auto_augment="invalid_policy"), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Неизвестная политика AutoAugment"):
|
||||
config.validate()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["mosaic", "fliplr", "erasing", "perspective"])
|
||||
def test_augmentation_probabilities_are_validated(field: str) -> None:
|
||||
augmentation = AugmentationConfig(**{field: 1.1})
|
||||
|
|
|
|||
0
tests/test_frontend.py
Normal file → Executable file
0
tests/test_frontend.py
Normal file → Executable file
0
tests/test_splitter.py
Normal file → Executable file
0
tests/test_splitter.py
Normal file → Executable file
0
tests/test_subprocess_runner.py
Normal file → Executable file
0
tests/test_subprocess_runner.py
Normal file → Executable file
0
tests/test_trainer.py
Normal file → Executable file
0
tests/test_trainer.py
Normal file → Executable file
0
uv.lock
generated
Normal file → Executable file
0
uv.lock
generated
Normal file → Executable file
Loading…
Add table
Reference in a new issue