6
0

project.py 3.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. #!/usr/bin/env python
  2. """project: Encapsulates a project with all its properties."""
  3. import json
  4. import os
  5. from ..pipeline import Pipeline
  6. from ..utils import Errorable
  7. from ..utils import Video
  8. class Project(Errorable):
  9. def __init__(self, project_root, config=None):
  10. Errorable.__init__(self)
  11. self._update_ui_fn = lambda: None
  12. new = config is not None
  13. # (1) Check if directory exists
  14. self.root = project_root
  15. if not os.path.exists(self.root) and not new:
  16. self._report_error("Cannot open directory: %s" % self.root)
  17. return
  18. elif not os.path.exists(self.root) and new:
  19. try:
  20. os.mkdir(self.root)
  21. except:
  22. self._report_error("Cannot create directory: %s" % self.root)
  23. return
  24. # (2) Load project json
  25. if new:
  26. self._config = config
  27. else:
  28. try:
  29. with open(os.path.join(self.root, 'project.json'), 'r') as c:
  30. self._config = json.load(c)
  31. except:
  32. self._report_error("Cannot open configuration for project: %s" % self.root)
  33. return
  34. # (3) Load detection/extraction pipeline
  35. pipeline_config = self._config["pipeline"]
  36. self._pipeline = Pipeline(pipeline_config)
  37. self._err_children += [self._pipeline]
  38. def save(self):
  39. try:
  40. with open(os.path.join(self.root, 'project.json'), 'w') as c:
  41. json.dump(self._config, c, indent=2)
  42. except:
  43. self._report_error("Cannot save configuration for project: %s" % self.root)
  44. def execute(self, jobs, callback=lambda progress: True):
  45. all_subjobs = []
  46. for job in jobs:
  47. job['prediction'] = {}
  48. job['jobinfo'] = {}
  49. for subjob in job['jobs']:
  50. if job['filetype'] == 'image':
  51. pipeline_subjob = {'filetype': 'image',
  52. 'filename': job['filename'],
  53. 'prediction': job['prediction'],
  54. 'jobinfo': job['jobinfo'],
  55. 'subjob': subjob}
  56. all_subjobs += [pipeline_subjob]
  57. # TODO split video into frames
  58. elif job['filetype'] == 'video':
  59. cap = Video(job['filename'])
  60. if cap.framecount > 0:
  61. job['prediction-by-frame'] = [{} for frame in range(cap.framecount)]
  62. job['jobinfo-by-frame'] = [{} for frame in range(cap.framecount)]
  63. for frame in range(cap.framecount):
  64. pipeline_subjob = {'filetype': 'video',
  65. 'cap': cap,
  66. 'frame': frame,
  67. 'filename': job['filename'],
  68. 'prediction': job['prediction-by-frame'][frame],
  69. 'jobinfo': job['jobinfo-by-frame'][frame],
  70. 'subjob': subjob}
  71. all_subjobs += [pipeline_subjob]
  72. self._pipeline.execute(all_subjobs, callback)
  73. callback(1)
  74. def close(self):
  75. self.detector.close()
  76. self.features.close()
  77. def close(self):
  78. self._pipeline.close()