make admin.py create_dataset more compatible to question_answering task

nusdbsystem · May 29, 2020 · 2f2c10c · 2f2c10c
1 parent cda68e1
commit 2f2c10c
Showing 1 changed file with 23 additions and 14 deletions.
diff --git a/singa_auto/admin/admin.py b/singa_auto/admin/admin.py
@@ -229,6 +229,22 @@ def create_dataset(self, user_id, name, task, data_file_path):
                         num_unlabeled_samples = num_samples - num_labeled_samples
                         break
             else:
+                if task == 'question_answering':
+                    stat = {}
+                    dataset = self._meta_store.create_dataset(name, task, size_bytes,
+                                                              store_dataset_id, user_id,
+                                                              stat)
+                    self._meta_store.commit()
+
+                    return {
+                        'id': dataset.id,
+                        'name': dataset.name,
+                        'task': dataset.task,
+                        'size_bytes': dataset.size_bytes,
+                        'store_dataset_id': dataset.store_dataset_id,
+                        'owner_id': dataset.owner_id,
+                        'stat': dataset.stat,
+                    }
                 # if the csv file was not provided in zip
                 with tempfile.TemporaryDirectory() as dir_path:
                     num_labeled_samples = len(dataset_zipfile.namelist())
@@ -250,21 +266,14 @@ def create_dataset(self, user_id, name, task, data_file_path):
                         img = Image.open(img_path)
                         img_size = str(img.size)
 
+            stat = {
+                'num_labeled_samples': num_labeled_samples,
+                'num_unlabeled_samples': num_unlabeled_samples,
+                'class_count': class_count.to_json(),
+                'ratio': ratio.to_json()
+            }
             if task == 'IMAGE_CLASSIFICATION':
-                stat = {
-                    'num_labeled_samples': num_labeled_samples,
-                    'num_unlabeled_samples': num_unlabeled_samples,
-                    'class_count': class_count.to_json(),
-                    'ratio': ratio.to_json(),
-                    'img_size': img_size
-                }
-            else:
-                stat = {
-                    'num_labeled_samples': num_labeled_samples,
-                    'num_unlabeled_samples': num_unlabeled_samples,
-                    'class_count': class_count.to_json(),
-                    'ratio': ratio.to_json()
-                }
+                stat['img_size'] = img_size
 
         print('begin saving to db')