TrainingPlatform_Django/datasets/views.py

227 lines
9.3 KiB
Python
Raw Normal View History

2025-06-03 16:10:11 +08:00
import requests
import os
import shutil
import glob
import zipfile
from django.shortcuts import render
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from django.core.files.storage import FileSystemStorage
from .models import Dataset, Project
from .tool import check_directory_structure_detect, delete_temp_dir, get_categories
from .minio_tools import upload_to_minio, delete_from_minio, get_dataset_link
from .config import TEMP_ROOT_DIR, LOCAL_IP, BUCKET_NAME
from concurrent.futures import ThreadPoolExecutor
executor = ThreadPoolExecutor(max_workers=10)
task_type_map = {'Detect':'Detection', 'Classify':'Classification'}
# Create your views here.
@csrf_exempt
def upload(request):
if request.method == 'POST':
# 检查文件是否存在
if 'file' in request.FILES:
dataset_name = request.POST.get('name')
user = request.POST.get('user')
task_type = request.POST.get('taskType')
size = request.POST.get('size')
description = request.POST.get('description')
uploaded_file = request.FILES.get('file')
## 检查数据库判断该数据集名称是否重复
if Dataset.objects.filter(name=dataset_name, user=user).exists():
return JsonResponse({'success': False, 'message': f'数据集 {dataset_name} 名称已存在'})
## 创建解压数据集的临时目录
temp_dir = TEMP_ROOT_DIR + '/' + user + '/' + dataset_name
# 如果目录已存在且不为空,则删除该目录
if os.path.exists(temp_dir) and os.listdir(temp_dir):
shutil.rmtree(temp_dir)
os.makedirs(temp_dir, exist_ok=True)
file_path = os.path.join(temp_dir, uploaded_file.name)
with open(file_path, 'wb') as f:
for chunk in uploaded_file.chunks():
f.write(chunk)
print(f"数据集 {dataset_name} 保存到临时目录成功")
## 开始解压数据集
try:
with zipfile.ZipFile(file_path, 'r') as zip_ref:
zip_ref.extractall(temp_dir)
print(f"数据集 {dataset_name} 解压到临时目录成功")
# 重命名解压后的文件夹
extracted_folder_name = os.path.splitext(uploaded_file.name)[0]
extracted_folder_path = os.path.join(temp_dir, extracted_folder_name)
new_folder_path = os.path.join(temp_dir, 'Detection')
os.rename(extracted_folder_path, new_folder_path)
except Exception as e:
print(f"数据集{dataset_name}解压失败", e)
return JsonResponse({'success': False, 'message': f'数据集 {dataset_name} 解压失败'})
# 根据 task_type 设置 data_path
if task_type == "Detect":
data_path = os.path.join(temp_dir, 'Detection')
elif task_type == "Classify":
data_path = os.path.join(temp_dir, 'Classification')
else:
data_path = os.path.join(temp_dir, 'Detection') # 默认值
# 根据 task_type 检查相应数据集的文件结构是否正确
flag = False
if task_type == "Detect":
flag, message = check_directory_structure_detect(data_path)
elif task_type == "Classify":
flag, message = check_directory_structure_detect(data_path)
else:
flag, message = check_directory_structure_detect(data_path)
if flag:
yaml_files = glob.glob(os.path.join(data_path, "*.yaml"))
yaml_path = yaml_files[0]
categories = get_categories(yaml_path)
print(f"数据集 {dataset_name} 结构检查成功")
else:
print(f"数据集 {dataset_name} 结构检查失败")
## 删除临时文件夹的所有数据
curr_temp_dir = os.path.join(os.path.join(TEMP_ROOT_DIR, user), dataset_name)
delete_temp_dir(curr_temp_dir)
return JsonResponse({'success': False, 'message': f'数据集 {dataset_name} 结构不符合要求: {message}'})
# 创建 Dataset 对象
dataset = Dataset(
name=dataset_name,
user=user, # 这里需要替换为实际的用户信息
task_type=task_type,
size=size,
number=0,
description=description,
categories = categories
)
future = executor.submit(upload_to_minio, data_path, dataset)
dataset.save()
if uploaded_file:
return JsonResponse({'success': True, 'message': '文件已成功上传'})
else:
return JsonResponse({'success': False, 'message': '没有文件上传'})
return JsonResponse({'success': False, 'message': '请求错误'})
def get_user_datasets(request):
user = request.GET.get('user')
datasets = Dataset.objects.filter(user=user)
return JsonResponse({'datasets': list(datasets.values())})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def create_project(request):
project_name = request.POST.get('name')
user = request.POST.get('user')
task_type = request.POST.get('taskType')
description = request.POST.get('description')
project = Project(name=project_name, user=user, task_type=task_type, description=description)
project.save()
return JsonResponse({'success': True, 'message': '项目已创建'})
def get_user_projects(request):
user = request.GET.get('user')
projects = Project.objects.filter(user=user)
return JsonResponse({'projects': list(projects.values())})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def delete_project(request):
user = request.GET.get('user')
project_name = request.GET.get('projectName')
project_deleted, _ = Project.objects.filter(user=user, name=project_name).delete()
if project_deleted > 0:
return JsonResponse({'success': True})
else:
return JsonResponse({'success': False})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def delete_dataset(request):
user = request.GET.get('user')
dataset_name = request.GET.get('datasetName')
project_deleted, _ = Dataset.objects.filter(user=user, name=dataset_name).delete()
delete_object_name = user + '/' + dataset_name
## 删除minio中的所有数据
future1 = executor.submit(delete_from_minio, delete_object_name)
## 删除临时文件夹的所有数据
curr_temp_dir = os.path.join(os.path.join(TEMP_ROOT_DIR, user), dataset_name)
future2 = executor.submit(delete_temp_dir, curr_temp_dir)
if project_deleted > 0:
return JsonResponse({'success': True})
else:
return JsonResponse({'success': False})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def get_project(request):
user = request.GET.get('user')
project_name = request.GET.get('projectName')
project = Project.objects.filter(user=user, name=project_name)
if project is not None:
return JsonResponse({'success': True, 'project': list(project.values())[0]})
else:
return JsonResponse({'success': False, 'project': None})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def get_dataset(request):
user = request.GET.get('user')
dataset_name = request.GET.get('datasetName')
dataset = Dataset.objects.filter(user=user, name=dataset_name)
if dataset is not None:
return JsonResponse({'success': True, 'dataset': list(dataset.values())[0]})
else:
return JsonResponse({'success': False, 'dataset': None})
@csrf_exempt # 仅在开发时使用,生产环境中请使用更安全的方式
def get_minio_links(request):
user = request.GET.get('user')
dataset_name = request.GET.get('datasetName')
next_image = request.GET.get('nextImage')
page_size = request.GET.get('pageSize')
if next_image != '':
next_image = os.path.relpath(next_image, f"http://{LOCAL_IP}:9000/{BUCKET_NAME}")
task_type = request.GET.get('taskType')
## 生成object_name路径用于访问minio
object_name = user + '/' + dataset_name + '/' + task_type + '/images'
http_links = get_dataset_link(object_name, next_image, page_size)
return JsonResponse({'success': True, 'links': http_links})
@csrf_exempt
def get_dataset_is_upload(request):
if request.method == 'GET':
user = request.GET.get('user')
dataset_name = request.GET.get('datasetName')
# 检查用户和数据集名称是否提供
if not user or not dataset_name:
return JsonResponse({'success': False, 'message': '用户或数据集名称未提供'})
# 查询数据库获取数据集
dataset = Dataset.objects.filter(user=user, name=dataset_name).first()
if dataset is not None:
return JsonResponse({'success': True, 'is_upload': dataset.is_upload})
else:
return JsonResponse({'success': False, 'message': '数据集未找到'})
return JsonResponse({'success': False, 'message': '请求方法错误'})