image mod tested

app.py

@@ -7,6 +7,8 @@ from rdkit.Chem import Draw
 import time
 import os
 import random
+from mod import check_img
+import string
 
 @st.dialog("Select From Library")
 def lib_modal():
@@ -49,10 +51,12 @@ mode = st.radio("Select file source:", ["Upload File", "Choose from Our Demo Lib
 if mode == "Upload File":
     st.session_state.pop("uploaded_file", None)
     uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "webp", "heic"])
-    col1.checkbox("Contribute To Public Library", value=True, help="If checked, images will be included in the PUBLIC library, and the image will be reviewed by our team and used for model training. When checked, do not upload any sensitive or personal data.")
+    contribute = col1.checkbox("Contribute To Public Library", value=True, help="If checked, images will be included in the PUBLIC library, and the image will be reviewed by our team and used for model training. When checked, do not upload any sensitive or personal data.")
 else:
+    contribute = False
     uploaded_file = None
-    st.warning("You are accessing a user-uploaded library. These images are unverified, unmoderated, and may include inappropriate content or incorrectly assembled files. Proceed at your own risk. For inquiries, contact us.")
+    st.warning("You are accessing a user-uploaded library. These images are unverified and may contain inappropriate content or incorrectly assembled files, despite some moderation. Proceed with caution. For inquiries, contact us.")
+
     from_library = st.button("Select from Library")
     if from_library:
         if st.session_state.show_modal:
@@ -88,5 +92,14 @@ if button:
         pubchem_url = "https://pubchem.ncbi.nlm.nih.gov/compound/{}".format(cid[0].cid)
         col.markdown("[PubChem]({})".format(pubchem_url))     
 
+    if contribute:
+        flagged = check_img(image)
+        if flagged:
+            st.warning("The image is flagged as inappropriate. Please upload a different image.")
+            st.stop()
+        else:
+            filename = "user_upload" + "".join(random.choices(string.ascii_uppercase, k=10)) + ".png"
+            image.save("lib/{}".format(filename))
+            st.success("The image is stored in the library.")
     st.markdown("---")
     st.markdown("Taken {} seconds".format(round(time.time() - start_time, 2)))

mod.py

@@ -0,0 +1,34 @@
+from openai import OpenAI
+import io
+import base64
+from PIL import Image
+
+client = OpenAI()
+
+def check_img(PIL_image):
+    buffer = io.BytesIO()
+    PIL_image.save(buffer, format="PNG")
+    buffer.seek(0)
+    img_bytes = buffer.read()
+    base64_bytes = base64.b64encode(img_bytes)
+    base64_string = base64_bytes.decode('utf-8')
+    base64_string = "data:image/png;base64," + base64_string
+
+    response = client.moderations.create(
+        model="omni-moderation-latest",
+        input=[
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": base64_string,
+                }
+            },
+        ],
+    )
+
+    return response.results[0].flagged
+
+if __name__ == "__main__":
+    image = Image.open("lib/11_0M0DY5zO.png")
+    flagged = check_img(image)
+    print("Flagged:", flagged)