Spaces:

J007acky
/

FaceRecog

Sleeping

FaceRecog / FastAPI.py

Update FastAPI.py

d1ce429 verified over 2 years ago

6.33 kB

from fastapi import FastAPI
from pydantic import BaseModel, HttpUrl
import requests
import face_recognition
import pickle
import cv2
import pyrebase
import os


# Initialize Firebase using Pyrebase
config = {
    "apiKey": "AIzaSyClnRJAnrJgAgkYjuYnlvu-CJ6Cxyklebo",
    "databaseURL": "https://console.firebase.google.com/project/socioverse-2025/database/socioverse-2025-default-rtdb/data/~2F",
    "authDomain": "socioverse-2025.firebaseapp.com",
    "projectId": "socioverse-2025",
    "storageBucket": "socioverse-2025.appspot.com",
    "messagingSenderId": "689574504641",
    "appId": "1:689574504641:web:a22f6a2fa343e4221acc40",
    "serviceAccount":"socioverse-2025-firebase-adminsdk-gcc6m-6bfb53e6d9.json"
}

firebase = pyrebase.initialize_app(config)
storage = firebase.storage()

# Define the folder containing face images in the Firebase Storage bucket
storage_folder = "Faces/"


app = FastAPI()

class ImgSave(BaseModel):
    image_url: HttpUrl
    user_name: str

class ImgInput(BaseModel):
    image_url: HttpUrl

class ImgOutput(BaseModel):
    user_id: list

class UserSaved(BaseModel):
    status: str

class UserDelete(BaseModel):
    label: str

class Message(BaseModel):
    message: str

class CleanPickle(BaseModel):
    confirm: bool


def recognize_face(image_url: HttpUrl) -> ImgOutput:

    storage.child().download("Faces/pkl/face_encodings.pkl","face_encodings.pkl")
    # Downloading image
    response = requests.get(image_url)
    with open("examp.jpg", 'wb') as file:
        file.write(response.content)

    # Load the stored face encodings and labels from the pickle file
    with open("face_encodings.pkl", "rb") as file:
        data = pickle.load(file)
        face_encodings = data["encodings"]
        labels = data["labels"]


        # Load the new image you want to recognize
        new_image = cv2.imread("examp.jpg")

        # Find face encodings in the new image
        new_face_encodings = face_recognition.face_encodings(new_image)

        if len(new_face_encodings) == 0:
            print("No faces found in the new image.")
            return ImgOutput(label=["unable to detect"])
        else:
            output_labels = []

            for new_face_encoding in new_face_encodings:
                # Compare the new face encoding to the stored encodings
                results = face_recognition.compare_faces(face_encodings, new_face_encoding)

                for i, result in enumerate(results):
                    if result:
                        output_labels.append(labels[i])

            os.remove("examp.jpg")

            if output_labels:
                return ImgOutput(user_id=output_labels)
            else:
                out = ["unable to detect"]
                return ImgOutput(user_id=out)



def add_face(image_url: HttpUrl,user_name : str):
    # Downloading image
    response = requests.get(image_url)
    with open("examp.jpg", 'wb') as file:
        file.write(response.content)

    # Load the stored face encodings and labels from the pickle file
    with open("face_encodings.pkl", "rb") as file:
        data = pickle.load(file)
        face_encodings = data["encodings"]
        labels = data["labels"]

    if user_name in labels:
        return "User already exists"
        
    # Load a new image you want to recognize
    new_image = cv2.imread("examp.jpg")

    # Convert the BGR image to RGB
    rgb_img = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)

    encode = face_recognition.face_encodings(new_image)[0]
    if len(encode) == 0:
        return "No face found"
    face_encodings.append(encode)
    labels.append(user_name)

    # Delete the temporary downloaded image
    os.remove("examp.jpg")

    # Save the encodings and labels to a pickle file
    data = {"encodings": face_encodings, "labels": labels}
    with open("face_encodings.pkl", "wb") as file:
        pickle.dump(data, file)

    # Upload the pickle file to Firebase Storage
    pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
    pkl_blob.put("face_encodings.pkl")
    return "User Saved"


