Spaces:

crangana
/

face-attribute-classifier

Sleeping

Rename app(3).py to app.py

a4b7561 verified 2 months ago

5.03 kB

import os
import base64
import requests
import gradio as gr
import spaces
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification, pipeline

# --- Roboflow face-detection key, read from a Space secret (not hardcoded) ---
ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")

# --- Load all three classifiers once at startup ---
# These run locally in this container via `transformers`, instead of calling
# Hugging Face's shared serverless Inference API (which no longer serves the
# gender/race models — only age currently has an attached provider).

def load_vit_pipeline(model_name):
    """
    Load explicitly as ViT rather than relying on AutoImageProcessor/AutoModel.
    These models were fine-tuned ~2 years ago using the standard HF ViT
    image-classification recipe; newer `transformers` versions sometimes
    fail to auto-detect older model configs, so we bypass Auto* entirely.
    """
    processor = ViTImageProcessor.from_pretrained(model_name)
    model = ViTForImageClassification.from_pretrained(model_name)
    return pipeline("image-classification", model=model, image_processor=processor)

print("Loading race/age/gender models — this can take a bit on first boot...")
race_pipe = load_vit_pipeline("crangana/trained-race")
age_pipe = load_vit_pipeline("crangana/trained-age")
gender_pipe = load_vit_pipeline("crangana/trained-gender")
print("Models loaded.")


def crop_face(image_path):
    """Detect and crop the face using Roboflow's hosted face detector."""
    if not ROBOFLOW_API_KEY:
        raise ValueError(
            "ROBOFLOW_API_KEY is not set. Add it under this Space's "
            "Settings → Repository secrets."
        )

    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")

    url = "https://detect.roboflow.com/face-det-npvna/1"
    params = {"api_key": ROBOFLOW_API_KEY}
    headers = {"Content-Type": "application/x-www-form-urlencoded"}

    response = requests.post(url, params=params, data=image_base64, headers=headers)

    try:
        crop_data = response.json()
    except ValueError:
        raise ValueError(
            f"Roboflow returned a non-JSON response (status {response.status_code}): "
            f"{response.text[:300]}"
        )

    print("Roboflow response:", crop_data)  # visible in the Space's Logs tab

    if "error" in crop_data:
        raise ValueError(f"Roboflow error: {crop_data['error']}")

    if response.status_code != 200:
        raise ValueError(f"Roboflow returned status {response.status_code}: {crop_data}")

    if not crop_data.get("predictions"):
        raise ValueError(
            "Roboflow ran successfully but found no face in this image. "
            "Try a clearer, front-facing photo."
        )

    pred = crop_data["predictions"][0]
    width, height = pred["width"], pred["height"]
    x = pred["x"] - width / 2
    y = pred["y"] - height / 2

    image = Image.open(image_path)
    cropped = image.crop((x, y, x + width, y + height))

    # Force an exact 224x224 square before it reaches the models — the
    # Roboflow bounding box isn't guaranteed to be square, and this
    # particular model's built-in preprocessing doesn't fully correct
    # for that on its own (was producing e.g. 227x224 instead of 224x224,
    # which the model's fixed patch grid can't accept).
    cropped = cropped.convert("RGB").resize((224, 224), Image.LANCZOS)

    cropped_path = "/tmp/cropped_face.jpg"
    cropped.save(cropped_path)
    return cropped_path


def get_top(results):
    """Pick the highest-confidence label from a list of {label, score} dicts."""
    best = max(results, key=lambda r: r["score"])
    return best["label"], best["score"]


@spaces.GPU
def classify(image):
    if image is None:
        return "No image provided.", "", ""

    tmp_path = "/tmp/input_face.jpg"
    image.convert("RGB").save(tmp_path)

    try:
        cropped_path = crop_face(tmp_path)
    except Exception as e:
        return f"Face detection failed: {e}", "—", "—"

    try:
        age_label, age_score = get_top(age_pipe(cropped_path))
        gender_label, gender_score = get_top(gender_pipe(cropped_path))
        race_label, race_score = get_top(race_pipe(cropped_path))
    except Exception as e:
        return f"Classification failed: {e}", "—", "—"

    return (
        f"{age_label} ({age_score:.0%})",
        f"{gender_label} ({gender_score:.0%})",
        f"{race_label} ({race_score:.0%})",
    )


demo = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="Upload a face photo"),
    outputs=[
        gr.Textbox(label="Age"),
        gr.Textbox(label="Gender"),
        gr.Textbox(label="Ethnicity"),
    ],
    title="Face Attribute Classifier",
    description="Predicts age, gender, and ethnicity from a face photo.",
    api_name="classify",  # names the programmatic API endpoint /gradio_api/call/classify
)

