diff --git a/arkindex/dataimport/models.py b/arkindex/dataimport/models.py index 1e7407bdab3e5cc3f5c10ae09db3e71d1aa2df8c..b2321c698a573b76e505aeefddcc5dcdbe80ccdb 100644 --- a/arkindex/dataimport/models.py +++ b/arkindex/dataimport/models.py @@ -75,13 +75,11 @@ class DataImport(IndexableModel): }, } - if self.mode == DataImportMode.Repository or self.ml_tools: - # Add the ML task if ML is required - for Git, let the Git import decide - tasks['ml'] = { - 'image': settings.ARKINDEX_ML_IMAGE, - 'command': 'arkindex-ml-analysis /data/ml_analysis.json', - 'parent': 'import', - } + for tool in self.ml_tools: + # Add the import task as parent of all ML tasks + task = tool.task_recipe + task['parent'] = 'import' + tasks['ml_{}'.format(tool.slug)] = task # Build the workflow in db recipe = settings.PONOS_RECIPE.copy() @@ -152,10 +150,10 @@ class DataImport(IndexableModel): tools = [] if 'classifier' in self.payload: tool = MLTool.get(settings.ML_CLASSIFIERS_DIR, MLToolType.Classifier, self.payload['classifier']) - tools.append((MLToolType.Classifier, tool.slug)) + tools.append(tool) if 'recognizer' in self.payload: tool = MLTool.get(settings.ML_CLASSIFIERS_DIR, MLToolType.Recognizer, self.payload['recognizer']) - tools.append((MLToolType.Recognizer, tool.slug)) + tools.append(tool) return tools diff --git a/arkindex/dataimport/tests/ml_tools/classifier/config.yml b/arkindex/dataimport/tests/ml_tools/classifier/config.yml index e1e9ddc0f315bae22bb525d813cd15a1d22fd2a1..b58ed1f7bd98ab80d92b17e5decef927fa42a046 100644 --- a/arkindex/dataimport/tests/ml_tools/classifier/config.yml +++ b/arkindex/dataimport/tests/ml_tools/classifier/config.yml @@ -3,6 +3,8 @@ name: Unit test classifier type: classifier slug: dummy_classifier version: 1.0.0 +docker: + image: some_image classes: - class_1 - class_2 diff --git a/arkindex/dataimport/tests/ml_tools/recognizer/config.yml b/arkindex/dataimport/tests/ml_tools/recognizer/config.yml index 83cb84d5877964be239d86d98772d814533b71f3..5ec28028d0ac258e772f0695d2bdfbf30f243b89 100644 --- a/arkindex/dataimport/tests/ml_tools/recognizer/config.yml +++ b/arkindex/dataimport/tests/ml_tools/recognizer/config.yml @@ -3,4 +3,6 @@ name: Unit test recognizer type: recognizer slug: dummy_recognizer version: 0.4.2 +docker: + image: some_image tesseract: [] diff --git a/arkindex/dataimport/tests/test_imports.py b/arkindex/dataimport/tests/test_imports.py index 5d42ea7c6279ec1d6ea3ec4e057c2c797d023071..6092cf4324259969b0212ef9911c0def275b37b0 100644 --- a/arkindex/dataimport/tests/test_imports.py +++ b/arkindex/dataimport/tests/test_imports.py @@ -137,6 +137,8 @@ class TestImports(FixtureAPITestCase): classifier_mock, recognizer_mock = MagicMock(), MagicMock() classifier_mock.slug = 'someclassifier' recognizer_mock.slug = 'somerecognizer' + classifier_mock.task_recipe = {'image': 'some_image1', 'command': 'some_command1'} + recognizer_mock.task_recipe = {'image': 'some_image2', 'command': 'some_command2'} ml_get_mock.side_effect = [classifier_mock, recognizer_mock] * 2 self.client.force_login(self.user) @@ -382,6 +384,8 @@ class TestImports(FixtureAPITestCase): classifier_mock, recognizer_mock = MagicMock(), MagicMock() classifier_mock.slug = 'someclassifier' recognizer_mock.slug = 'somerecognizer' + classifier_mock.task_recipe = {'image': 'some_image1', 'command': 'some_command1'} + recognizer_mock.task_recipe = {'image': 'some_image2', 'command': 'some_command2'} ml_get_mock.side_effect = [classifier_mock, recognizer_mock] * 3 self.client.force_login(self.user)