change group option behavior.

qfgaohao · Nov 1, 2018 · 2a222a9 · 2a222a9
1 parent dca078f
commit 2a222a9
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/open_images_downloader.py b/open_images_downloader.py
@@ -146,16 +146,15 @@ def parse_args():
         if not args.include_depiction:
             annotations = annotations.loc[annotations['IsDepiction'] != 1, :]
 
-        # TODO MAKE IT MORE EFFICIENT
-        #filter by IsGroupOf
         filtered = []
         for class_name, group_filter, percentage in zip(class_names, group_filters, percentages):
             sub = annotations.loc[annotations['ClassName'] == class_name, :]
-            if group_filter == "group":
-                sub = sub.loc[sub['IsGroupOf'] == 1, :]
-            elif group_filter == '~group':
-                sub = sub.loc[sub['IsGroupOf'] == 0, :]
             excluded_images |= set(sub['ImageID'].sample(frac=1 - percentage))
+
+            if group_filter == '~group':
+                excluded_images |= set(sub.loc[sub['IsGroupOf'] == 1, 'ImageID'])
+            elif group_filter == 'group':
+                excluded_images |= set(sub.loc[sub['IsGroupOf'] == 0, 'ImageID'])
             filtered.append(sub)
 
         annotations = pd.concat(filtered)