fix safety check

app.py

@@ -262,12 +262,11 @@ def handle_image_upload_and_gallery_hf(uploaded_files):
 
         ind_small = 0
         for uploaded_file in uploaded_files:
-
+            uploaded_file.seek(0)
             if SAFE.check_for_inappropriate_content(uploaded_file):
                 clear_image_uploads()
                 report_violation(uploaded_file.name, is_hf=st.session_state['is_hf'])
                 st.error("Warning: You uploaded an image that violates our terms of service.")
-                return True
 
             
             # Determine the file type
@@ -391,7 +390,7 @@ def content_input_images(col_left, col_right):
     
     with col_right:
         if st.session_state.is_hf:
-            result = handle_image_upload_and_gallery_hf(uploaded_files)
+            handle_image_upload_and_gallery_hf(uploaded_files)
 
         else:
             st.session_state['view_local_gallery'] = st.toggle("View Image Gallery",)

vouchervision/OCR_google_cloud_vision.py

@@ -792,58 +792,6 @@ class OCREngine:
         except:
             pass
 
-# class SafetyCheck():
-#     def __init__(self, is_hf) -> None:
-#         self.is_hf = is_hf
-#         self.set_client()
-
-#     def set_client(self):
-#         if self.is_hf:
-#             self.client = vision.ImageAnnotatorClient(credentials=self.get_google_credentials())
-#         else:
-#             self.client = vision.ImageAnnotatorClient(credentials=self.get_google_credentials())
-
-
-#     def get_google_credentials(self):
-#         creds_json_str = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')
-#         credentials = service_account.Credentials.from_service_account_info(json.loads(creds_json_str))
-#         return credentials
-    
-#     def check_for_inappropriate_content(self, file_stream):
-#         LEVEL = 2
-#         content = file_stream.read()
-#         image = vision.Image(content=content)
-#         response = self.client.safe_search_detection(image=image)
-#         safe = response.safe_search_annotation
-
-#         likelihood_name = (
-#             "UNKNOWN",
-#             "VERY_UNLIKELY",
-#             "UNLIKELY",
-#             "POSSIBLE",
-#             "LIKELY",
-#             "VERY_LIKELY",
-#         )
-#         print("Safe search:")
-
-#         print(f"    adult*: {likelihood_name[safe.adult]}")
-#         print(f"    medical*: {likelihood_name[safe.medical]}")
-#         print(f"    spoofed: {likelihood_name[safe.spoof]}")
-#         print(f"    violence*: {likelihood_name[safe.violence]}")
-#         print(f"    racy: {likelihood_name[safe.racy]}")
-
-#         # Check the levels of adult, violence, racy, etc. content.
-#         if (safe.adult > LEVEL or
-#             safe.medical > LEVEL or
-#             # safe.spoof > LEVEL or
-#             safe.violence > LEVEL #or
-#             # safe.racy > LEVEL
-#             ):
-#             print("Found violation")
-#             return True  # The image violates safe search guidelines.
-
-#         print("Found NO violation")
-#         return False  # The image is considered safe.
 class SafetyCheck():
     def __init__(self, is_hf) -> None:
         self.is_hf = is_hf
@@ -859,15 +807,14 @@ class SafetyCheck():
         creds_json_str = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')
         credentials = service_account.Credentials.from_service_account_info(json.loads(creds_json_str))
         return credentials
-
+    
     def check_for_inappropriate_content(self, file_stream):
         LEVEL = 2
+        # content = file_stream.read()
+        file_stream.seek(0)  # Reset file stream position to the beginning
         content = file_stream.read()
         image = vision.Image(content=content)
-        feature = vision.Feature(type=vision.Feature.Type.SAFE_SEARCH_DETECTION)
-
-        request = vision.AnnotateImageRequest(image=image, features=[feature])
-        response = self.client.annotate_image(request=request)
+        response = self.client.safe_search_detection(image=image)
         safe = response.safe_search_annotation
 
         likelihood_name = (
@@ -879,6 +826,7 @@ class SafetyCheck():
             "VERY_LIKELY",
         )
         print("Safe search:")
+
         print(f"    adult*: {likelihood_name[safe.adult]}")
         print(f"    medical*: {likelihood_name[safe.medical]}")
         print(f"    spoofed: {likelihood_name[safe.spoof]}")
@@ -888,7 +836,10 @@ class SafetyCheck():
         # Check the levels of adult, violence, racy, etc. content.
         if (safe.adult > LEVEL or
             safe.medical > LEVEL or
-            safe.violence > LEVEL):
+            # safe.spoof > LEVEL or
+            safe.violence > LEVEL #or
+            # safe.racy > LEVEL
+            ):
             print("Found violation")
             return True  # The image violates safe search guidelines.