app.py

import streamlit as st
import requests
import os

# Load API key from Hugging Face secrets
api_key = os.getenv("KEY")

# Streamlit UI
st.title("Head Extraction Tool")
st.write("Extract heads from images using AI.")

uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])

def extract_head(image):
    url = "https://head-extraction1.p.rapidapi.com/head-extraction"
    headers = {
        "x-rapidapi-key": api_key,
        "x-rapidapi-host": "head-extraction1.p.rapidapi.com"
    }
    files = {"image": (image.name, image, image.type)}
    response = requests.post(url, files=files, headers=headers)
    return response.json()

if st.button("Extract Head"):
    if uploaded_file:
        with st.spinner("Extracting head..."):
            result = extract_head(uploaded_file)
            
            # Extract image URL from JSON response
            head_url = result.get("message", {}).get("head_image_url")
            if head_url:
                st.image(head_url, caption="Extracted Head", use_container_width=True)
                
                # Provide download button
                st.download_button(
                    label="Download Head Image",
                    data=requests.get(head_url).content,
                    file_name="extracted_head.png",
                    mime="image/png"
                )
            else:
                st.error("Failed to extract head. No image URL found in response.")
                st.write("Response Data:", result)  # Debugging: Show response data
    else:
        st.warning("Please upload an image.")