simplify UI

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.JPEG filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -9,4 +9,50 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Generative Data Augmentation Demo
+
+Main GitHub Repo: [Generative Data Augmentation](https://github.com/zhulinchng/generative-data-augmentation) | Image Classification Demo: [Generative Augmented Classifiers](https://huggingface.co/spaces/czl/generative-augmented-classifiers).
+
+This demo is created as part of the 'Investigating the Effectiveness of Generative Diffusion Models in Synthesizing Images for Data Augmentation in Image Classification' dissertation.
+
+The user can augment an image by interpolating between two prompts, and specify the number of interpolation steps and the specific step to generate the image.
+
+## Demo Usage Instructions
+
+1. Upload an image.
+2. Enter the two prompts to interpolate between, the first prompt should contain the desired class of the augmented image, the second prompt should contain the undesired class (i.e., confusing class).
+
+## Configuration
+
+- Total Interpolation Steps: The number of steps to interpolate between the two prompts.
+- Interpolation Step: The specific step to generate the image.
+- Example for 10 steps:
+
+    ```python
+        Total:  [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        Sampled:          4
+    ```
+
+- Seed: Seed value for reproducibility.
+- Negative Prompt: Prompt to guide the model away from generating the image.
+- Width, Height: The dimensions of the generated image.
+- Guidance Scale: The scale of the guide the model on how closely to follow the prompts.
+
+## Metadata
+
+[SSIM Score](https://lightning.ai/docs/torchmetrics/stable/image/structural_similarity.html): Structural Similarity Index (SSIM) score between the original and generated image, ranges from 0 to 1.
+[CLIP Score](https://lightning.ai/docs/torchmetrics/stable/multimodal/clip_score.html): CLIP similarity score between the original and generated image, ranges from 0 to 100.
+
+## Local Setup
+
+```bash
+git clone https://huggingface.co/spaces/czl/generative-data-augmentation-demo
+cd generative-data-augmentation-demo
+# Setup the data directory structure as shown above
+conda create --name $env_name python=3.11.* # Replace $env_name with your environment name
+conda activate $env_name
+# Visit PyTorch website https://pytorch.org/get-started/previous-versions/#v212 for PyTorch installation instructions.
+pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url # Obtain the correct URL from the PyTorch website
+pip install -r requirements.txt
+python app.py
+```

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import numpy as np
 import torch
 import torchvision.transforms as transforms
+from PIL import Image
 from torchmetrics.functional.image import structural_similarity_index_measure as ssim
 from transformers import CLIPModel, CLIPProcessor
 
@@ -45,8 +46,6 @@ def infer(
     interpolation_step,
     num_inference_steps,
     num_interpolation_steps,
-    sample_mid_interpolation,
-    remove_n_middle,
 ):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -55,38 +54,15 @@ def infer(
         assert num_interpolation_steps % 2 == 0
     except AssertionError:
         raise ValueError("num_interpolation_steps must be an even number")
-    try:
-        assert sample_mid_interpolation % 2 == 0
-    except AssertionError:
-        raise ValueError("sample_mid_interpolation must be an even number")
-    try:
-        assert remove_n_middle % 2 == 0
-    except AssertionError:
-        raise ValueError("remove_n_middle must be an even number")
-    try:
-        assert num_interpolation_steps >= sample_mid_interpolation
-    except AssertionError:
-        raise ValueError(
-            "num_interpolation_steps must be greater than or equal to sample_mid_interpolation"
-        )
-    try:
-        assert num_interpolation_steps >= 2 and sample_mid_interpolation >= 2
-    except AssertionError:
-        raise ValueError(
-            "num_interpolation_steps and sample_mid_interpolation must be greater than or equal to 2"
-        )
-    try:
-        assert sample_mid_interpolation - remove_n_middle >= 2
-    except AssertionError:
-        raise ValueError(
-            "sample_mid_interpolation must be greater than or equal to remove_n_middle + 2"
-        )
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     prompts = [prompt1, prompt2]
     generator = torch.Generator().manual_seed(seed)
 
+    sample_mid_interpolation = num_interpolation_steps
+    remove_n_middle = 0
+
     interpolated_prompt_embeds, prompt_metadata = synth.interpolatePrompts(
         prompts,
         pipe,
@@ -116,7 +92,6 @@ def infer(
     ).to(device)
     embed_pairs = zip(interpolated_prompt_embeds, interpolated_negative_prompts_embeds)
     embed_pairs_list = list(embed_pairs)
-    print(len(embed_pairs_list))
     # offset step by -1
     prompt_embeds, negative_prompt_embeds = embed_pairs_list[interpolation_step - 1]
     preprocess_input = transforms.Compose(
@@ -127,7 +102,7 @@ def infer(
         npe = negative_prompt_embeds[None, ...]
     else:
         npe = None
-    image = pipe(
+    images_list = pipe(
         height=height,
         width=width,
         num_images_per_prompt=1,
@@ -138,7 +113,13 @@ def infer(
         generator=generator,
         latents=latents,
         image=input_img_tensor,
-    ).images[0]
+    )
+    if images_list["nsfw_content_detected"][0]:
+        image = Image.open("samples/unsafe.jpeg")
+        return image, seed, "Unsafe content detected", "Unsafe content detected"
+    else:
+        image = images_list.images[0]
+
     pred_image = transforms.ToTensor()(image).unsqueeze(0)
     ssim_score = ssim(pred_image, input_img_tensor).item()
     real_inputs = clip_processor(
@@ -163,25 +144,17 @@ def infer(
 
 
 examples1 = [
-    "A photo of a chain saw, chainsaw",
+    "A photo of a garbage truck, dustcart",
     "A photo of a Shih-Tzu, a type of dog",
 ]
 examples2 = [
-    "A photo of a golf ball",
+    "A photo of a cassette player",
     "A photo of a beagle, a type of dog",
 ]
 
 
 def update_steps(total_steps, interpolation_step):
-    if interpolation_step > total_steps:
-        return gr.update(maximum=total_steps // 2, value=total_steps)
-    return gr.update(maximum=total_steps // 2)
-
-
-def update_sampling_steps(total_steps, sample_steps):
-    # if sample_steps > total_steps:
-    #     return gr.update(value=total_steps)
-    return gr.update(value=total_steps)
+    return gr.update(maximum=total_steps)
 
 
 def update_format(image_format):
@@ -211,7 +184,7 @@ with gr.Blocks(title="Generative Date Augmentation Demo") as demo:
                     label="Prompt for the image to synthesize. (Actual class)",
                     show_label=True,
                     max_lines=1,
-                    placeholder="Enter your first prompt",
+                    placeholder="Enter Prompt for the image to synthesize. (Actual class)",
                     container=False,
                 )
             with gr.Row():
@@ -219,32 +192,44 @@ with gr.Blocks(title="Generative Date Augmentation Demo") as demo:
                     label="Prompt to augment against. (Confusing class)",
                     show_label=True,
                     max_lines=1,
-                    placeholder="Enter your second prompt",
+                    placeholder="Enter Prompt to augment against. (Confusing class)",
                     container=False,
                 )
             with gr.Row():
                 gr.Examples(
-                    examples=examples1, inputs=[prompt1], label="Example for Prompt 1"
+                    examples=[
+                        "samples/n03417042_5234.JPEG",
+                        "samples/n02086240_2799.JPEG",
+                    ],
+                    inputs=[input_image],
+                    label="Example Images",
+                )
+                gr.Examples(
+                    examples=examples1,
+                    inputs=[prompt1],
+                    label="Example for Prompt 1 (Actual class)",
                 )
                 gr.Examples(
-                    examples=examples2, inputs=[prompt2], label="Example for Prompt 2"
+                    examples=examples2,
+                    inputs=[prompt2],
+                    label="Example for Prompt 2 (Confusing class)",
                 )
 
             with gr.Row():
-                interpolation_step = gr.Slider(
-                    label="Specific Interpolation Step",
-                    minimum=1,
-                    maximum=8,
-                    step=1,
-                    value=8,
-                )
                 num_interpolation_steps = gr.Slider(
-                    label="Total interpolation steps",
+                    label="Total Interpolation Steps",
                     minimum=2,
-                    maximum=32,
+                    maximum=128,
                     step=2,
                     value=16,
                 )
+                interpolation_step = gr.Slider(
+                    label="Sample Interpolation Step",
+                    minimum=1,
+                    maximum=16,
+                    step=1,
+                    value=8,
+                )
                 num_interpolation_steps.change(
                     fn=update_steps,
                     inputs=[num_interpolation_steps, interpolation_step],
@@ -305,27 +290,6 @@ with gr.Blocks(title="Generative Date Augmentation Demo") as demo:
                         step=1,
                         value=25,
                     )
-                with gr.Row():
-                    sample_mid_interpolation = gr.Slider(
-                        label="Number of sampling steps in the middle of interpolation",
-                        minimum=2,
-                        maximum=80,
-                        step=2,
-                        value=16,
-                    )
-                    num_interpolation_steps.change(
-                        fn=update_sampling_steps,
-                        inputs=[num_interpolation_steps, sample_mid_interpolation],
-                        outputs=[sample_mid_interpolation],
-                    )
-                with gr.Row():
-                    remove_n_middle = gr.Slider(
-                        label="Number of middle steps to remove from interpolation",
-                        minimum=0,
-                        maximum=80,
-                        step=2,
-                        value=0,
-                    )
                 with gr.Row():
                     image_type = gr.Radio(
                         choices=[
@@ -372,6 +336,10 @@ Note: Running on CPU will take longer (approx. 6 minutes with default settings).
 This demo is created as part of the 'Investigating the Effectiveness of Generative Diffusion Models in Synthesizing Images for Data Augmentation in Image Classification' dissertation.
 
 The user can augment an image by interpolating between two prompts, and specify the number of interpolation steps and the specific step to generate the image.
+
+View the files used in this demo [here](https://huggingface.co/spaces/czl/generative-data-augmentation-demo/tree/main).
+
+Note: Safety checker is enabled to prevent unsafe content from being displayed in this public demo.
             """
             )
         run_button.click(
@@ -389,27 +357,8 @@ The user can augment an image by interpolating between two prompts, and specify
                 interpolation_step,
                 num_inference_steps,
                 num_interpolation_steps,
-                sample_mid_interpolation,
-                remove_n_middle,
             ],
             outputs=[result, show_seed, ssim_score, cos_sim],
         )
 
 demo.queue().launch(show_error=True)
-
-"""
-    input_image,
-    prompt1,
-    prompt2,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    interpolation_step,
-    num_inference_steps,
-    num_interpolation_steps,
-    sample_mid_interpolation,
-    remove_n_middle,
-    """

tools/synth.py

@@ -157,7 +157,6 @@ def pipe_img(
         scheduler=scheduler,
         torch_dtype=torch.float32,
         use_safetensors=use_safetensors,
-        safety_checker=None,
     ).to(device)
     if cpu_offload:
         pipe.enable_model_cpu_offload()