Create app.py

app.py

@@ -0,0 +1,67 @@
+from transformers import AutoProcessor, ShieldGemma2ForImageClassification
+from PIL import Image
+import requests
+import torch
+import gradio as gr
+import spaces
+model_id = "google/shieldgemma-2-4b-it"
+model = ShieldGemma2ForImageClassification.from_pretrained(model_id).to("cuda")
+processor = AutoProcessor.from_pretrained(model_id)
+
+@spaces.GPU()
+def infer(image, policies, policy_descriptions):
+  policies = policies.split(";")
+  policy_descriptions = policy_descriptions.split(";")
+  custom_policies = dict(zip(policies, policy_descriptions))
+  print(custom_policies)
+
+  inputs = processor(
+      images=[image],
+      custom_policies=custom_policies,
+      policies=policies,
+      return_tensors="pt",
+  ).to(model.device)
+  
+  with torch.inference_mode():
+      output = model(**inputs)
+      print(output.probabilities)
+  
+  outs = {}
+  for idx, policy in enumerate(output.probabilities):
+    yes_prob = policy[0]
+    no_prob = policy[1]
+
+    outs[f"Yes for {policies[idx]}"] = yes_prob
+    outs[f"No for {policies[idx]}"] = no_prob
+    print(outs)
+
+  return outs
+
+
+IMG = """
+  <img src='https://storage.googleapis.com/gweb-developer-goog-blog-assets/images/SheildGemma2_WagtailBlog_RD1_V01a.original.png' />
+  """
+
+
+with gr.Blocks() as demo:
+  gr.Markdown(IMG)
+  gr.Markdown("## ShieldGemma2 for Multimodal Safety")
+  gr.Markdown("ShieldGemma2 is a safety moderation model for vision language models. It can detect unsafe images. To use it, simply input an image, and provide policies. A policy is a description of what should be detected, and also provide names of policies.")
+  gr.Markdown("You can test it with an image of your choice and example policies provided below.")
+ 
+  with gr.Row():
+    with gr.Column():
+      image = gr.Image(type="pil")
+      policies = gr.Textbox(label = "Name of policies separated by semicolon")
+      policy_descriptions = gr.Textbox(label = "Description of policies separated by semicolon", lines=10)
+      btn = gr.Button("Submit")
+    with gr.Column():
+      outputs = gr.Label()
+    
+    btn.click(fn=infer, inputs=[image, policies, policy_descriptions], outputs=outputs)
+  gr.Examples(examples = [["Sexually Explicit content;Dangerous content;Violent content", "The image shall not contain content that depicts explicit or graphic sexual acts.; The image shall not contain content that facilitates or encourages activities that could cause real-world harm (e.g., building firearms and explosive devices, promotion of terrorism, instructions for suicide).;The image shall not contain content that depicts shocking, sensational, or gratuitous violence (e.g., excessive blood and gore, gratuitous violence against animals, extreme injury or moment of death)."]],
+              inputs = [policies, policy_descriptions])
+
+
+demo.launch(debug=True)
+