dataset.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239
  1. # copytrue (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from ..utils import (set_folder_status, get_folder_status, DatasetStatus,
  15. TaskStatus, is_available, DownloadStatus,
  16. PretrainedModelStatus, ProjectType)
  17. from threading import Thread
  18. import random
  19. from .utils import copy_directory
  20. import traceback
  21. import shutil
  22. import psutil
  23. import pickle
  24. import os
  25. import os.path as osp
  26. import time
  27. import json
  28. import base64
  29. from .. import workspace_pb2 as w
  30. def create_dataset(data, workspace):
  31. """
  32. 创建dataset
  33. """
  34. create_time = time.time()
  35. time_array = time.localtime(create_time)
  36. create_time = time.strftime("%Y-%m-%d %H:%M:%S", time_array)
  37. id = workspace.max_dataset_id + 1
  38. if id < 10000:
  39. did = 'D%04d' % id
  40. else:
  41. did = 'D{}'.format(id)
  42. assert not did in workspace.datasets, "【数据集创建】ID'{}'已经被占用.".format(did)
  43. path = osp.join(workspace.path, 'datasets', did)
  44. if osp.exists(path):
  45. if not osp.isdir(path):
  46. os.remove(path)
  47. else:
  48. shutil.rmtree(path)
  49. os.makedirs(path)
  50. set_folder_status(path, DatasetStatus.XEMPTY)
  51. workspace.max_dataset_id = id
  52. ds = w.Dataset(
  53. id=did,
  54. name=data['name'],
  55. desc=data['desc'],
  56. type=data['dataset_type'],
  57. create_time=create_time,
  58. path=path)
  59. workspace.datasets[did].CopyFrom(ds)
  60. return {'status': 1, 'did': did}
  61. def import_dataset(data, workspace, monitored_processes, load_demo_proc_dict):
  62. """导入数据集到工作目录,包括数据检查和拷贝
  63. Args:
  64. data为dict, key包括
  65. 'did':数据集id,'path': 原数据集目录路径,
  66. 'demo'(可选): 该数据集为demo数据集
  67. """
  68. dataset_id = data['did']
  69. source_path = data['path']
  70. assert dataset_id in workspace.datasets, "数据集ID'{}'不存在.".format(dataset_id)
  71. dataset_type = workspace.datasets[dataset_id].type
  72. dataset_path = workspace.datasets[dataset_id].path
  73. valid_dataset_type = [
  74. 'classification', 'detection', 'segmentation', 'instance_segmentation',
  75. 'remote_segmentation'
  76. ]
  77. assert dataset_type in valid_dataset_type, "无法识别的数据类型{}".format(
  78. dataset_type)
  79. from .operate import import_dataset
  80. process = import_dataset(dataset_id, dataset_type, dataset_path,
  81. source_path)
  82. monitored_processes.put(process.pid)
  83. if 'demo' in data:
  84. prj_type = getattr(ProjectType, dataset_type)
  85. if prj_type not in load_demo_proc_dict:
  86. load_demo_proc_dict[prj_type] = []
  87. load_demo_proc_dict[prj_type].append(process)
  88. return {'status': 1}
  89. def delete_dataset(data, workspace):
  90. """删除dataset。
  91. Args:
  92. data为dict,key包括
  93. 'did'数据集id
  94. """
  95. dataset_id = data['did']
  96. assert dataset_id in workspace.datasets, "数据集ID'{}'不存在.".format(dataset_id)
  97. counter = 0
  98. for key in workspace.projects:
  99. if workspace.projects[key].did == dataset_id:
  100. counter += 1
  101. assert counter == 0, "无法删除数据集,当前仍被{}个项目中使用中,请先删除相关项目".format(counter)
  102. path = workspace.datasets[dataset_id].path
  103. if osp.exists(path):
  104. shutil.rmtree(path)
  105. del workspace.datasets[dataset_id]
  106. return {'status': 1}
  107. def get_dataset_status(data, workspace):
  108. """获取数据集当前状态
  109. Args:
  110. data为dict, key包括
  111. 'did':数据集id
  112. """
  113. from .operate import get_dataset_status
  114. dataset_id = data['did']
  115. assert dataset_id in workspace.datasets, "数据集ID'{}'不存在.".format(dataset_id)
  116. dataset_type = workspace.datasets[dataset_id].type
  117. dataset_path = workspace.datasets[dataset_id].path
  118. dataset_name = workspace.datasets[dataset_id].name
  119. dataset_desc = workspace.datasets[dataset_id].desc
  120. dataset_create_time = workspace.datasets[dataset_id].create_time
  121. status, message = get_dataset_status(dataset_id, dataset_type,
  122. dataset_path)
  123. dataset_pids = list()
  124. for key in workspace.projects:
  125. if dataset_id == workspace.projects[key].did:
  126. dataset_pids.append(workspace.projects[key].id)
  127. attr = {
  128. "type": dataset_type,
  129. "id": dataset_id,
  130. "name": dataset_name,
  131. "path": dataset_path,
  132. "desc": dataset_desc,
  133. "create_time": dataset_create_time,
  134. "pids": dataset_pids
  135. }
  136. return {
  137. 'status': 1,
  138. 'id': dataset_id,
  139. 'dataset_status': status.value,
  140. 'message': message,
  141. 'attr': attr
  142. }
  143. def list_datasets(workspace):
  144. """
  145. 列出数据集列表,可根据request中的参数进行筛选
  146. """
  147. from .operate import get_dataset_status
  148. dataset_list = list()
  149. for key in workspace.datasets:
  150. dataset_type = workspace.datasets[key].type
  151. dataset_id = workspace.datasets[key].id
  152. dataset_name = workspace.datasets[key].name
  153. dataset_path = workspace.datasets[key].path
  154. dataset_desc = workspace.datasets[key].desc
  155. dataset_create_time = workspace.datasets[key].create_time
  156. status, message = get_dataset_status(dataset_id, dataset_type,
  157. dataset_path)
  158. attr = {
  159. "type": dataset_type,
  160. "id": dataset_id,
  161. "name": dataset_name,
  162. "path": dataset_path,
  163. "desc": dataset_desc,
  164. "create_time": dataset_create_time,
  165. 'dataset_status': status.value,
  166. 'message': message
  167. }
  168. dataset_list.append({"id": dataset_id, "attr": attr})
  169. return {'status': 1, "datasets": dataset_list}
  170. def get_dataset_details(data, workspace):
  171. """获取数据集详情
  172. Args:
  173. data为dict, key包括
  174. 'did':数据集id
  175. Return:
  176. details(dict): 'file_info': 全量数据集文件与标签映射表,'label_info': 标签与全量数据集文件映射表,
  177. 'labels': 标签列表,'train_files': 训练集文件列表, 'val_files': 验证集文件列表,
  178. 'test_files': 测试集文件列表
  179. """
  180. from .operate import get_dataset_details
  181. dataset_id = data['did']
  182. assert dataset_id in workspace.datasets, "数据集ID'{}'不存在.".format(dataset_id)
  183. dataset_path = workspace.datasets[dataset_id].path
  184. details = get_dataset_details(dataset_path)
  185. return {'status': 1, 'details': details}
  186. def split_dataset(data, workspace):
  187. """将数据集切分为训练集、验证集和测试集
  188. Args:
  189. data为dict, key包括
  190. 'did':数据集id, 'val_split': 验证集比例, 'test_split': 测试集比例
  191. """
  192. from .operate import split_dataset
  193. from .operate import get_dataset_details
  194. dataset_id = data['did']
  195. assert dataset_id in workspace.datasets, "数据集ID'{}'不存在.".format(dataset_id)
  196. dataset_type = workspace.datasets[dataset_id].type
  197. dataset_path = workspace.datasets[dataset_id].path
  198. val_split = data['val_split']
  199. test_split = data['test_split']
  200. split_dataset(dataset_id, dataset_type, dataset_path, val_split,
  201. test_split)
  202. return {'status': 1}
  203. def img_base64(data):
  204. """将数据集切分为训练集、验证集和测试集
  205. Args:
  206. data为dict, key包括
  207. 'path':图片绝对路径
  208. """
  209. path = data['path']
  210. print(path)
  211. with open(path, 'rb') as f:
  212. base64_data = base64.b64encode(f.read())
  213. base64_str = str(base64_data, 'utf-8')
  214. return {'status': 1, 'img_data': base64_str}