Project.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343
  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.Collection import Collection
  6. from pycs.database.File import File
  7. from pycs.database.Label import Label
  8. from pycs.database.LabelProvider import LabelProvider
  9. from pycs.database.Model import Model
  10. class Project:
  11. """
  12. database class for projects
  13. """
  14. def __init__(self, database, row):
  15. self.database = database
  16. self.identifier = row[0]
  17. self.name = row[1]
  18. self.description = row[2]
  19. self.created = row[3]
  20. self.model_id = row[4]
  21. self.label_provider_id = row[5]
  22. self.root_folder = row[6]
  23. self.external_data = bool(row[7])
  24. self.data_folder = row[8]
  25. def model(self) -> Model:
  26. """
  27. get the model this project is associated with
  28. :return: model
  29. """
  30. return self.database.model(self.model_id)
  31. def label_provider(self) -> Optional[LabelProvider]:
  32. """
  33. get the label provider this project is associated with
  34. :return: label provider
  35. """
  36. if self.label_provider_id is not None:
  37. return self.database.label_provider(self.label_provider_id)
  38. return None
  39. def labels(self) -> List[Label]:
  40. """
  41. get a list of labels associated with this project
  42. :return: list of labels
  43. """
  44. with closing(self.database.con.cursor()) as cursor:
  45. cursor.execute('SELECT * FROM labels WHERE project = ?', [self.identifier])
  46. return list(map(
  47. lambda row: Label(self.database, row),
  48. cursor.fetchall()
  49. ))
  50. def label(self, identifier: int) -> Optional[Label]:
  51. """
  52. get a label using its unique identifier
  53. :param identifier: unique identifier
  54. :return: label
  55. """
  56. with closing(self.database.con.cursor()) as cursor:
  57. cursor.execute('SELECT * FROM labels WHERE id = ? AND project = ?',
  58. (identifier, self.identifier))
  59. row = cursor.fetchone()
  60. if row is not None:
  61. return Label(self.database, row)
  62. return None
  63. def create_label(self, name: str, reference: str = 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. :return: created or edited label, insert
  70. """
  71. created = int(time())
  72. with closing(self.database.con.cursor()) as cursor:
  73. cursor.execute('''
  74. INSERT INTO labels (project, created, reference, name)
  75. VALUES (?, ?, ?, ?)
  76. ON CONFLICT (project, reference) DO
  77. UPDATE SET name = ?
  78. ''', (self.identifier, created, reference, name, name))
  79. # lastrowid is 0 if on conflict clause applies.
  80. # If this is the case we do an extra query to receive the row id.
  81. if cursor.lastrowid > 0:
  82. row_id = cursor.lastrowid
  83. insert = True
  84. else:
  85. cursor.execute('SELECT id FROM labels WHERE project = ? AND reference = ?',
  86. (self.identifier, reference))
  87. row_id = cursor.fetchone()[0]
  88. insert = False
  89. return self.label(row_id), insert
  90. def collections(self) -> List[Collection]:
  91. """
  92. get a list of collections associated with this project
  93. :return: list of collections
  94. """
  95. with closing(self.database.con.cursor()) as cursor:
  96. cursor.execute('SELECT * FROM collections WHERE project = ? ORDER BY position ASC',
  97. [self.identifier])
  98. return list(map(
  99. lambda row: Collection(self.database, row),
  100. cursor.fetchall()
  101. ))
  102. def collection(self, identifier: int) -> Optional[Collection]:
  103. """
  104. get a collection using its unique identifier
  105. :param identifier: unique identifier
  106. :return: collection
  107. """
  108. with closing(self.database.con.cursor()) as cursor:
  109. cursor.execute('SELECT * FROM collections WHERE id = ? AND project = ?',
  110. (identifier, self.identifier))
  111. row = cursor.fetchone()
  112. if row is not None:
  113. return Collection(self.database, row)
  114. return None
  115. def create_collection(self,
  116. reference: str,
  117. name: str,
  118. description: str,
  119. position: int,
  120. autoselect: bool):
  121. autoselect = 1 if autoselect else 0
  122. with closing(self.database.con.cursor()) as cursor:
  123. cursor.execute('''
  124. INSERT INTO collections
  125. (project, reference, name, description, position, autoselect)
  126. VALUES (?, ?, ?, ?, ?, ?)
  127. ON CONFLICT (project, reference) DO
  128. UPDATE SET name = ?, description = ?, position = ?, autoselect = ?
  129. ''', (self.identifier, reference, name, description, position, autoselect,
  130. name, description, position, autoselect))
  131. # lastrowid is 0 if on conflict clause applies.
  132. # If this is the case we do an extra query to receive the row id.
  133. if cursor.lastrowid > 0:
  134. row_id = cursor.lastrowid
  135. insert = True
  136. else:
  137. cursor.execute('SELECT id FROM collections WHERE project = ? AND reference = ?',
  138. (self.identifier, reference))
  139. row_id = cursor.fetchone()[0]
  140. insert = False
  141. return self.collection(row_id), insert
  142. def remove(self) -> None:
  143. """
  144. remove this project from the database
  145. :return:
  146. """
  147. with closing(self.database.con.cursor()) as cursor:
  148. cursor.execute('DELETE FROM projects WHERE id = ?', [self.identifier])
  149. def set_name(self, name: str) -> None:
  150. """
  151. set this projects name
  152. :param name: new name
  153. :return:
  154. """
  155. with closing(self.database.con.cursor()) as cursor:
  156. cursor.execute('UPDATE projects SET name = ? WHERE id = ?', (name, self.identifier))
  157. self.name = name
  158. def set_description(self, description: str) -> None:
  159. """
  160. set this projects description
  161. :param description: new description
  162. :return:
  163. """
  164. with closing(self.database.con.cursor()) as cursor:
  165. cursor.execute('UPDATE projects SET description = ? WHERE id = ?',
  166. (description, self.identifier))
  167. self.description = description
  168. def count_files(self) -> int:
  169. """
  170. count files associated with this project
  171. :return: count
  172. """
  173. with closing(self.database.con.cursor()) as cursor:
  174. cursor.execute('SELECT COUNT(*) FROM files WHERE project = ?', [self.identifier])
  175. return cursor.fetchone()[0]
  176. def files(self, offset=0, limit=-1) -> List[File]:
  177. """
  178. get a list of files associated with this project
  179. :param offset: file offset
  180. :param limit: file limit
  181. :return: list of files
  182. """
  183. with closing(self.database.con.cursor()) as cursor:
  184. cursor.execute('SELECT * FROM files WHERE project = ? ORDER BY id ASC LIMIT ? OFFSET ?',
  185. (self.identifier, limit, offset))
  186. return list(map(
  187. lambda row: File(self.database, row),
  188. cursor.fetchall()
  189. ))
  190. def files_without_results(self) -> List[File]:
  191. """
  192. get a list of files without associated results
  193. :return: list of files
  194. """
  195. with closing(self.database.con.cursor()) as cursor:
  196. cursor.execute('''
  197. SELECT files.*
  198. FROM files
  199. LEFT JOIN results ON files.id = results.file
  200. WHERE files.project = ? AND results.id IS NULL
  201. ORDER BY id ASC
  202. ''', [self.identifier])
  203. return list(map(
  204. lambda row: File(self.database, row),
  205. cursor.fetchall()
  206. ))
  207. def count_files_without_collection(self) -> int:
  208. """
  209. count files associated with this project but with no collection
  210. :return: count
  211. """
  212. with closing(self.database.con.cursor()) as cursor:
  213. cursor.execute('SELECT COUNT(*) FROM files WHERE project = ? AND collection IS NULL',
  214. [self.identifier])
  215. return cursor.fetchone()[0]
  216. def files_without_collection(self, offset=0, limit=-1) -> List[File]:
  217. """
  218. get a list of files without not associated with any collection
  219. :return: list of files
  220. """
  221. with closing(self.database.con.cursor()) as cursor:
  222. cursor.execute('''
  223. SELECT * FROM files
  224. WHERE files.project = ? AND files.collection IS NULL
  225. ORDER BY id ASC
  226. LIMIT ? OFFSET ?
  227. ''', (self.identifier, limit, offset))
  228. return list(map(
  229. lambda row: File(self.database, row),
  230. cursor.fetchall()
  231. ))
  232. def file(self, identifier) -> Optional[File]:
  233. """
  234. get a file using its unique identifier
  235. :param identifier: unique identifier
  236. :return: file
  237. """
  238. with closing(self.database.con.cursor()) as cursor:
  239. cursor.execute('SELECT * FROM files WHERE id = ? AND project = ?',
  240. (identifier, self.identifier))
  241. row = cursor.fetchone()
  242. if row is not None:
  243. return File(self.database, row)
  244. return None
  245. def add_file(self, uuid: str, file_type: str, name: str, extension: str, size: int,
  246. filename: str, frames: int = None, fps: float = None) -> Tuple[File, bool]:
  247. """
  248. add a file to this project
  249. :param uuid: unique identifier which is used for temporary files
  250. :param file_type: file type (either image or video)
  251. :param name: file name
  252. :param extension: file extension
  253. :param size: file size
  254. :param filename: actual name in filesystem
  255. :param frames: frame count
  256. :param fps: frames per second
  257. :return: file
  258. """
  259. created = int(time())
  260. path = join(self.data_folder, filename + extension)
  261. with closing(self.database.con.cursor()) as cursor:
  262. cursor.execute('''
  263. INSERT INTO files (
  264. uuid, project, type, name, extension, size, created, path, frames, fps
  265. )
  266. VALUES (
  267. ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
  268. )
  269. ON CONFLICT (project, path) DO
  270. UPDATE SET type = ?, name = ?, extension = ?, size = ?, frames = ?, fps = ?
  271. ''', (uuid, self.identifier, file_type, name, extension, size, created, path, frames,
  272. fps, file_type, name, extension, size, frames, fps))
  273. # lastrowid is 0 if on conflict clause applies.
  274. # If this is the case we do an extra query to receive the row id.
  275. if cursor.lastrowid > 0:
  276. row_id = cursor.lastrowid
  277. insert = True
  278. else:
  279. cursor.execute('SELECT id FROM files WHERE project = ? AND path = ?',
  280. (self.identifier, path))
  281. row_id = cursor.fetchone()[0]
  282. insert = False
  283. return self.file(row_id), insert