Project.py 10.0 KB

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