refine

modelscope · Dec 6, 2023 · d97bce1 · d97bce1
1 parent a858e3d
commit d97bce1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/data_juicer/ops/filter/face_area_filter.py b/data_juicer/ops/filter/face_area_filter.py
@@ -65,7 +65,8 @@ def compute_stats(self, sample, context=False):
 
         # there is no image in this sample, still default ratio 0.0
         if self.image_key not in sample or not sample[self.image_key]:
-            sample[Fields.stats][StatsKeys.face_ratios] = [0.0]
+            sample[Fields.stats][StatsKeys.face_ratios] = np.empty(0,
+                                                                   dtype=float)
             return sample
 
         # load images
@@ -108,7 +109,7 @@ def compute_stats(self, sample, context=False):
             img_area = images[key].width * images[key].height
             # Calculate the max face ratio for the current image
             max_face_ratios.append(
-                max([w * h / img_area for _, _, w, h in dets]))
+                max([w * h / img_area for _, _, w, h in dets], default=0.0))
         sample[Fields.stats][StatsKeys.face_ratios] = max_face_ratios
 
         return sample

diff --git a/tests/ops/filter/test_face_area_filter.py b/tests/ops/filter/test_face_area_filter.py
@@ -25,7 +25,7 @@ def _run_face_area_filter(self,
                                          column=[{}] * dataset.num_rows)
         dataset = dataset.map(op.compute_stats, num_proc=num_proc)
         dataset = dataset.filter(op.process, num_proc=num_proc)
-        dataset = dataset.remove_columns('__dj__stats__')
+        dataset = dataset.remove_columns(Fields.stats)
         res_list = dataset.to_list()
         self.assertEqual(res_list, target_list)