elfarash's picture
Update output format to user friendly Gradio Label
5a52f0b
Raw
History Blame Contribute Delete
3.38 kB
import os
import json
import base64
import requests
import gradio as gr
from dotenv import load_dotenv
# Load environment variables (useful for local testing)
load_dotenv()
# The API Key must be set in the HuggingFace Space "Secrets"
API_KEY = os.getenv("HIVE_API_KEY")
HIVE_API_URL = "https://api.thehive.ai/api/v3/hive/ai-generated-and-deepfake-content-detection"
def detect_ai_content(image_filepath):
if not API_KEY:
raise gr.Error("HIVE_API_KEY environment variable is not set. Please add it to your HuggingFace Space Secrets.")
if image_filepath is None:
raise gr.Error("No image provided.")
try:
# Read the image file and encode it to Base64
with open(image_filepath, "rb") as image_file:
file_bytes = image_file.read()
encoded_string = base64.b64encode(file_bytes).decode("utf-8")
# Format the Base64 string for the Hive API
media_base64 = f"data:image/jpeg;base64,{encoded_string}"
payload = {
"media_metadata": True,
"input": [{"media_base64": media_base64}]
}
headers = {
"Authorization": f"Bearer {API_KEY}",
"Content-Type": "application/json",
"Accept": "application/json"
}
response = requests.post(
HIVE_API_URL,
headers=headers,
data=json.dumps(payload),
timeout=60
)
# Raise an exception if the response indicates an HTTP error
response.raise_for_status()
data = response.json()
real_val = 0.0
fake_val = 0.0
# Parse the JSON to extract the required classes
if "output" in data and len(data["output"]) > 0:
classes = data["output"][0].get("classes", [])
for c in classes:
class_name = c.get("class", "")
if class_name == "not_ai_generated" or class_name == "none":
real_val = c.get("value", 0.0)
elif class_name == "ai_generated":
fake_val = c.get("value", 0.0)
# If the API returned specific model percentages (like midjourney, dall-e)
# but no overarching 'ai_generated' tag, we sum all fake probabilities
if fake_val == 0.0 and real_val > 0.0:
fake_val = sum(c.get("value", 0.0) for c in classes if c.get("class") not in ["not_ai_generated", "none", "inconclusive", "not_ai_generated_audio", "ai_generated_audio"])
return {"Real": real_val, "Fake": fake_val}
except requests.exceptions.RequestException as e:
error_msg = f"API Error: {e}"
if hasattr(e, 'response') and e.response is not None:
error_msg += f" | Response: {e.response.text}"
raise gr.Error(error_msg)
except Exception as e:
raise gr.Error(f"Internal Error: {e}")
# Define the Gradio Interface
demo = gr.Interface(
fn=detect_ai_content,
inputs=gr.Image(type="filepath", label="Upload an Image"),
outputs=gr.Label(label="Detection Results"),
title="Hive AI Content Detection",
description="Upload an image to detect if it is AI-generated or contains deepfakes using the Hive AI API."
)
# Launch the app
if __name__ == "__main__":
demo.launch()