ExecuteExternalStorage.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. import os
  2. import uuid
  3. from flask import abort
  4. from flask import make_response
  5. from flask import request
  6. from flask.views import View
  7. from pycs import db
  8. from pycs.database.Project import Project
  9. from pycs.frontend.notifications.NotificationManager import NotificationManager
  10. from pycs.jobs.JobGroupBusyException import JobGroupBusyException
  11. from pycs.jobs.JobRunner import JobRunner
  12. from pycs.util.FileOperations import file_info
  13. from pycs.util.FileOperations import find_images
  14. class ExecuteExternalStorage(View):
  15. """
  16. find media files stored in a projects data_folder
  17. """
  18. # pylint: disable=arguments-differ
  19. methods = ['POST']
  20. def __init__(self, nm: NotificationManager, jobs: JobRunner):
  21. # pylint: disable=invalid-name
  22. self.nm = nm
  23. self.jobs = jobs
  24. def dispatch_request(self, user: str, project_id: int):
  25. # pylint: disable=unused-argument
  26. # extract request data
  27. data = request.get_json(force=True)
  28. if not data.get('execute', False):
  29. return abort(400)
  30. # find project
  31. project = Project.get_or_404(project_id)
  32. if not project.external_data:
  33. return abort(400, "External data is not set!")
  34. # execute label provider and add labels to project
  35. try:
  36. ExecuteExternalStorage.find_media_files(self.nm, self.jobs, project)
  37. except JobGroupBusyException:
  38. return abort(400, "Job is already running!")
  39. return make_response()
  40. @staticmethod
  41. def find_media_files(nm: NotificationManager, jobs: JobRunner, project: Project):
  42. """
  43. start a job that finds media files in the projects data_folder and adds them to the
  44. database afterwards
  45. :param nm: notification manager object
  46. :param jobs: job runner object
  47. :param project: project
  48. :return:
  49. """
  50. # pylint: disable=invalid-name
  51. # find lists the given data folder and prepares item dictionaries
  52. def find(data_folder):
  53. image_paths = find_images(data_folder)
  54. length = len(image_paths)
  55. elements = []
  56. current = 0
  57. for file_path in image_paths:
  58. file_folder = str(file_path.parent)
  59. file_name = file_path.stem
  60. file_extension = file_path.suffix
  61. file_size = os.path.getsize(str(file_path))
  62. try:
  63. ftype, frames, fps = file_info(file_folder, file_name, file_extension)
  64. except ValueError:
  65. continue
  66. file_uuid = str(uuid.uuid1())
  67. file_attrs = dict(
  68. uuid=file_uuid,
  69. file_type=ftype,
  70. name=file_name,
  71. extension=file_extension,
  72. size=file_size,
  73. filename=file_path.relative_to(data_folder).with_name(file_name),
  74. frames=frames,
  75. fps=fps)
  76. elements.append(file_attrs)
  77. current += 1
  78. if len(elements) >= 200:
  79. yield elements, current, length
  80. elements.clear()
  81. if len(elements) > 0:
  82. yield elements, current, length
  83. # progress inserts elements into the database and fires events
  84. def progress(elements, current, length):
  85. with db.session.begin_nested():
  86. for file_attrs in elements:
  87. file, is_new = project.add_file(commit=False, **file_attrs)
  88. if is_new:
  89. nm.create_file(file)
  90. return current / length
  91. # run job with given functions
  92. jobs.run(project,
  93. 'Find Media Files',
  94. project.name,
  95. f'{project.id}/find-files',
  96. find, project.data_folder,
  97. progress=progress)