# Function to delete a user
def delete_user(user_name: str):
    # Load the stored face encodings and labels from the pickle file
    with open("face_encodings.pkl", "rb") as file:
        data = pickle.load(file)
        face_encodings = data["encodings"]
        labels = data["labels"]

    if user_name in labels:
        index = labels.index(user_name)
        del labels[index]
        del face_encodings[index]

        # Save the updated encodings and labels to the pickle file
        data = {"encodings": face_encodings, "labels": labels}
        with open("face_encodings.pkl", "wb") as file:
            pickle.dump(data, file)

        # Upload the updated pickle file to Firebase Storage
        pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
        pkl_blob.put("face_encodings.pkl")

        return {"message": f"User '{user_name}' deleted successfully."}
    else:
        return {"message": f"User '{user_name}' not found."}


def clean_pickle(confirm: bool):
    if confirm:
        # Remove the pickle file
        if os.path.exists("face_encodings.pkl"):
            os.remove("face_encodings.pkl")

        # Create an empty pickle file
        with open("face_encodings.pkl", "wb") as file:
            data = {"encodings": [], "labels": []}
            pickle.dump(data, file)

        # Upload the empty pickle file to Firebase Storage
        pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
        pkl_blob.put("face_encodings.pkl")

        return {"message": "Pickle file cleaned and uploaded successfully."}
    else:
        return {"message": "Confirmation required to clean the pickle file."}


@app.post('/')
async def scoring_endpoint(item:ImgInput):
    result = recognize_face(item.image_url)
    return result


@app.post('/user/')
async def scoring_endpoint(item:ImgSave):
    results = add_face(item.image_url, item.user_name)
    return Message(message=results)


@app.delete('/user_delete/')
async def scoring_endpoint(item: UserDelete):
    result = delete_user(item.label)
    return Message(message=result["message"])


