Adds Age Classifier and NSFW Classifier

app.py

@@ -136,6 +136,87 @@ def skintoneplot(hex_codes):
             ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
     return fig
 
+def age_detector(image):
+    """
+    A function that detects the age from an image.
+
+    Args:
+        image: The input image for age detection.
+
+    Returns:
+        str: The detected age label from the image.
+    """
+    pipe = pipeline('image-classification', model="dima806/faces_age_detection", device=0)
+    result = pipe(image)
+    max_score_item = max(result, key=lambda item: item['score'])
+    return max_score_item['label']
+
+def ageplot(agelist):
+    """
+    A function that plots age-related data based on the given list of age categories.
+
+    Args:
+        agelist (list): A list of age categories ("YOUNG", "MIDDLE", "OLD").
+
+    Returns:
+        fig: A matplotlib figure object representing the age plot.
+    """
+    order = ["YOUNG", "MIDDLE", "OLD"]
+    words = sorted(agelist, key=lambda x: order.index(x))
+    colors = {"YOUNG": "skyblue", "MIDDLE": "royalblue", "OLD": "darkblue"}
+    word_colors = [colors[word] for word in words]
+    fig, axes = plt.subplots(2, 5, figsize=(5,5))
+    plt.subplots_adjust(hspace=0.1, wspace=0.1)
+    for i, ax in enumerate(axes.flat):
+        ax.set_axis_off()
+        ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=word_colors[i]))
+    return fig
+
+def is_nsfw(image):
+    """
+    A function that checks if the input image is not safe for work (NSFW) by classifying it using 
+    an image classification pipeline and returning the label with the highest score.
+    
+    Args:
+        image: The input image to be classified.
+        
+    Returns:
+        str: The label of the NSFW category with the highest score.
+    """
+    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+    result = classifier(image)
+    max_score_item = max(result, key=lambda item: item['score'])
+    return max_score_item['label']
+
+def nsfwplot(nsfwlist):
+    """
+    Generates a plot of NSFW categories based on a list of NSFW labels.
+
+    Args:
+        nsfwlist (list): A list of NSFW labels ("normal" or "nsfw").
+
+    Returns:
+        fig: A matplotlib figure object representing the NSFW plot.
+
+    Raises:
+        None
+
+    This function takes a list of NSFW labels and generates a plot with a grid of 2 rows and 5 columns. 
+    Each label is sorted based on a predefined order and assigned a color. The plot is then created using matplotlib, 
+    with each cell representing an NSFW label. The color of each cell is determined by the corresponding label's color. 
+    The function returns the generated figure object.
+    """
+    order = ["normal", "nsfw"]
+    words = sorted(nsfwlist, key=lambda x: order.index(x))
+    colors = {"normal": "mistyrose", "nsfw": "red"}
+    word_colors = [colors[word] for word in words]
+    fig, axes = plt.subplots(2, 5, figsize=(5,5))
+    plt.subplots_adjust(hspace=0.1, wspace=0.1)
+    for i, ax in enumerate(axes.flat):
+        ax.set_axis_off()
+        ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=word_colors[i]))
+    return fig
+
 @spaces.GPU(duration=200)
 def generate_images_plots(prompt, model_name):
     global pipeline_text2image
@@ -145,6 +226,8 @@ def generate_images_plots(prompt, model_name):
     images = [getimgen(prompt, model_name) for _ in range(10)]
     genders = []
     skintones = []
+    ages = []
+    nsfws = []
     for image, i in zip(images, range(10)):
         prompt_prefix = "photo of a "
         caption = blip_caption_image(image, prefix=prompt_prefix)
@@ -156,7 +239,9 @@ def generate_images_plots(prompt, model_name):
         except:
             skintones.append(None)
         genders.append(genderfromcaption(caption))
-    return images, skintoneplot(skintones), genderplot(genders)
+        ages.append(age)
+        nsfws.append(nsfw)
+    return images, skintoneplot(skintones), genderplot(genders), ageplot(ages), nsfwplot(nsfws)
 
 with gr.Blocks(title="Skin Tone and Gender bias in Text-to-Image Generation Models") as demo:
     gr.Markdown("# Skin Tone and Gender bias in Text to Image Models")
@@ -166,7 +251,8 @@ In this demo, we explore the potential biases in text-to-image models by generat
 1. **Image Generation**: For each prompt, 10 images are generated using the selected model.
 2. **Gender Detection**: The [BLIP caption generator](https://huggingface.co/Salesforce/blip-image-captioning-large) is used to elicit gender markers by identifying words like "man," "boy," "woman," and "girl" in the captions.
 3. **Skin Tone Classification**: The [skin-tone-classifier library](https://github.com/ChenglongMa/SkinToneClassifier) is used to extract the skin tones of the generated subjects.
-
+4. **Age Detection**: The [Faces Age Detection model](https://huggingface.co/dima806/faces_age_detection) is used to identify the age of the generated subjects.
+5. **NSFW Detection**: The [Falconsai/nsfw_image_detection](https://huggingface.co/Falconsai/nsfw_image_detection) model is used to identify whether the generated images are NSFW (not safe for work).
 
 #### Visualization
 
@@ -174,7 +260,9 @@ We create visual grids to represent the data:
 
 - **Skin Tone Grids**: Skin tones are plotted as exact hex codes rather than using the Fitzpatrick scale, which can be [problematic and limiting for darker skin tones](https://arxiv.org/pdf/2309.05148).
 - **Gender Grids**: Light green denotes men, dark green denotes women, and grey denotes cases where the BLIP caption did not specify a binary gender.
-
+- **Age Grids**: Light blue denotes people between 18 and 30, blue denotes people between 30 and 50, and dark blue denotes people older than 50.
+- **NSFW Grids**: Light red denotes SFW images, and dark red denotes NSFW images.
+                
 This demo provides an insightful look into how current text-to-image models handle sensitive attributes, shedding light on areas for improvement and further study. 
 [Here is an article](https://medium.com/@evijit/analysis-of-ai-generated-images-of-indian-people-for-colorism-and-sexism-b80ff946759f) showing how this space can be used to perform such analyses, using colorism and sexism in India as an example.
 ''')
@@ -204,6 +292,9 @@ This demo provides an insightful look into how current text-to-image models hand
     with gr.Row(equal_height=True):
         skinplot = gr.Plot(label="Skin Tone")
         genplot = gr.Plot(label="Gender")
-    btn.click(generate_images_plots, inputs=[prompt, model_dropdown], outputs=[gallery, skinplot, genplot])
+    with gr.Row(equal_height=True):
+        agesplot = gr.Plot(label="Age")
+        nsfwsplot = gr.Plot(label="NSFW")
+    btn.click(generate_images_plots, inputs=[prompt, model_dropdown], outputs=[gallery, skinplot, genplot, agesplot, nsfwsplot])
 
 demo.launch(debug=True)