6
0

Project.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. import os
  2. import shutil
  3. import typing as T
  4. import warnings
  5. from datetime import datetime
  6. from pycs import db
  7. from pycs.database.base import NamedBaseModel
  8. from pycs.database.Collection import Collection
  9. from pycs.database.File import File
  10. from pycs.database.Label import Label
  11. from pycs.database.util import commit_on_return
  12. class Project(NamedBaseModel):
  13. """ DB Model for projects """
  14. description = db.Column(db.String)
  15. created = db.Column(db.DateTime, default=datetime.utcnow,
  16. index=True, nullable=False)
  17. model_id = db.Column(
  18. db.Integer,
  19. db.ForeignKey("model.id", ondelete="SET NULL"))
  20. label_provider_id = db.Column(
  21. db.Integer,
  22. db.ForeignKey("label_provider.id", ondelete="SET NULL"))
  23. root_folder = db.Column(db.String, nullable=False, unique=True)
  24. external_data = db.Column(db.Boolean, nullable=False)
  25. data_folder = db.Column(db.String, nullable=False)
  26. # contraints
  27. __table_args__ = ()
  28. # relationships to other models
  29. files = db.relationship(
  30. "File",
  31. backref="project",
  32. lazy="dynamic",
  33. passive_deletes=True,
  34. )
  35. labels = db.relationship(
  36. "Label",
  37. backref="project",
  38. lazy="dynamic",
  39. passive_deletes=True,
  40. )
  41. collections = db.relationship(
  42. "Collection",
  43. backref="project",
  44. lazy="dynamic",
  45. passive_deletes=True,
  46. )
  47. serialize_only = NamedBaseModel.serialize_only + (
  48. "created",
  49. "description",
  50. "model_id",
  51. "label_provider_id",
  52. "root_folder",
  53. "external_data",
  54. "data_folder",
  55. )
  56. @commit_on_return
  57. def delete(self) -> T.Tuple[dict, dict]:
  58. # pylint: disable=unexpected-keyword-arg
  59. dump = super().delete(commit=False)
  60. model_dump = {}
  61. if self.model_id is not None:
  62. # pylint: disable=unexpected-keyword-arg
  63. model_dump = self.model.delete(commit=False)
  64. if os.path.exists(self.root_folder):
  65. # remove from file system
  66. shutil.rmtree(self.root_folder)
  67. return dump, model_dump
  68. def label(self, identifier: int) -> T.Optional[Label]:
  69. """
  70. get a label using its unique identifier
  71. :param identifier: unique identifier
  72. :return: label
  73. """
  74. return self.labels.filter(Label.id == identifier).one_or_none()
  75. def label_by_reference(self, reference: str) -> T.Optional[Label]:
  76. """
  77. get a label using its reference string
  78. :param reference: reference string
  79. :return: label
  80. """
  81. return self.labels.filter(Label.reference == reference).one_or_none()
  82. def file(self, identifier: int) -> T.Optional[Label]:
  83. """
  84. get a file using its unique identifier
  85. :param identifier: unique identifier
  86. :return: file
  87. """
  88. return self.files.filter(File.id == identifier).one_or_none()
  89. def label_tree(self) -> T.List[Label]:
  90. """
  91. get a list of root labels associated with this project
  92. :return: list of labels
  93. """
  94. warnings.warn("Check performance of this method!")
  95. # pylint: disable=no-member
  96. return self.labels.filter(Label.parent_id.is_(None)).all()
  97. def label_tree_original(self):
  98. """
  99. get a list of root labels associated with this project
  100. :return: list of labels
  101. """
  102. raise NotImplementedError
  103. # pylint: disable=unreachable
  104. # pylint: disable=pointless-string-statement
  105. """
  106. with closing(self.database.con.cursor()) as cursor:
  107. cursor.execute('''
  108. WITH RECURSIVE
  109. tree AS (
  110. SELECT labels.* FROM labels
  111. WHERE project = ? AND parent IS NULL
  112. UNION ALL
  113. SELECT labels.* FROM labels
  114. JOIN tree ON labels.parent = tree.id
  115. )
  116. SELECT * FROM tree
  117. ''', [self.id])
  118. result = []
  119. lookup = {}
  120. for row in cursor.fetchall():
  121. label = TreeNodeLabel(self.database, row)
  122. lookup[label.id] = label
  123. if label.parent_id is None:
  124. result.append(label)
  125. else:
  126. lookup[label.parent_id].children.append(label)
  127. return result
  128. """
  129. def collection(self, identifier: int) -> T.Optional[Collection]:
  130. """
  131. get a collection using its unique identifier
  132. :param identifier: unique identifier
  133. :return: collection
  134. """
  135. return self.collections.filter(Collection.id == identifier).one_or_none()
  136. def collection_by_reference(self, reference: str) -> T.Optional[Collection]:
  137. """
  138. get a collection using its unique identifier
  139. :param identifier: unique identifier
  140. :return: collection
  141. """
  142. return self.collections.filter(Collection.reference == reference).one_or_none()
  143. @commit_on_return
  144. def create_label(self, name: str,
  145. reference: str = None,
  146. parent: T.Optional[T.Union[int, str, Label]] = None,
  147. hierarchy_level: str = None) -> T.Tuple[T.Optional[Label], bool]:
  148. """
  149. create a label for this project. If there is already a label with the same reference
  150. in the database its name is updated.
  151. :param name: label name
  152. :param reference: label reference
  153. :param parent: parent label. Either a reference string, a Label id or a Label instance
  154. :param hierarchy_level: hierarchy level name
  155. :return: created or edited label, insert
  156. """
  157. label = None
  158. is_new = False
  159. if reference is not None:
  160. label = Label.query.filter_by(project_id=self.id, reference=reference).one_or_none()
  161. if label is None:
  162. label = Label.new(commit=False, project_id=self.id, reference=reference)
  163. is_new = True
  164. label.set_name(name, commit=False)
  165. label.set_parent(parent, commit=False)
  166. label.hierarchy_level = hierarchy_level
  167. return label, is_new
  168. # pylint: disable=too-many-arguments
  169. @commit_on_return
  170. def create_collection(self,
  171. reference: str,
  172. name: str,
  173. description: str,
  174. position: int,
  175. autoselect: bool) -> T.Tuple[Collection, bool]:
  176. """
  177. create a new collection associated with this project
  178. :param reference: collection reference string
  179. :param name: collection name
  180. :param description: collection description
  181. :param position: position in menus
  182. :param autoselect: automatically select this collection on session load
  183. :return: collection object, insert
  184. """
  185. collection, is_new = Collection.get_or_create(
  186. project_id=self.id, reference=reference)
  187. collection.name = name
  188. collection.description = description
  189. collection.position = position
  190. collection.autoselect = autoselect
  191. return collection, is_new
  192. # pylint: disable=too-many-arguments
  193. @commit_on_return
  194. def add_file(self,
  195. uuid: str,
  196. file_type: str,
  197. name: str,
  198. extension: str,
  199. size: int,
  200. filename: str,
  201. frames: int = None,
  202. fps: float = None) -> T.Tuple[File, bool]:
  203. """
  204. add a file to this project
  205. :param uuid: unique identifier which is used for temporary files
  206. :param file_type: file type (either image or video)
  207. :param name: file name
  208. :param extension: file extension
  209. :param size: file size
  210. :param filename: actual name in filesystem
  211. :param frames: frame count
  212. :param fps: frames per second
  213. :return: file
  214. """
  215. path = os.path.join(self.data_folder, f"{filename}{extension}")
  216. file, is_new = File.get_or_create(
  217. project_id=self.id, path=path)
  218. file.uuid = uuid
  219. file.type = file_type
  220. file.name = name
  221. file.extension = extension
  222. file.size = size
  223. file.frames = frames
  224. file.fps = fps
  225. return file, is_new
  226. def get_files(self, offset: int = 0, limit: int = -1) -> T.List[File]:
  227. """
  228. get an iterator of files associated with this project
  229. :param offset: file offset
  230. :param limit: file limit
  231. :return: iterator of files
  232. """
  233. return self.files.order_by(File.id).offset(offset).limit(limit)
  234. def _files_without_results(self):
  235. """
  236. get files without any results
  237. :return: a query object
  238. """
  239. # pylint: disable=no-member
  240. return self.files.filter(~File.results.any())
  241. def count_files_without_results(self) -> int:
  242. """
  243. count files without associated results
  244. :return: count
  245. """
  246. return self._files_without_results().count()
  247. def files_without_results(self) -> T.List[File]:
  248. """
  249. get a list of files without associated results
  250. :return: list of files
  251. """
  252. return self._files_without_results().all()
  253. def _files_without_collection(self, offset: int = 0, limit: int = -1):
  254. """
  255. get files without a collection
  256. :return: a query object
  257. """
  258. # pylint: disable=no-member
  259. return self.get_files(offset, limit).filter(File.collection_id.is_(None))
  260. def files_without_collection(self, offset: int = 0, limit: int = -1) -> T.List[File]:
  261. """
  262. get a list of files without a collection
  263. :return: list of files
  264. """
  265. return self._files_without_collection(offset=offset, limit=limit).all()
  266. def count_files_without_collection(self) -> int:
  267. """
  268. count files associated with this project but without a collection
  269. :return: count
  270. """
  271. return self._files_without_collection().count()