@app.delete('/clean/')
async def clean_pickle_endpoint(item: CleanPickle):
    result = clean_pickle(item.confirm)
    return result

	from fastapi import FastAPI
	from pydantic import BaseModel, HttpUrl
	import requests
	import face_recognition
	import pickle
	import cv2
	import pyrebase
	import os


	# Initialize Firebase using Pyrebase
	config = {
	"apiKey": "AIzaSyClnRJAnrJgAgkYjuYnlvu-CJ6Cxyklebo",
	"databaseURL": "https://console.firebase.google.com/project/socioverse-2025/database/socioverse-2025-default-rtdb/data/~2F",
	"authDomain": "socioverse-2025.firebaseapp.com",
	"projectId": "socioverse-2025",
	"storageBucket": "socioverse-2025.appspot.com",
	"messagingSenderId": "689574504641",
	"appId": "1:689574504641:web:a22f6a2fa343e4221acc40",
	"serviceAccount":"socioverse-2025-firebase-adminsdk-gcc6m-6bfb53e6d9.json"
	}

	firebase = pyrebase.initialize_app(config)
	storage = firebase.storage()

	# Define the folder containing face images in the Firebase Storage bucket
	storage_folder = "Faces/"


	app = FastAPI()

	class ImgSave(BaseModel):
	image_url: HttpUrl
	user_name: str

	class ImgInput(BaseModel):
	image_url: HttpUrl

	class ImgOutput(BaseModel):
	user_id: list

	class UserSaved(BaseModel):
	status: str

	class UserDelete(BaseModel):
	label: str

	class Message(BaseModel):
	message: str

	class CleanPickle(BaseModel):
	confirm: bool


	def recognize_face(image_url: HttpUrl) -> ImgOutput:

	storage.child().download("Faces/pkl/face_encodings.pkl","face_encodings.pkl")
	# Downloading image
	response = requests.get(image_url)
	with open("examp.jpg", 'wb') as file:
	file.write(response.content)

	# Load the stored face encodings and labels from the pickle file
	with open("face_encodings.pkl", "rb") as file:
	data = pickle.load(file)
	face_encodings = data["encodings"]
	labels = data["labels"]


	# Load the new image you want to recognize
	new_image = cv2.imread("examp.jpg")

	# Find face encodings in the new image
	new_face_encodings = face_recognition.face_encodings(new_image)

	if len(new_face_encodings) == 0:
	print("No faces found in the new image.")
	return ImgOutput(label=["unable to detect"])
	else:
	output_labels = []

	for new_face_encoding in new_face_encodings:
	# Compare the new face encoding to the stored encodings
	results = face_recognition.compare_faces(face_encodings, new_face_encoding)

	for i, result in enumerate(results):
	if result:
	output_labels.append(labels[i])

	os.remove("examp.jpg")

	if output_labels:
	return ImgOutput(user_id=output_labels)
	else:
	out = ["unable to detect"]
	return ImgOutput(user_id=out)



	def add_face(image_url: HttpUrl,user_name : str):
	# Downloading image
	response = requests.get(image_url)
	with open("examp.jpg", 'wb') as file:
	file.write(response.content)

	# Load the stored face encodings and labels from the pickle file
	with open("face_encodings.pkl", "rb") as file:
	data = pickle.load(file)
	face_encodings = data["encodings"]
	labels = data["labels"]

	if user_name in labels:
	return "User already exists"

	# Load a new image you want to recognize
	new_image = cv2.imread("examp.jpg")

	# Convert the BGR image to RGB
	rgb_img = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)

	encode = face_recognition.face_encodings(new_image)[0]
	if len(encode) == 0:
	return "No face found"
	face_encodings.append(encode)
	labels.append(user_name)

	# Delete the temporary downloaded image
	os.remove("examp.jpg")

	# Save the encodings and labels to a pickle file
	data = {"encodings": face_encodings, "labels": labels}
	with open("face_encodings.pkl", "wb") as file:
	pickle.dump(data, file)

	# Upload the pickle file to Firebase Storage
	pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
	pkl_blob.put("face_encodings.pkl")
	return "User Saved"


	# Function to delete a user
	def delete_user(user_name: str):
	# Load the stored face encodings and labels from the pickle file
	with open("face_encodings.pkl", "rb") as file:
	data = pickle.load(file)
	face_encodings = data["encodings"]
	labels = data["labels"]

	if user_name in labels:
	index = labels.index(user_name)
	del labels[index]
	del face_encodings[index]

	# Save the updated encodings and labels to the pickle file
	data = {"encodings": face_encodings, "labels": labels}
	with open("face_encodings.pkl", "wb") as file:
	pickle.dump(data, file)

	# Upload the updated pickle file to Firebase Storage
	pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
	pkl_blob.put("face_encodings.pkl")

	return {"message": f"User '{user_name}' deleted successfully."}
	else:
	return {"message": f"User '{user_name}' not found."}


	def clean_pickle(confirm: bool):
	if confirm:
	# Remove the pickle file
	if os.path.exists("face_encodings.pkl"):
	os.remove("face_encodings.pkl")

	# Create an empty pickle file
	with open("face_encodings.pkl", "wb") as file:
	data = {"encodings": [], "labels": []}
	pickle.dump(data, file)

	# Upload the empty pickle file to Firebase Storage
	pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
	pkl_blob.put("face_encodings.pkl")

	return {"message": "Pickle file cleaned and uploaded successfully."}
	else:
	return {"message": "Confirmation required to clean the pickle file."}


	@app.post('/')
	async def scoring_endpoint(item:ImgInput):
	result = recognize_face(item.image_url)
	return result


	@app.post('/user/')
	async def scoring_endpoint(item:ImgSave):
	results = add_face(item.image_url, item.user_name)
	return Message(message=results)


	@app.delete('/user_delete/')
	async def scoring_endpoint(item: UserDelete):
	result = delete_user(item.label)
	return Message(message=result["message"])


	@app.delete('/clean/')
	async def clean_pickle_endpoint(item: CleanPickle):
	result = clean_pickle(item.confirm)
	return result