if __name__ == "__main__":
    demo.launch()

	import os
	import base64
	import requests
	import gradio as gr
	import spaces
	from PIL import Image
	from transformers import ViTImageProcessor, ViTForImageClassification, pipeline

	# --- Roboflow face-detection key, read from a Space secret (not hardcoded) ---
	ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")

	# --- Load all three classifiers once at startup ---
	# These run locally in this container via `transformers`, instead of calling
	# Hugging Face's shared serverless Inference API (which no longer serves the
	# gender/race models — only age currently has an attached provider).

	def load_vit_pipeline(model_name):
	"""
	Load explicitly as ViT rather than relying on AutoImageProcessor/AutoModel.
	These models were fine-tuned ~2 years ago using the standard HF ViT
	image-classification recipe; newer `transformers` versions sometimes
	fail to auto-detect older model configs, so we bypass Auto* entirely.
	"""
	processor = ViTImageProcessor.from_pretrained(model_name)
	model = ViTForImageClassification.from_pretrained(model_name)
	return pipeline("image-classification", model=model, image_processor=processor)

	print("Loading race/age/gender models — this can take a bit on first boot...")
	race_pipe = load_vit_pipeline("crangana/trained-race")
	age_pipe = load_vit_pipeline("crangana/trained-age")
	gender_pipe = load_vit_pipeline("crangana/trained-gender")
	print("Models loaded.")


	def crop_face(image_path):
	"""Detect and crop the face using Roboflow's hosted face detector."""
	if not ROBOFLOW_API_KEY:
	raise ValueError(
	"ROBOFLOW_API_KEY is not set. Add it under this Space's "
	"Settings → Repository secrets."
	)

	with open(image_path, "rb") as image_file:
	image_base64 = base64.b64encode(image_file.read()).decode("utf-8")

	url = "https://detect.roboflow.com/face-det-npvna/1"
	params = {"api_key": ROBOFLOW_API_KEY}
	headers = {"Content-Type": "application/x-www-form-urlencoded"}

	response = requests.post(url, params=params, data=image_base64, headers=headers)

	try:
	crop_data = response.json()
	except ValueError:
	raise ValueError(
	f"Roboflow returned a non-JSON response (status {response.status_code}): "
	f"{response.text[:300]}"
	)

	print("Roboflow response:", crop_data) # visible in the Space's Logs tab

	if "error" in crop_data:
	raise ValueError(f"Roboflow error: {crop_data['error']}")

	if response.status_code != 200:
	raise ValueError(f"Roboflow returned status {response.status_code}: {crop_data}")

	if not crop_data.get("predictions"):
	raise ValueError(
	"Roboflow ran successfully but found no face in this image. "
	"Try a clearer, front-facing photo."
	)

	pred = crop_data["predictions"][0]
	width, height = pred["width"], pred["height"]
	x = pred["x"] - width / 2
	y = pred["y"] - height / 2

	image = Image.open(image_path)
	cropped = image.crop((x, y, x + width, y + height))

	# Force an exact 224x224 square before it reaches the models — the
	# Roboflow bounding box isn't guaranteed to be square, and this
	# particular model's built-in preprocessing doesn't fully correct
	# for that on its own (was producing e.g. 227x224 instead of 224x224,
	# which the model's fixed patch grid can't accept).
	cropped = cropped.convert("RGB").resize((224, 224), Image.LANCZOS)

	cropped_path = "/tmp/cropped_face.jpg"
	cropped.save(cropped_path)
	return cropped_path


	def get_top(results):
	"""Pick the highest-confidence label from a list of {label, score} dicts."""
	best = max(results, key=lambda r: r["score"])
	return best["label"], best["score"]


	@spaces.GPU
	def classify(image):
	if image is None:
	return "No image provided.", "", ""

	tmp_path = "/tmp/input_face.jpg"
	image.convert("RGB").save(tmp_path)

	try:
	cropped_path = crop_face(tmp_path)
	except Exception as e:
	return f"Face detection failed: {e}", "—", "—"

	try:
	age_label, age_score = get_top(age_pipe(cropped_path))
	gender_label, gender_score = get_top(gender_pipe(cropped_path))
	race_label, race_score = get_top(race_pipe(cropped_path))
	except Exception as e:
	return f"Classification failed: {e}", "—", "—"

	return (
	f"{age_label} ({age_score:.0%})",
	f"{gender_label} ({gender_score:.0%})",
	f"{race_label} ({race_score:.0%})",
	)


	demo = gr.Interface(
	fn=classify,
	inputs=gr.Image(type="pil", label="Upload a face photo"),
	outputs=[
	gr.Textbox(label="Age"),
	gr.Textbox(label="Gender"),
	gr.Textbox(label="Ethnicity"),
	],
	title="Face Attribute Classifier",
	description="Predicts age, gender, and ethnicity from a face photo.",
	api_name="classify", # names the programmatic API endpoint /gradio_api/call/classify
	)

	if __name__ == "__main__":
	demo.launch()