Project.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249
  1. import os
  2. import typing as T
  3. from contextlib import closing
  4. from datetime import datetime
  5. from pycs import db
  6. from pycs.database.base import NamedBaseModel
  7. from pycs.database.Collection import Collection
  8. from pycs.database.File import File
  9. from pycs.database.Label import Label
  10. class Project(NamedBaseModel):
  11. description = db.Column(db.String)
  12. created = db.Column(db.DateTime, default=datetime.utcnow,
  13. index=True, nullable=False)
  14. model_id = db.Column(
  15. db.Integer,
  16. db.ForeignKey("model.id", ondelete="SET NULL"))
  17. label_provider_id = db.Column(
  18. db.Integer,
  19. db.ForeignKey("label_provider.id", ondelete="SET NULL"))
  20. root_folder = db.Column(db.String, nullable=False, unique=True)
  21. external_data = db.Column(db.Boolean, nullable=False)
  22. data_folder = db.Column(db.String, nullable=False)
  23. # contraints
  24. __table_args__ = ()
  25. # relationships to other models
  26. files = db.relationship("File",
  27. backref="project",
  28. lazy="dynamic",
  29. passive_deletes=True)
  30. labels = db.relationship("Label",
  31. backref="project",
  32. lazy="dynamic",
  33. passive_deletes=True)
  34. collections = db.relationship("Collection",
  35. backref="project",
  36. lazy="dynamic",
  37. passive_deletes=True)
  38. serialize_only = (
  39. "id",
  40. "name",
  41. "created",
  42. "description",
  43. "model_id",
  44. "label_provider_id",
  45. "root_folder",
  46. "external_data",
  47. "data_folder",
  48. )
  49. def label(self, id: int) -> T.Optional[Label]:
  50. """
  51. get a label using its unique identifier
  52. :param identifier: unique identifier
  53. :return: label
  54. """
  55. return self.labels.filter_by(id=id).one_or_none()
  56. def file(self, id: int) -> T.Optional[Label]:
  57. """
  58. get a file using its unique identifier
  59. :param identifier: unique identifier
  60. :return: file
  61. """
  62. return self.files.filter_by(id=id).one_or_none()
  63. def collection(self, id: int) -> T.Optional[Collection]:
  64. """
  65. get a collection using its unique identifier
  66. :param identifier: unique identifier
  67. :return: collection
  68. """
  69. return self.collections.filter_by(id=id).one_or_none()
  70. def collection_by_reference(self, reference: str) -> T.Optional[Collection]:
  71. """
  72. get a collection using its unique identifier
  73. :param identifier: unique identifier
  74. :return: collection
  75. """
  76. return self.collections.filter_by(reference=reference).one_or_none()
  77. def create_label(self, name: str, reference: str = None,
  78. parent_id: int = None) -> T.Tuple[T.Optional[Label], bool]:
  79. """
  80. create a label for this project. If there is already a label with the same reference
  81. in the database its name is updated.
  82. :param name: label name
  83. :param reference: label reference
  84. :param parent_id: parent's identifier
  85. :return: created or edited label, insert
  86. """
  87. label, is_new = Label.get_or_create(project=self, reference=reference)
  88. label.set_name(name)
  89. label.set_parent(parent_id)
  90. self.commit()
  91. return label, is_new
  92. def create_collection(self,
  93. reference: str,
  94. name: str,
  95. description: str,
  96. position: int,
  97. autoselect: bool):
  98. collection, is_new = Collection.get_or_create(project=self, reference=reference)
  99. collection.name = name
  100. collection.description = description
  101. collection.position = position
  102. collection.autoselect = autoselect
  103. self.commit()
  104. return collection, is_new
  105. def add_file(self, uuid: str, file_type: str, name: str, extension: str, size: int,
  106. filename: str, frames: int = None, fps: float = None) -> T.Tuple[File, bool]:
  107. """
  108. add a file to this project
  109. :param uuid: unique identifier which is used for temporary files
  110. :param file_type: file type (either image or video)
  111. :param name: file name
  112. :param extension: file extension
  113. :param size: file size
  114. :param filename: actual name in filesystem
  115. :param frames: frame count
  116. :param fps: frames per second
  117. :return: file
  118. """
  119. path = os.path.join(self.data_folder, filename + extension)
  120. file, is_new = File.get_or_create(project=self, path=path)
  121. file.uuid = uuid
  122. file.type = file_type
  123. file.name = name
  124. file.extension = extension
  125. file.size = size
  126. file.frames = frames
  127. file.fps = fps
  128. self.commit()
  129. return file, is_new
  130. def set_description(self, description: str):
  131. """
  132. set this projects description
  133. :param description: new description
  134. :return:
  135. """
  136. self.description = description
  137. self.commit()
  138. def count_files(self) -> int:
  139. """
  140. count files associated with this project
  141. :return: count
  142. """
  143. return self.files.count()
  144. def get_files(self, offset: int = 0, limit: int = -1) -> T.Iterator[File]:
  145. """
  146. get an iterator of files associated with this project
  147. :param offset: file offset
  148. :param limit: file limit
  149. :return: iterator of files
  150. """
  151. return self.files.order_by(File.id).offset(offset).limit(limit)
  152. def count_files_without_results(self) -> int:
  153. """
  154. count files without associated results
  155. :return: count
  156. """
  157. raise NotImplementedError
  158. with closing(self.database.con.cursor()) as cursor:
  159. cursor.execute('''
  160. SELECT COUNT(*)
  161. FROM files
  162. LEFT JOIN results ON files.id = results.file
  163. WHERE files.project = ? AND results.id IS NULL
  164. ''', [self.id])
  165. return cursor.fetchone()[0]
  166. def files_without_results(self) -> T.Iterator[File]:
  167. """
  168. get an iterator of files without associated results
  169. :return: list of files
  170. """
  171. raise NotImplementedError
  172. with closing(self.database.con.cursor()) as cursor:
  173. cursor.execute('''
  174. SELECT files.*
  175. FROM files
  176. LEFT JOIN results ON files.id = results.file
  177. WHERE files.project = ? AND results.id IS NULL
  178. ORDER BY id ASC
  179. ''', [self.id])
  180. for row in cursor:
  181. yield File(self.database, row)
  182. def files_without_collection(self, offset: int = 0, limit: int = -1) -> T.Iterator[File]:
  183. """
  184. get an iterator of files without not associated with any collection
  185. :return: list of files
  186. """
  187. return self.get_files(offset, limit).filter(File.collection_id == None)
  188. def count_files_without_collection(self) -> int:
  189. """
  190. count files associated with this project but with no collection
  191. :return: count
  192. """
  193. return self.files_without_collection().count()