WIP

app.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import random
 from functools import partial
 
 import gradio as gr
@@ -14,9 +15,11 @@ logger.setLevel(logging.INFO)
 load_dotenv()
 
 # dataset = load_dataset("detection-datasets/coco")
-it_dataset = load_dataset(
-    "imagenet-1k", split="train", streaming=True, trust_remote_code=True
-).shuffle(42)
+it_dataset = (
+    load_dataset("imagenet-1k", split="train", streaming=True, trust_remote_code=True)
+    .shuffle(42)
+    .take(1000)
+)
 
 
 def gen_from_iterable_dataset(iterable_ds):
@@ -26,6 +29,10 @@ def gen_from_iterable_dataset(iterable_ds):
     yield from iterable_ds
 
 
+dataset = Dataset.from_generator(
+    partial(gen_from_iterable_dataset, it_dataset), features=it_dataset.features
+)
+
 # imagenet_categories_data.json is a JSON file containing a hierarchy of ImageNet categories.
 # We want to take all categories under "artifact, artefact".
 # Each node has this structure:
@@ -58,14 +65,17 @@ def filter_imgs_by_label(x):
     """
     Filter out the images that have label -1
     """
+    logger.info(f'label: {x["label"]} (present: {x["label"] in artifact_categories})')
     return x["label"] in artifact_categories
 
 
-it_dataset = it_dataset.take(1000).filter(filter_imgs_by_label)
-dataset = Dataset.from_generator(
-    partial(gen_from_iterable_dataset, it_dataset), features=it_dataset.features
-)
-dataset_iterable = iter(dataset)
+dataset = dataset.filter(filter_imgs_by_label)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+load_dotenv()
 
 
 def get_user_prompt():
@@ -74,15 +84,11 @@ def get_user_prompt():
     machine_labels = []
     human_labels = []
     for i in range(3):
-        data = next(dataset_iterable)
-        logger.info(f"Data: {data}")
+        data = dataset[random.randint(0, len(dataset) - 1)]
         images.append(data["image"])
         # Get the label as a human readable string
         machine_labels.append(data["label"])
-        logger.info(dataset)
-        human_label = dataset.features["label"].int2str(data["label"]) + str(
-            data["label"]
-        )
+        human_label = dataset.features["label"].int2str(data["label"])
         human_labels.append(human_label)
     return {
         "images": images,
@@ -94,7 +100,7 @@ def get_user_prompt():
 hf_writer = gr.HuggingFaceDatasetSaver(
     hf_token=os.environ["HF_TOKEN"], dataset_name="acmc/maker-faire-bot", private=True
 )
-csv_writer = gr.CSVLogger(simplify_file_data=True)
+csv_writer = gr.CSVLogger()
 
 theme = gr.themes.Default(primary_hue="cyan", secondary_hue="fuchsia")
 
@@ -137,15 +143,14 @@ with gr.Blocks(theme=theme) as demo:
                 btn = gr.Button("Change", variant="secondary")
 
                 def change_image(user_prompt):
-                    data = next(dataset_iterable)
-                    logger.info(user_prompt)
+                    logger.info(f"Current user prompt: {user_prompt}")
+                    data = dataset[random.randint(0, len(dataset) - 1)]
                     user_prompt = user_prompt.copy()
                     user_prompt["images"][i] = data["image"]
                     user_prompt["machine_labels"][i] = data["label"]
                     user_prompt["human_labels"][i] = dataset.features["label"].int2str(
                         data["label"]
                     )
-                    logger.info(user_prompt)
                     return (
                         user_prompt,
                         user_prompt["images"][i],
@@ -192,20 +197,39 @@ with gr.Blocks(theme=theme) as demo:
     submit_btn = gr.Button("Submit", variant="primary")
 
     def log_results(prompt, object, explanation):
-        csv_writer.flag([prompt, object, explanation])
-        hf_writer.flag([prompt, object, explanation])
+        logger.info(f"logging - Prompt: {prompt}")
+        csv_writer.flag(
+            [
+                {
+                    "machine_labels": prompt["machine_labels"],
+                    "human_labels": prompt["human_labels"],
+                },
+                object,
+                explanation,
+            ]
+        )
+        hf_writer.flag(
+            [
+                {
+                    "machine_labels": prompt["machine_labels"],
+                    "human_labels": prompt["human_labels"],
+                },
+                object,
+                explanation,
+            ]
+        )
 
     submit_btn.click(
         log_results,
         inputs=[user_prompt, user_answer_object, user_answer_explanation],
-        preprocess=False,
+        preprocess=True,
     )
 
     new_prompt_btn = gr.Button("New Prompt", variant="secondary")
     new_prompt_btn.click(
         get_user_prompt,
         outputs=[user_prompt],
-        preprocess=False,
+        # preprocess=True,
     )
 
     gr.Markdown(

requirements.txt

@@ -1,3 +1,2 @@
 datasets==2.19.0
-gradio==4.28.0
 python-dotenv==1.0.1