125 lines
5.0 KiB
Python
125 lines
5.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO, 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/model_imdb_cross_person_4.22_99.46.pth.tar")
|
|
args.add_argument("--detector_weights", default="models/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/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(img_bgr):
|
|
try:
|
|
logging.info("Processing image")
|
|
# 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 {} |