Project.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352
  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 = Label.query.get(project=self, reference=reference)
  152. is_new = False
  153. if label is None:
  154. label = Label.new(project=self, reference=reference)
  155. is_new = True
  156. label.set_name(name, commit=False)
  157. label.set_parent(parent, commit=False)
  158. label.hierarchy_level = hierarchy_level
  159. return label, is_new
  160. # pylint: disable=too-many-arguments
  161. @commit_on_return
  162. def create_collection(self,
  163. reference: str,
  164. name: str,
  165. description: str,
  166. position: int,
  167. autoselect: bool) -> T.Tuple[Collection, bool]:
  168. """
  169. create a new collection associated with this project
  170. :param reference: collection reference string
  171. :param name: collection name
  172. :param description: collection description
  173. :param position: position in menus
  174. :param autoselect: automatically select this collection on session load
  175. :return: collection object, insert
  176. """
  177. collection, is_new = Collection.get_or_create(
  178. project_id=self.id, reference=reference)
  179. collection.name = name
  180. collection.description = description
  181. collection.position = position
  182. collection.autoselect = autoselect
  183. return collection, is_new
  184. # pylint: disable=too-many-arguments
  185. @commit_on_return
  186. def add_file(self,
  187. uuid: str,
  188. file_type: str,
  189. name: str,
  190. extension: str,
  191. size: int,
  192. filename: str,
  193. frames: int = None,
  194. fps: float = None) -> T.Tuple[File, bool]:
  195. """
  196. add a file to this project
  197. :param uuid: unique identifier which is used for temporary files
  198. :param file_type: file type (either image or video)
  199. :param name: file name
  200. :param extension: file extension
  201. :param size: file size
  202. :param filename: actual name in filesystem
  203. :param frames: frame count
  204. :param fps: frames per second
  205. :return: file
  206. """
  207. path = os.path.join(self.data_folder, f"{filename}{extension}")
  208. file, is_new = File.get_or_create(
  209. project_id=self.id, path=path)
  210. file.uuid = uuid
  211. file.type = file_type
  212. file.name = name
  213. file.extension = extension
  214. file.size = size
  215. file.frames = frames
  216. file.fps = fps
  217. return file, is_new
  218. def get_files(self, offset: int = 0, limit: int = -1) -> T.List[File]:
  219. """
  220. get an iterator of files associated with this project
  221. :param offset: file offset
  222. :param limit: file limit
  223. :return: iterator of files
  224. """
  225. return self.files.order_by(File.id).offset(offset).limit(limit)
  226. def _files_without_results(self):
  227. """
  228. get files without any results
  229. :return: a query object
  230. """
  231. # pylint: disable=no-member
  232. return self.files.filter(~File.results.any())
  233. def count_files_without_results(self) -> int:
  234. """
  235. count files without associated results
  236. :return: count
  237. """
  238. return self._files_without_results().count()
  239. def files_without_results(self) -> T.List[File]:
  240. """
  241. get a list of files without associated results
  242. :return: list of files
  243. """
  244. return self._files_without_results().all()
  245. def _files_without_collection(self, offset: int = 0, limit: int = -1):
  246. """
  247. get files without a collection
  248. :return: a query object
  249. """
  250. # pylint: disable=no-member
  251. return self.get_files(offset, limit).filter(File.collection_id.is_(None))
  252. def files_without_collection(self, offset: int = 0, limit: int = -1) -> T.List[File]:
  253. """
  254. get a list of files without a collection
  255. :return: list of files
  256. """
  257. return self._files_without_collection(offset=offset, limit=limit).all()
  258. def count_files_without_collection(self) -> int:
  259. """
  260. count files associated with this project but without a collection
  261. :return: count
  262. """
  263. return self._files_without_collection().count()