Refactor app.py and utils.py

app.py

@@ -1,7 +1,13 @@
 import glob
 import gradio as gr
 from huggingface_hub import get_token
-from utils import load_model, load_image_from_url, inference, load_badges
+from utils import (
+    load_model,
+    load_image_from_url,
+    inference,
+    load_badges,
+    count_flagged_images_from_csv,
+)
 from flagging import myHuggingFaceDatasetSaver
 
 
@@ -33,12 +39,19 @@ model.iou = 0.4
 model.max_det = 100
 model.agnostic = True  # NMS class-agnostic
 
-# This callback will be used to flag images
+# Flagging
 dataset_name = "SEA-AI/crowdsourced-sea-images"
 hf_writer = myHuggingFaceDatasetSaver(get_token(), dataset_name)
 
-with gr.Blocks(css=css) as demo:
-    badges = gr.HTML(load_badges(dataset_name, trials=1))
+
+def get_flagged_count():
+    """Count flagged images in dataset."""
+    return count_flagged_images_from_csv(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()))
     title = gr.HTML(TITLE)
 
     with gr.Row():
@@ -68,17 +81,22 @@ with gr.Blocks(css=css) as demo:
         cache_examples=True,
     )
 
-    # add components to clear
+    # add components to clear when clear button is clicked
     clear.add([img_input, img_url, img_output])
 
     # event listeners
     img_url.change(load_image_from_url, [img_url], img_input)
-    submit.click(lambda image: inference(model, image), [img_input], img_output)
+    submit.click(
+        lambda image: inference(model, image),
+        [img_input],
+        img_output,
+        api_name="inference",
+    )
 
     # event listeners with decorators
-    @img_output.change(inputs=[img_output], outputs=[flag, notice])
-    def show_hide(img_output):
-        visible = img_output is not None
+    @img_output.change(inputs=[img_output], outputs=[flag, notice], show_api=False)
+    def show_hide(_img_ouput):
+        visible = _img_ouput is not None
         return {
             flag: gr.Button("Flag", visible=visible, interactive=True),
             notice: gr.Markdown(value=NOTICE, visible=visible),
@@ -87,17 +105,21 @@ with gr.Blocks(css=css) as demo:
     # This needs to be called prior to the first call to callback.flag()
     hf_writer.setup([img_input], "flagged")
 
-    # We can choose which components to flag (in this case, we'll flag all)
-    flag.click(lambda: gr.Info("Thank you for contributing!")).then(
-        lambda: {flag: gr.Button("Flag", interactive=False)}, [], [flag]
+    # Sequential logic when flag button is clicked
+    flag.click(lambda: gr.Info("Thank you for contributing!"), show_api=False).then(
+        lambda: {flag: gr.Button("Flag", interactive=False)}, [], [flag], show_api=False
     ).then(
         lambda *args: hf_writer.flag(args),
         [img_input, flag],
         [],
         preprocess=False,
+        show_api=False,
     ).then(
-        lambda: load_badges(dataset_name), [], badges
+        lambda: load_badges(get_flagged_count()), [], badges, show_api=False
     )
 
+    # called during initial load in browser
+    demo.load(lambda: load_badges(get_flagged_count()), [], badges, show_api=False)
+
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()  # show_api=False)

utils.py

@@ -1,21 +1,24 @@
+import time
 import requests
 from io import BytesIO
 import numpy as np
+import pandas as pd
 from PIL import Image
 import yolov5
 from yolov5.utils.plots import Annotator, colors
 import gradio as gr
 from huggingface_hub import get_token
-import time
 
 
 def load_model(model_path, img_size=640):
+    """Load model from HuggingFace Hub."""
     model = yolov5.load(model_path, hf_token=get_token())
     model.img_size = img_size  # add img_size attribute
     return model
 
 
 def load_image_from_url(url):
+    """Load image from URL."""
     if not url:  # empty or None
         return gr.Image(interactive=True)
     try:
@@ -27,6 +30,7 @@ def load_image_from_url(url):
 
 
 def inference(model, image):
+    """Run inference on image and return annotated image."""
     results = model(image, size=model.img_size)
     annotator = Annotator(np.asarray(image))
     for *box, _, cls in reversed(results.pred[0]):
@@ -36,7 +40,9 @@ def inference(model, image):
     return annotator.im
 
 
-def count_flagged_images(dataset_name, trials=10):
+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}"
 
@@ -58,8 +64,15 @@ def count_flagged_images(dataset_name, trials=10):
     return 0
 
 
-def load_badges(dataset_name, trials=10):
-    n = count_flagged_images(dataset_name, trials)
+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"""
         <p style="display: flex">
         <img alt="" src="https://img.shields.io/badge/SEA.AI-beta-blue">