Project.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251
  1. from contextlib import closing
  2. from os.path import join
  3. from time import time
  4. from typing import List, Optional, Tuple
  5. from pycs.database.File import File
  6. from pycs.database.Label import Label
  7. from pycs.database.LabelProvider import LabelProvider
  8. from pycs.database.Model import Model
  9. class Project:
  10. """
  11. database class for projects
  12. """
  13. def __init__(self, database, row):
  14. self.database = database
  15. self.identifier = row[0]
  16. self.name = row[1]
  17. self.description = row[2]
  18. self.created = row[3]
  19. self.model_id = row[4]
  20. self.label_provider_id = row[5]
  21. self.root_folder = row[6]
  22. self.external_data = bool(row[7])
  23. self.data_folder = row[8]
  24. def model(self) -> Model:
  25. """
  26. get the model this project is associated with
  27. :return: model
  28. """
  29. return self.database.model(self.model_id)
  30. def label_provider(self) -> Optional[LabelProvider]:
  31. """
  32. get the label provider this project is associated with
  33. :return: label provider
  34. """
  35. if self.label_provider_id is not None:
  36. return self.database.label_provider(self.label_provider_id)
  37. return None
  38. def labels(self) -> List[Label]:
  39. """
  40. get a list of labels associated with this project
  41. :return: list of labels
  42. """
  43. with closing(self.database.con.cursor()) as cursor:
  44. cursor.execute('SELECT * FROM labels WHERE project = ?', [self.identifier])
  45. return list(map(
  46. lambda row: Label(self.database, row),
  47. cursor.fetchall()
  48. ))
  49. def label(self, identifier: int) -> Optional[Label]:
  50. """
  51. get a label using its unique identifier
  52. :param identifier: unique identifier
  53. :return: label
  54. """
  55. with closing(self.database.con.cursor()) as cursor:
  56. cursor.execute('SELECT * FROM labels WHERE id = ? AND project = ?',
  57. (identifier, self.identifier))
  58. row = cursor.fetchone()
  59. if row is not None:
  60. return Label(self.database, row)
  61. return None
  62. def create_label(self, name: str, reference: str = None,
  63. parent_id: int = None) -> Tuple[Optional[Label], bool]:
  64. """
  65. create a label for this project. If there is already a label with the same reference
  66. in the database its name is updated.
  67. :param name: label name
  68. :param reference: label reference
  69. :param parent_id: parent's identifier
  70. :return: created or edited label, insert
  71. """
  72. created = int(time())
  73. with closing(self.database.con.cursor()) as cursor:
  74. cursor.execute('''
  75. INSERT INTO labels (project, parent, created, reference, name)
  76. VALUES (?, ?, ?, ?, ?)
  77. ON CONFLICT (project, reference) DO
  78. UPDATE SET parent = ?, name = ?
  79. ''', (self.identifier, parent_id, created, reference, name, parent_id, name))
  80. # lastrowid is 0 if on conflict clause applies.
  81. # If this is the case we do an extra query to receive the row id.
  82. if cursor.lastrowid > 0:
  83. row_id = cursor.lastrowid
  84. insert = True
  85. else:
  86. cursor.execute('SELECT id FROM labels WHERE project = ? AND reference = ?',
  87. (self.identifier, reference))
  88. row_id = cursor.fetchone()[0]
  89. insert = False
  90. return self.label(row_id), insert
  91. def remove(self) -> None:
  92. """
  93. remove this project from the database
  94. :return:
  95. """
  96. with closing(self.database.con.cursor()) as cursor:
  97. cursor.execute('DELETE FROM projects WHERE id = ?', [self.identifier])
  98. def set_name(self, name: str) -> None:
  99. """
  100. set this projects name
  101. :param name: new name
  102. :return:
  103. """
  104. with closing(self.database.con.cursor()) as cursor:
  105. cursor.execute('UPDATE projects SET name = ? WHERE id = ?', (name, self.identifier))
  106. self.name = name
  107. def set_description(self, description: str) -> None:
  108. """
  109. set this projects description
  110. :param description: new description
  111. :return:
  112. """
  113. with closing(self.database.con.cursor()) as cursor:
  114. cursor.execute('UPDATE projects SET description = ? WHERE id = ?',
  115. (description, self.identifier))
  116. self.description = description
  117. def count_files(self) -> int:
  118. """
  119. count files associated with this project
  120. :return: count
  121. """
  122. with closing(self.database.con.cursor()) as cursor:
  123. cursor.execute('SELECT COUNT(*) FROM files WHERE project = ?', [self.identifier])
  124. return cursor.fetchone()[0]
  125. def files(self, offset=0, limit=-1) -> List[File]:
  126. """
  127. get a list of files associated with this project
  128. :param offset: file offset
  129. :param limit: file limit
  130. :return: list of files
  131. """
  132. with closing(self.database.con.cursor()) as cursor:
  133. cursor.execute('SELECT * FROM files WHERE project = ? ORDER BY id ASC LIMIT ? OFFSET ?',
  134. (self.identifier, limit, offset))
  135. return list(map(
  136. lambda row: File(self.database, row),
  137. cursor.fetchall()
  138. ))
  139. def files_without_results(self) -> List[File]:
  140. """
  141. get a list of files without associated results
  142. :return: list of files
  143. """
  144. with closing(self.database.con.cursor()) as cursor:
  145. cursor.execute('''
  146. SELECT files.*
  147. FROM files
  148. LEFT JOIN results ON files.id = results.file
  149. WHERE files.project = ? AND results.id IS NULL
  150. ORDER BY id ASC
  151. ''', [self.identifier])
  152. return list(map(
  153. lambda row: File(self.database, row),
  154. cursor.fetchall()
  155. ))
  156. def file(self, identifier) -> Optional[File]:
  157. """
  158. get a file using its unique identifier
  159. :param identifier: unique identifier
  160. :return: file
  161. """
  162. with closing(self.database.con.cursor()) as cursor:
  163. cursor.execute('SELECT * FROM files WHERE id = ? AND project = ?',
  164. (identifier, self.identifier))
  165. row = cursor.fetchone()
  166. if row is not None:
  167. return File(self.database, row)
  168. return None
  169. def add_file(self, uuid: str, file_type: str, name: str, extension: str, size: int,
  170. filename: str, frames: int = None, fps: float = None) -> Tuple[File, bool]:
  171. """
  172. add a file to this project
  173. :param uuid: unique identifier which is used for temporary files
  174. :param file_type: file type (either image or video)
  175. :param name: file name
  176. :param extension: file extension
  177. :param size: file size
  178. :param filename: actual name in filesystem
  179. :param frames: frame count
  180. :param fps: frames per second
  181. :return: file
  182. """
  183. created = int(time())
  184. path = join(self.data_folder, filename + extension)
  185. with closing(self.database.con.cursor()) as cursor:
  186. cursor.execute('''
  187. INSERT INTO files (
  188. uuid, project, type, name, extension, size, created, path, frames, fps
  189. )
  190. VALUES (
  191. ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
  192. )
  193. ON CONFLICT (project, path) DO
  194. UPDATE SET type = ?, name = ?, extension = ?, size = ?, frames = ?, fps = ?
  195. ''', (uuid, self.identifier, file_type, name, extension, size, created, path, frames,
  196. fps, file_type, name, extension, size, frames, fps))
  197. # lastrowid is 0 if on conflict clause applies.
  198. # If this is the case we do an extra query to receive the row id.
  199. if cursor.lastrowid > 0:
  200. row_id = cursor.lastrowid
  201. insert = True
  202. else:
  203. cursor.execute('SELECT id FROM files WHERE project = ? AND path = ?',
  204. (self.identifier, path))
  205. row_id = cursor.fetchone()[0]
  206. insert = False
  207. return self.file(row_id), insert