import os import json import math import pyzipper import numpy as np import gradio as gr import numpy as np import opennsfw2 as n2 import base64 from annoy import AnnoyIndex from deepface.commons import functions from deepface.basemodels import Facenet512 from fastcore.all import * from fastai.vision.all import * os.environ["DEEPFACE_HOME"] = "." yahooNsfwModel = n2.make_open_nsfw_model() # load the face model model = Facenet512.loadModel() input_shape_x, input_shape_y = functions.find_input_shape(model) index = AnnoyIndex(512, "euclidean") index.load(f"face.db") ANNOY_INDEX = json.load(open(f"face.json")) with pyzipper.AESZipFile("persons.zip") as zf: password = os.getenv("VISAGE_KEY", "").encode("ascii") zf.setpassword(password) PERFORMER_DB = json.loads(zf.read("performers.json")) ## Prediction functions def image_search_performer(image, threshold=20.0, results=3): """Search for a performer in an image Returns a list of performers with at least following keys: - id: the performer's id - distance: the distance between the face in the image and the performer's face - confidence: a confidence score between 0 and 100 - hits: the number of times the performer was found in our database """ image_array = np.array(image) try : img = functions.preprocess_face( img=image_array, target_size=(input_shape_x, input_shape_y), detector_backend="retinaface", align=True, ) img = functions.normalize_input(img, normalization="Facenet2018") face = model.predict(img)[0] return search_performer(face, threshold, results) except Exception as e: print(e, "for img", image) return [] def search_performer(vector, threshold=20.0, results=3): threshold = threshold or 20.0 results = results or 3 ids, distances = index.get_nns_by_vector( vector, 50, search_k=10000, include_distances=True ) persons = {} for p, distance in zip(ids, distances): id = ANNOY_INDEX[p] if id in persons: persons[id]["hits"] += 1 persons[id]["distance"] -= 0.5 persons[id]["confidence"] = normalize_confidence_from_distance( persons[id]["distance"], threshold ) continue persons[id] = { "id": id, "distance": round(distance, 2), "confidence": normalize_confidence_from_distance(distance, threshold), "hits": 1, } if id in PERFORMER_DB: persons[id].update(PERFORMER_DB.get(id)) persons = sorted(persons.values(), key=lambda x: x["distance"]) persons = [p for p in persons if p["distance"] < threshold] return persons[:results] def normalize_confidence_from_distance(distance, threshold=20.0): """Normalize confidence to 0-100 scale""" confidence = face_distance_to_conf(distance, threshold) return int(((confidence - 0.0) / (1.0 - 0.0)) * (100.0 - 0.0) + 0.0) def face_distance_to_conf(face_distance, face_match_threshold=20.0): """Using a face distance, calculate a similarity confidence value""" if face_distance > face_match_threshold: # The face is far away, so give it a low confidence range = 1.0 - face_match_threshold linear_val = (1.0 - face_distance) / (range * 2.0) return linear_val else: # The face is close, so give it a high confidence range = face_match_threshold linear_val = 1.0 - (face_distance / (range * 2.0)) # But adjust this value by a curve so that we don't get a linear # transition from close to far. We want it to be more confident # the closer it is. return linear_val + ((1.0 - linear_val) * math.pow((linear_val - 0.5) * 2, 0.2)) def predict(image, vtt): vtt = base64.b64decode(vtt.replace("data:text/vtt;base64,", "")) sprite = PILImage.create(image) pre_process_data = [] for left, top, right, bottom in getVTToffsets(vtt): cut_frame = sprite.crop((left, top, left + right, top + bottom)) image = n2.preprocess_image(cut_frame, n2.Preprocessing.YAHOO) pre_process_data.append( (np.expand_dims(image, axis=0), cut_frame, (left, top, right, bottom)) ) offsets = [] images = [] tensors = [i[0] for i in pre_process_data] predictions = yahooNsfwModel.predict(np.vstack(tensors)) for i, prediction in enumerate(predictions): if prediction[0] < 0.5: images.append(PILImage.create(np.asarray(pre_process_data[i][1]))) offsets.append(pre_process_data[i][2]) persons = {} for image in images: personList = image_search_performer(image) for person in personList: person_id = person["id"] if person_id not in persons: persons[person_id] = person else: existing_person = persons[person_id] existing_person["hits"] += person["hits"] if person["distance"] < existing_person["distance"]: existing_person["distance"] = person["distance"] if person["confidence"] > existing_person["confidence"]: existing_person["confidence"] = person["confidence"] return persons def getVTToffsets(vtt): left = top = right = bottom = None for line in vtt.decode("utf-8").split("\n"): line = line.strip() if "xywh=" in line: left, top, right, bottom = line.split("xywh=")[-1].split(",") left, top, right, bottom = ( int(left), int(top), int(right), int(bottom), ) else: continue if not left: continue yield left, top, right, bottom image_search = gr.Interface( fn=image_search_performer, inputs=[ gr.components.Image(), gr.components.Slider(label="threshold", minimum=0.0, maximum=30.0, value=20.0), gr.components.Slider(label="results", minimum=0, maximum=50, value=3, step=1), ], outputs=gr.outputs.JSON(label=""), title="Who is in the photo?", description="Upload an image of a person and we'll tell you who it is.", ) sprite_search = gr.Interface( fn=predict, inputs=[ gr.Image(), gr.Textbox(label="VTT file"), ], outputs=gr.JSON(label=""), ).launch(enable_queue=True, server_name="0.0.0.0") gr.TabbedInterface([image_search, sprite_search]).launch( enable_queue=True, server_name="0.0.0.0" )