113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
import os
|
|
from fastapi import APIRouter, UploadFile, File, Form, HTTPException
|
|
from pathlib import Path
|
|
from database import new_session, User, CSVFile
|
|
from repository import CSVFileRepository, UserRepository, H5ModelRepository, ModelStatisticsRepository
|
|
import shutil
|
|
import pandas as pd
|
|
from sqlalchemy import select
|
|
from schemas import UserCreate, ModelStatisticsCreate
|
|
|
|
router = APIRouter()
|
|
|
|
UPLOAD_FOLDER_CSV = Path("uploads/csv")
|
|
UPLOAD_FOLDER_CSV.mkdir(parents=True, exist_ok=True)
|
|
|
|
UPLOAD_FOLDER_MODELS = "uploads/models"
|
|
os.makedirs(UPLOAD_FOLDER_MODELS, exist_ok=True)
|
|
|
|
# Регистрация пользователя
|
|
@router.post("/users/")
|
|
async def create_user(user: UserCreate):
|
|
user_id = await UserRepository.create_user(user)
|
|
return {"user_id": user_id}
|
|
|
|
# Обновлённый метод для загрузки CSV файла
|
|
@router.post("/upload/csv/")
|
|
async def upload_csv(user_id: int = Form(...), file: UploadFile = File(...)):
|
|
# Получаем данные пользователя
|
|
async with new_session() as session:
|
|
user = await session.get(User, user_id)
|
|
if not user:
|
|
return {"error": "Пользователь не найден"}
|
|
username = user.username
|
|
|
|
# Формируем путь для сохранения
|
|
user_folder = UPLOAD_FOLDER_CSV / str(user_id)
|
|
user_folder.mkdir(parents=True, exist_ok=True)
|
|
file_path = user_folder / file.filename
|
|
|
|
# Сохраняем файл
|
|
with open(file_path, "wb") as buffer:
|
|
shutil.copyfileobj(file.file, buffer)
|
|
|
|
# Сохраняем данные в БД, заменяя обратные слэши на прямые
|
|
normalized_path = str(file_path).replace("\\", "/")
|
|
await CSVFileRepository.upload_file(user_id=user_id, file_path=normalized_path)
|
|
return {"message": "CSV файл загружен", "file_path": normalized_path}
|
|
|
|
|
|
# Загрузка H5 модели
|
|
@router.post("/upload/h5/")
|
|
async def upload_h5_model(user_id: int = Form(...), file: UploadFile = File(...)):
|
|
file_path = os.path.join(UPLOAD_FOLDER_MODELS, file.filename)
|
|
with open(file_path, "wb") as buffer:
|
|
shutil.copyfileobj(file.file, buffer)
|
|
|
|
await H5ModelRepository.add_model(user_id=user_id, model_path=file_path)
|
|
return {"message": "H5 модель загружена", "file_path": file_path}
|
|
|
|
# Добавление статистики модели
|
|
@router.post("/models/statistics/")
|
|
async def add_model_statistics(stats: ModelStatisticsCreate):
|
|
await ModelStatisticsRepository.add_statistics(stats)
|
|
return {"message": "Статистика модели сохранена"}
|
|
|
|
|
|
# Получение строк из CSV файла
|
|
@router.get("/csv/rows/")
|
|
async def get_csv_rows(
|
|
user_id: int,
|
|
file_name: str,
|
|
start_row: int,
|
|
end_row: int
|
|
):
|
|
# Проверка существования пользователя
|
|
async with new_session() as session:
|
|
user = await session.get(User, user_id)
|
|
if not user:
|
|
raise HTTPException(status_code=404, detail="Пользователь не найден")
|
|
|
|
# Формируем путь к файлу
|
|
user_folder = UPLOAD_FOLDER_CSV / str(user_id)
|
|
file_path = user_folder / file_name
|
|
|
|
# Проверка файла в БД, заменяя обратные слэши на прямые
|
|
normalized_file_path = str(file_path).replace("\\", "/")
|
|
result = await session.execute(
|
|
select(CSVFile).where(
|
|
CSVFile.user_id == user_id,
|
|
CSVFile.file_path == normalized_file_path
|
|
)
|
|
)
|
|
csv_file = result.scalars().first()
|
|
|
|
if not csv_file or not file_path.exists():
|
|
raise HTTPException(status_code=404, detail=f"Файл '{file_name}' не найден для пользователя '{user_id}'")
|
|
|
|
# Читаем файл и извлекаем строки
|
|
try:
|
|
if file_name.endswith(".xlsx"):
|
|
df = pd.read_excel(file_path)
|
|
elif file_name.endswith(".csv"):
|
|
df = pd.read_csv(file_path)
|
|
else:
|
|
raise HTTPException(status_code=400, detail="Неподдерживаемый формат файла")
|
|
|
|
rows = df.iloc[start_row:end_row].to_dict(orient="records")
|
|
return {"rows": rows}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Ошибка при чтении файла: {str(e)}")
|
|
|