Static assets are missing. Place index.html under static/.
", status_code=404, ) return HTMLResponse(content=index_file.read_text(encoding="utf-8")) def get_sessions_dir() -> Path: path = Path("runs") / "sessions" path.mkdir(parents=True, exist_ok=True) return path def write_json_atomic(path: Path, data: dict[str, Any]) -> None: """Replace a JSON file atomically so concurrent readers never see partial data.""" temp_path: Path | None = None try: with tempfile.NamedTemporaryFile( "w", prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, delete=False, encoding="utf-8", ) as temp_file: temp_path = Path(temp_file.name) json.dump(data, temp_file, ensure_ascii=False, indent=2) os.replace(temp_path, path) temp_path = None finally: if temp_path is not None: try: temp_path.unlink() except FileNotFoundError: pass def parse_training_config(config_data: dict[str, Any]) -> TrainingConfig: config = TrainingConfig.from_dict(config_data) config.validate() return config def parse_export_config(config_data: dict[str, Any]) -> dict[str, Any]: config = ExportConfig.from_mapping(config_data) return { "model": str(config.model), "format": config.export_format, "imgsz": config.imgsz, "half": config.half, "int8": config.int8, "dynamic": config.dynamic, "simplify": config.simplify, "batch": config.batch, "workspace": config.workspace, } def get_session_path(name: str, *, allow_last_run: bool = True) -> Path: if SESSION_NAME_PATTERN.fullmatch(name) is None: raise HTTPException( status_code=400, detail="Имя сессии может содержать только латинские буквы, цифры, '_' и '-'.", ) if not allow_last_run and name == "last_run": raise HTTPException(status_code=400, detail="Имя 'last_run' зарезервировано.") return get_sessions_dir() / f"{name}.json" @app.get("/api/config/defaults") async def get_defaults(): # Return defaults by instantiating with dummy paths and serializing defaults = TrainingConfig(dataset="coco8.yaml", model="yolo11n.pt") return defaults.to_dict() @app.get("/api/sessions") async def list_sessions(): sessions_dir = get_sessions_dir() files = sessions_dir.glob("*.json") names = [f.stem for f in files if f.name != "last_run.json"] return sorted(names) @app.get("/api/datasets") async def list_datasets(): datasets_dir = Path("datasets") if not datasets_dir.exists(): return [] items = [] try: for path in datasets_dir.iterdir(): if path.is_dir() and not path.name.startswith("."): items.append({ "name": path.name, "path": str(path.absolute()), "type": "directory" }) elif path.is_file() and path.suffix.lower() in (".yaml", ".yml"): items.append({ "name": path.name, "path": str(path.absolute()), "type": "yaml" }) except Exception as e: logger.error(f"Failed to list datasets: {e}") return sorted(items, key=lambda x: x["name"]) @app.get("/api/models") async def list_models(): items = [] # 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.relative_to(runs_dir).as_posix(), "path": str(path.absolute()), "source": "runs" }) for path in runs_dir.rglob("*.pth"): items.append({ "name": path.relative_to(runs_dir).as_posix(), "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"]) @app.get("/api/sessions/{name}") async def load_session(name: str): file_path = get_session_path(name) if not file_path.exists(): raise HTTPException(status_code=404, detail="Сессия не найдена.") try: with file_path.open("r", encoding="utf-8") as f: return json.load(f) except Exception as exc: raise HTTPException(status_code=500, detail=f"Не удалось загрузить сессию: {exc}") @app.post("/api/sessions/{name}") async def save_session(name: str, config_data: dict[str, Any]): file_path = get_session_path(name, allow_last_run=False) try: config = parse_training_config(config_data) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) except Exception as exc: raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}") try: write_json_atomic(file_path, config.to_dict()) return {"message": "Сессия успешно сохранена."} except Exception as exc: raise HTTPException(status_code=500, detail=f"Не удалось сохранить сессию: {exc}") @app.delete("/api/sessions/{name}") async def delete_session(name: str): file_path = get_session_path(name, allow_last_run=False) if not file_path.exists(): raise HTTPException(status_code=404, detail="Сессия не найдена.") try: file_path.unlink() return {"message": "Сессия удалена."} except Exception as exc: raise HTTPException(status_code=500, detail=f"Не удалось удалить сессию: {exc}") @app.get("/api/train/status") async def get_status(): with manager._lock: return manager.state.snapshot() @app.post("/api/train/start") async def start_training(config_data: dict[str, Any]): try: config = parse_training_config(config_data) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) except Exception as exc: raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}") try: manager.start_training(config) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) # Only a configuration accepted for execution becomes the last run. try: sessions_dir = get_sessions_dir() last_run_path = sessions_dir / "last_run.json" write_json_atomic(last_run_path, config.to_dict()) except Exception as exc: logger.error(f"Failed to auto-save last run: {exc}") return {"message": "Обучение запущено."} @app.post("/api/train/stop") async def stop_training(): manager.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) future.add_done_callback(TrainingManager._log_broadcast_failure) except Exception: coroutine.close() logger.exception("Failed to schedule export WebSocket broadcast") 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 process = self._process if process is not None: if process.poll() is not None: # The subprocess has already completed; let its real return code # determine the final status instead of relabeling it cancelled. return try: process.terminate() except Exception: logger.warning("Failed to terminate export subprocess", exc_info=True) return self.state.status = "stopping" self.state.stop_requested = True 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 process: subprocess.Popen[str] | None = 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: if self.state.stop_requested: cancelled_before_start = True else: cancelled_before_start = False process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, ) self._process = process if cancelled_before_start: self._finalize_process_result(0) return if process is None or process.stdout is None: raise RuntimeError("Не удалось открыть stdout процесса экспорта.") buffer = "" while True: char = 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 process.wait() rc = process.returncode with self._lock: if self._process is process: self._process = None self._finalize_process_result(rc) except Exception as exc: logger.exception("Error in export process thread:") if process is not None: try: if process.poll() is None: process.kill() process.wait(timeout=5) except Exception: pass with self._lock: if self._process is process: 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__": status_changed = False with self._lock: if self.state.status == "preparing": self.state.status = "exporting" status_changed = True if status_changed: 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: snapshot = export_manager.state.snapshot() return { "status": snapshot["status"], "output_dir": snapshot["output_dir"], "logs": snapshot["logs"], } @app.post("/api/export/start") async def start_export(config_data: dict[str, Any]): try: validated_data = parse_export_config(config_data) export_manager.start_export(validated_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: snapshot = export_manager.state.snapshot() state_dict = { "type": "init", "status": snapshot["status"], "output_dir": snapshot["output_dir"], "logs": [log.split(":", 1) for log in snapshot["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() manager.add_websocket(websocket) # Send current state upon connection with manager._lock: snapshot = manager.state.snapshot() state_dict = { "type": "init", "status": snapshot["status"], "epoch": snapshot["epoch"], "total_epochs": snapshot["total_epochs"], "output_dir": snapshot["output_dir"], "metrics": snapshot["metrics"], # We format log items for the UI "logs": [log.split(":", 1) for log in snapshot["logs"] if ":" in log], } try: await websocket.send_text(json.dumps(state_dict)) while True: # Keep connection alive; discard incoming messages await websocket.receive_text() except WebSocketDisconnect: pass except Exception: logger.warning("WebSocket connection failed", exc_info=True) finally: manager.remove_websocket(websocket) def main() -> None: parser = argparse.ArgumentParser(description="YOLO Train Studio Web UI") parser.add_argument("--host", default="127.0.0.1", help="Host address to bind to") parser.add_argument("--port", type=int, default=8000, help="Port to bind to") args = parser.parse_args() # Headless matplotlib os.environ["MPLBACKEND"] = "Agg" logger.info(f"Starting server on http://{args.host}:{args.port}") uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": main()