ImageNet

app.py

@@ -1,27 +1,94 @@
+import json
 import logging
 import os
+from functools import partial
 
 import gradio as gr
+from datasets import Dataset, load_dataset
 from dotenv import load_dotenv
 
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
+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)
+
+
+def gen_from_iterable_dataset(iterable_ds):
+    """
+    Convert an iterable dataset to a generator
+    """
+    yield from iterable_ds
+
+
+# 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:
+# {
+#     "id": 1,
+#     "name": "entity",
+#     "children": ...
+# }
+with open("imagenet_categories_data.json") as f:
+    data = json.load(f)
+
+    # Recursively find all categories under "artifact, artefact".
+    # We want to get all the "index" values of the leaf nodes. Nodes that are not leaf nodes have a "children" key.
+    def find_categories(node):
+        if "children" in node:
+            for child in node["children"]:
+                yield from find_categories(child)
+        elif "index" in node:
+            yield node["index"]
+
+    broad_categories = data["children"]
+    artifact_category = next(
+        filter(lambda x: x["name"] == "artifact, artefact", broad_categories)
+    )
+    artifact_categories = list(find_categories(artifact_category))
+    logger.info(f"Artifact categories: {artifact_categories}")
+
+
+def filter_imgs_by_label(x):
+    """
+    Filter out the images that have label -1
+    """
+    print(f'label: {x["label"]}')
+    return x["label"] in artifact_categories
+
+
+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)
+
 
 def get_user_prompt():
+    # Pick the first 3 images and labels
+    images = []
+    machine_labels = []
+    human_labels = []
+    for i in range(3):
+        data = next(dataset_iterable)
+        logger.info(f"Data: {data}")
+        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_labels.append(human_label)
     return {
-        "images": [
-            "images/1.jpeg",
-            "images/1.jpeg",
-            "images/1.jpeg",
-        ],
-        "labels": [
-            "A roll of string",
-            "A camera",
-            "A loudspeaker",
-        ],
+        "images": images,
+        "machine_labels": machine_labels,
+        "human_labels": human_labels,
     }
 
 
@@ -56,16 +123,45 @@ with gr.Blocks(theme=theme) as demo:
     gr.Markdown(
         """We want to teach the Maker Faire Bot some creativity. Help us get ideas on what you'd build!"""
     )
+    image_components = []
     with gr.Row(variant="panel") as row:
         for i in range(len(user_prompt.value["images"])):
             with gr.Column(variant="default") as col:
-                gr.Image(
+                img = gr.Image(
                     user_prompt.value["images"][i],
-                    label=user_prompt.value["labels"][i],
+                    label=user_prompt.value["human_labels"][i],
                     interactive=False,
                     show_download_button=False,
                     show_share_button=False,
                 )
+                image_components.append(img)
+                btn = gr.Button("Change", variant="secondary")
+
+                def change_image(user_prompt):
+                    data = next(dataset_iterable)
+                    logger.info(user_prompt)
+                    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],
+                        gr.update(
+                            label=user_prompt["human_labels"][i],
+                        ),
+                    )
+
+                btn.click(
+                    change_image,
+                    inputs=[user_prompt],
+                    outputs=[user_prompt, img, img],
+                    preprocess=True,
+                    postprocess=True,
+                )
 
     user_answer_object = gr.Textbox(
         autofocus=True,
@@ -106,6 +202,13 @@ with gr.Blocks(theme=theme) as demo:
         preprocess=False,
     )
 
+    new_prompt_btn = gr.Button("New Prompt", variant="secondary")
+    new_prompt_btn.click(
+        get_user_prompt,
+        outputs=[user_prompt],
+        preprocess=False,
+    )
+
     gr.Markdown(
         """
         This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our [policy](https://makerfaire.com/privacy).

requirements.txt

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