Refactor flagged image counting logic

app.py

@@ -6,7 +6,7 @@ from utils import (
     load_image_from_url,
     inference,
     load_badges,
-    count_flagged_images_from_csv,
+    FlaggedCounter,
 )
 from flagging import myHuggingFaceDatasetSaver
 
@@ -42,16 +42,12 @@ model.agnostic = True  # NMS class-agnostic
 # Flagging
 dataset_name = "SEA-AI/crowdsourced-sea-images"
 hf_writer = myHuggingFaceDatasetSaver(get_token(), dataset_name)
-
-
-def get_flagged_count():
-    """Count flagged images in dataset."""
-    return count_flagged_images_from_csv(dataset_name)
+flagged_counter = FlaggedCounter(dataset_name)
 
 
 theme = gr.themes.Default(primary_hue=gr.themes.colors.indigo)
 with gr.Blocks(theme=theme, css=css) as demo:
-    badges = gr.HTML(load_badges(get_flagged_count()))
+    badges = gr.HTML(load_badges(flagged_counter.count()))
     title = gr.HTML(TITLE)
 
     with gr.Row():
@@ -115,11 +111,11 @@ with gr.Blocks(theme=theme, css=css) as demo:
         preprocess=False,
         show_api=False,
     ).then(
-        lambda: load_badges(get_flagged_count()), [], badges, show_api=False
+        lambda: load_badges(flagged_counter.count()), [], badges, show_api=False
     )
 
     # called during initial load in browser
-    demo.load(lambda: load_badges(get_flagged_count()), [], badges, show_api=False)
+    demo.load(lambda: load_badges(flagged_counter.count()), [], badges, show_api=False)
 
 if __name__ == "__main__":
     demo.queue().launch()  # show_api=False)

utils.py

@@ -1,6 +1,7 @@
 import time
 import requests
 from io import BytesIO
+from dataclasses import dataclass
 import numpy as np
 import pandas as pd
 from PIL import Image
@@ -40,37 +41,6 @@ def inference(model, image):
     return annotator.im
 
 
-def count_flagged_images_via_api(dataset_name, trials=10):
-    """Count flagged images via API. Might be slow."""
-
-    headers = {"Authorization": f"Bearer {get_token()}"}
-    API_URL = f"https://datasets-server.huggingface.co/size?dataset={dataset_name}"
-
-    def query():
-        response = requests.get(API_URL, headers=headers, timeout=5)
-        return response.json()
-
-    for i in range(trials):
-        try:
-            data = query()
-            if "error" not in data and data["size"]["dataset"]["num_rows"] > 0:
-                print(f"[{i+1}/{trials}] {data}")
-                return data["size"]["dataset"]["num_rows"]
-        except Exception:
-            pass
-        print(f"[{i+1}/{trials}] {data}")
-        time.sleep(5)
-
-    return 0
-
-
-def count_flagged_images_from_csv(dataset_name):
-    """Count flagged images from CSV. Fast but relies on local files."""
-    dataset_name = dataset_name.split("/")[-1]
-    df = pd.read_csv(f"./flagged/{dataset_name}/data.csv")
-    return len(df)
-
-
 def load_badges(n):
     """Load badges."""
     return f"""
@@ -80,3 +50,52 @@ def load_badges(n):
         <img alt="" src="https://img.shields.io/badge/%F0%9F%96%BC%EF%B8%8F-{n}-green">
         </p>
         """
+
+
+@dataclass
+class FlaggedCounter:
+    """Count flagged images in dataset."""
+
+    dataset_name: str
+    headers: dict = None
+
+    def __post_init__(self):
+        self.API_URL = (
+            f"https://datasets-server.huggingface.co/size?dataset={self.dataset_name}"
+        )
+        self.trials = 10
+        if self.headers is None:
+            self.headers = {"Authorization": f"Bearer {get_token()}"}
+
+    def query(self):
+        """Query API."""
+        response = requests.get(self.API_URL, headers=self.headers, timeout=5)
+        return response.json()
+
+    def from_query(self, data):
+        """Count flagged images via API. Might be slow."""
+        for i in range(self.trials):
+            try:
+                data = self.query()
+                if "error" not in data and data["size"]["dataset"]["num_rows"] > 0:
+                    print(f"[{i+1}/{self.trials}] {data}")
+                    return data["size"]["dataset"]["num_rows"]
+            except Exception:
+                pass
+            print(f"[{i+1}/{self.trials}] {data}")
+            time.sleep(5)
+
+        return 0
+
+    def from_csv(self):
+        """Count flagged images from CSV. Fast but relies on local files."""
+        dataset_name = self.dataset_name.split("/")[-1]
+        df = pd.read_csv(f"./flagged/{dataset_name}/data.csv")
+        return len(df)
+
+    def count(self):
+        """Count flagged images."""
+        try:
+            return self.from_csv()
+        except FileNotFoundError:
+            return self.from_query(self.query())