SquadDetective / app.py
Brasd99's picture
Initial commit
8434bec
raw
history blame
3.35 kB
import json
import zipfile
import numpy as np
import cv2
import os
import gradio as gr
from deepface import DeepFace
from ultralytics import YOLO
import urllib.request
with open("config.json", "r") as f:
config = json.load(f)
FACE_DIST_TRESH = config['FACE_DIST_TRESH']
FACE_DET_TRESH = config['FACE_DET_TRESH']
YOLO_WEIGHTS_URL = config['YOLO_WEIGHTS_URL']
yolo_weights_filename = os.path.basename(YOLO_WEIGHTS_URL)
if not os.path.exists(yolo_weights_filename):
urllib.request.urlretrieve(YOLO_WEIGHTS_URL, yolo_weights_filename)
model = YOLO(yolo_weights_filename)
def find_distance(base_face, check_face):
result = DeepFace.verify(base_face, check_face, enforce_detection=False)
return result['distance']
def find_faces(image):
outputs = model(image)
faces = []
for box in outputs[0].boxes:
if float(box.conf) >= FACE_DET_TRESH:
x, y, w, h = [int(coord) for coord in box.xywh[0]]
x_center, y_center = x + w / 2, y + h / 2
x1 = int(x_center - w)
y1 = int(y_center - h)
crop_img = image[y1:y1+h, x1:x1+w]
faces.append(crop_img)
return faces
def load_images_from_zip(zip_path):
images = []
with zipfile.ZipFile(zip_path, 'r') as zip_file:
for file_name in zip_file.namelist():
with zip_file.open(file_name) as file:
img_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
img = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
images.append(img)
return [img for img in images if img is not None]
def create_image(images):
if not images:
return np.zeros((1, 1, 3), dtype=np.uint8)
max_width = max([img.shape[1] for img in images])
total_height = sum([img.shape[0] for img in images])
new_image = np.zeros((total_height, max_width, 3), dtype=np.uint8)
y_offset = 0
for img in images:
new_image[y_offset:y_offset+img.shape[0], :img.shape[1], :] = img
y_offset += img.shape[0]
new_image = cv2.resize(new_image, (0, 0), fx=0.5, fy=0.5)
return new_image
def check(avatars_zip, photos_zip):
avatars = load_images_from_zip(avatars_zip.name)
photos = load_images_from_zip(photos_zip.name)
not_found_faces = []
for photo in photos:
photo = cv2.cvtColor(photo, cv2.COLOR_RGB2BGR)
input_faces = find_faces(photo)
for input_face in input_faces:
avatars_checked = 0
for avatar in avatars:
avatar = cv2.cvtColor(avatar, cv2.COLOR_RGB2BGR)
distance = find_distance(avatar, input_face)
if distance > FACE_DIST_TRESH:
break
else:
avatars_checked+=1
if avatars_checked == len(avatars):
not_found_faces.append(input_face)
return create_image(not_found_faces)
title = '<h1 style="text-align:center">FootballChecker</h1>'
with gr.Blocks(theme='soft', title='FootballChecker') as blocks:
gr.HTML(title)
gr.Markdown('Face checking toolkit')
with gr.Row():
avatars = gr.inputs.File(label="Avatar Images (zip)")
photos = gr.inputs.File(label="Check Photos (zip)")
inputs = [avatars, photos]
process_button = gr.Button('Check images')
outputs=gr.outputs.Image(type="numpy", label="Results")
process_button.click(fn=check, inputs=inputs, outputs=outputs)
blocks.launch()