Schools NL, Ghost post utils, nude + age detection

app_cv/cv_processor.py

@@ -0,0 +1,132 @@
+import cv2
+import numpy as np
+import logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s', handlers=[logging.StreamHandler()])
+
+# Age
+from mivolo.predictor import Predictor
+import argparse
+# Nudity
+from nudenet import NudeDetector
+
+
+class CV():
+    def __init__(self):
+        args = argparse.ArgumentParser()
+        args.add_argument("--device", type=str, default="cpu")
+        args.add_argument("--checkpoint", default="models/mivolo/model_imdb_cross_person_4.22_99.46.pth.tar")
+        args.add_argument("--detector_weights", default="models/mivolo/yolov8x_person_face.pt")
+        args.add_argument("--with-persons", action="store_true", default=False, help="If set model will run with persons, if available")
+        args.add_argument("--disable-faces", action="store_true", default=False, help="If set model will use only persons if available")
+        args.add_argument("--draw", action="store_true", default=False, help="If set, resulted images will be drawn")
+        args = args.parse_args([])
+        # Initialize
+        self.predictor_age = Predictor(args)
+
+        # Initialize
+        self.nude_detector = NudeDetector(model_path="models/nude_detector/640m.onnx", inference_resolution=640)
+        # detector = NudeDetector(model_path="downloaded_640m.onnx path", inference_resolution=640)
+        # https://github.com/notAI-tech/NudeNet?tab=readme-ov-file#available-models
+
+        # All labels list
+        self.nudity_all_labels = [
+            "FEMALE_GENITALIA_COVERED",
+            "FACE_FEMALE",
+            "BUTTOCKS_EXPOSED",
+            "FEMALE_BREAST_EXPOSED",
+            "FEMALE_GENITALIA_EXPOSED",
+            "MALE_BREAST_EXPOSED",
+            "ANUS_EXPOSED",
+            "FEET_EXPOSED",
+            "BELLY_COVERED",
+            "FEET_COVERED",
+            "ARMPITS_COVERED",
+            "ARMPITS_EXPOSED",
+            "FACE_MALE",
+            "BELLY_EXPOSED",
+            "MALE_GENITALIA_EXPOSED",
+            "ANUS_COVERED",
+            "FEMALE_BREAST_COVERED",
+            "BUTTOCKS_COVERED",
+        ]
+        # Classes of interest
+        self.nudity_classes_of_interest = ["BUTTOCKS_EXPOSED", "FEMALE_BREAST_EXPOSED", "FEMALE_GENITALIA_EXPOSED", "ANUS_EXPOSED", "MALE_GENITALIA_EXPOSED"]
+
+    def _censor(self, image_bgr, detections):
+        # Copy original image
+        image_bgr_censored = image_bgr.copy()
+
+        for detection in detections:
+            box = detection["box"]
+            x, y, w, h = box[0], box[1], box[2], box[3]
+            # Change these pixels to pure black
+            image_bgr_censored[y : y + h, x : x + w] = (0, 0, 0)
+
+        return image_bgr_censored
+
+    def process_image(self, image_bgr):
+        ###################################################################
+        # Predict
+        detected_objects, out_img = self.predictor_age.recognize(image_bgr)
+        logging.debug("#persons: {}, #faces: {}".format(detected_objects.n_persons, detected_objects.n_faces))
+
+        # Num faces and persons detected
+        detected_objects.n_faces, detected_objects.n_persons
+        # Association
+        detected_objects.associate_faces_with_persons()
+
+        # detected_objects.face_to_person_map
+        # {2: 1, 3: 0}
+        # detected_objects.ages
+        # [None, None, 27.18, 23.77]
+        age_predictions = [e for e in detected_objects.ages if e is not None]
+
+        # Crops of faces & persons
+        # crops = detected_objects.collect_crops(img)
+        any_minor_present = any([ a < 18 for a in detected_objects.ages if a is not None ])
+        ###################################################################
+
+        ###################################################################
+        # Predict
+        nude_detections = self.nude_detector.detect(image_bgr)
+        logging.debug("Nude detections: {}".format(nude_detections))
+        # Filter by classes of interest
+        nude_detections = [ detection for detection in nude_detections if detection["class"] in self.nudity_classes_of_interest ]
+        # Nude detections present?
+        any_nude_detection = len(nude_detections) > 0
+        ###################################################################
+
+        ###################################################################
+        # Censor image
+        censored_img_bgr = self._censor(image_bgr, nude_detections)
+        # Plot age predictions on censored image
+        output_image = detected_objects.plot(img=censored_img_bgr)
+        ###################################################################
+
+        results = {
+            "any_minor_present": any_minor_present,
+            "any_nude_detection": any_nude_detection,
+            "nudity_detections": nude_detections,
+            "age_predictions": age_predictions,
+            "image": output_image,
+        }
+        return results
+
+def process(image_bytes):
+    try:
+        logging.info("Processing image")
+        # Convert bytes to NumPy array
+        img_array = np.frombuffer(image_bytes.getvalue(), dtype=np.uint8)
+        # Decode image using OpenCV
+        img_bgr = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+
+        if img_bgr is None:
+            return {}
+        # Process
+        results = CV().process_image(img_bgr)
+        logging.info("Returning results")
+        return results
+        
+    except Exception as e:
+        logging.warning("Error processing image: {}".format(str(e)))
+        return {}