diff --git a/Dockerfile b/Dockerfile index 6a19f7c..a98fc63 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,37 +1,22 @@ -FROM python:3.11-slim +FROM ultralytics/ultralytics:latest-gpu -# 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 diff --git a/docker-compose.yml b/docker-compose.yml index b711824..b88030b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -12,11 +12,11 @@ services: - ./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] restart: unless-stopped diff --git a/pyproject.toml b/pyproject.toml index 048d9f4..1821b12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/src/yolo_webui/app.py b/src/yolo_webui/app.py index 52936ab..2fa54ce 100644 --- a/src/yolo_webui/app.py +++ b/src/yolo_webui/app.py @@ -410,21 +410,41 @@ 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() diff --git a/src/yolo_webui/export_runner.py b/src/yolo_webui/export_runner.py new file mode 100644 index 0000000..31693e3 --- /dev/null +++ b/src/yolo_webui/export_runner.py @@ -0,0 +1,72 @@ +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 ", + 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) + + 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}", flush=True) + + exported_path = model.export( + format=export_format, + imgsz=imgsz, + half=half, + int8=int8, + dynamic=dynamic + ) + + 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()) diff --git a/src/yolo_webui/static/app.js b/src/yolo_webui/static/app.js index 5cbf1cb..375b0fc 100644 --- a/src/yolo_webui/static/app.js +++ b/src/yolo_webui/static/app.js @@ -131,6 +131,34 @@ 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.style.display = 'none'; + v.classList.remove('active'); + }); + + btn.classList.add('active'); + const viewId = btn.dataset.view; + const view = document.getElementById(viewId); + view.style.display = 'flex'; + 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 +166,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 +761,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 +769,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 +993,178 @@ 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 + }; + + 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 +1173,7 @@ document.addEventListener('DOMContentLoaded', () => { }).then(() => { loadSessionsList(); connectWebSocket(); + connectExportWebSocket(); initChart(); }); }); diff --git a/src/yolo_webui/static/index.html b/src/yolo_webui/static/index.html index bb488df..ef751b1 100644 --- a/src/yolo_webui/static/index.html +++ b/src/yolo_webui/static/index.html @@ -38,6 +38,12 @@
+
+ + +
+ +
@@ -400,6 +406,129 @@
+ + +
diff --git a/src/yolo_webui/static/style.css b/src/yolo_webui/static/style.css index 11c1d87..79963e7 100644 --- a/src/yolo_webui/static/style.css +++ b/src/yolo_webui/static/style.css @@ -349,6 +349,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;