ExecuteExternalStorage.py 3.8 KB

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