app.py

# import gradio as gr
# import requests
# import base64
# from PIL import Image
# import io

# gr.set_page_config(page_title="AI Image Detector", page_icon="🔍")

# gr.title("AI Image Detector")
# gr.write("Upload an image to check if it's AI-generated")


# api_key = NVIDIA_API_KEY
# gr.session_state.api_key = api_key

# def process_image(image_bytes, api_key):
#     header_auth = f"Bearer {api_key}"
#     invoke_url = "https://ai.api.nvidia.com/v1/cv/hive/ai-generated-image-detection"
    
#     # Convert image bytes to base64
#     image_b64 = base64.b64encode(image_bytes).decode()
    
#     payload = {
#         "input": [f"data:image/png;base64,{image_b64}"]
#     }
#     headers = {
#         "Content-Type": "application/json",
#         "Authorization": header_auth,
#         "Accept": "application/json",
#     }
    
#     try:
#         response = requests.post(invoke_url, headers=headers, json=payload)
#         response.raise_for_status()
#         result = response.json()
        
#         # Check if response contains the expected structure
#         if 'data' in result and len(result['data']) > 0:
#             first_result = result['data'][0]
#             if 'is_ai_generated' in first_result:
#                 return {
#                     'confidence': first_result['is_ai_generated'],
#                     'sources': first_result.get('possible_sources', {}),
#                     'status': first_result.get('status', 'UNKNOWN')
#                 }
        
#         gr.error("Unexpected response format from API")
#         return None
            
#     except requests.exceptions.RequestException as e:
#         gr.error(f"Error processing image: {str(e)}")
#         return None

# # File uploader
# uploaded_file = gr.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])

# if uploaded_file is not None and api_key:
#     # Display the uploaded image
#     image = Image.open(uploaded_file)
#     gr.image(image, caption="Uploaded Image", use_container_width=True)
    
#     # Convert image to bytes
#     img_byte_arr = io.BytesIO()
#     image.save(img_byte_arr, format=image.format)
#     img_byte_arr = img_byte_arr.getvalue()
    
#     # Process the image
#     with gr.spinner("Analyzing image..."):
#         result = process_image(img_byte_arr, api_key)
        
#         if result and result['status'] == 'SUCCESS':
#             confidence = result['confidence']
#             sources = result['sources']
            
#             gr.write("---")
#             gr.write("### Result")
            
#             # Determine if image is AI-generated (using 50% threshold)
#             is_ai_generated = "Yes" if confidence >= 0.5 else "No"
            
#             # Display result with appropriate styling
#             if is_ai_generated == "Yes":
#                 gr.error(f"Is this image AI-generated? **{is_ai_generated}**")
                
#                 # Show top 3 possible sources if AI-generated
#                 if sources:
#                     gr.write("Top possible AI models used:")
#                     sorted_sources = sorted(sources.items(), key=lambda x: x[1], reverse=True)[:3]
#                     for source, prob in sorted_sources:
#                         if prob > 0.01:  # Only show sources with >1% probability
#                             gr.write(f"- {source}: {prob:.1%}")
#             else:
#                 gr.success(f"Is this image AI-generated? **{is_ai_generated}**")
            
#             # Show confidence score in smaller text
#             gr.caption(f"Confidence score: {confidence:.2%}")

# elif not api_key and uploaded_file is not None:
#     gr.warning("Please enter your NVIDIA API key first")

# # Add footer with instructions
# gr.markdown("---")
# gr.markdown("""
# ---
# ### How to use:

# 1. Upload an image (PNG, JPG, or JPEG)
# 2. Wait for the analysis result
# 3. Get a ** Yes/No ** answer based on whether the image is AI-generated

# """)

import gradio as gr
import requests
import base64
from PIL import Image
import io

def process_image(image):
    """Process the uploaded image using NVIDIA's AI detection API"""
    # Get API key from environment variable for security
    api_key = os.getenv('NVIDIA_API_KEY')
    if not api_key:
        raise ValueError("NVIDIA API key not found in environment variables")
    
    header_auth = f"Bearer {api_key}"
    invoke_url = "https://ai.api.nvidia.com/v1/cv/hive/ai-generated-image-detection"
    
    # Convert PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    image_bytes = img_byte_arr.getvalue()
    
    # Convert image bytes to base64
    image_b64 = base64.b64encode(image_bytes).decode()
    
    payload = {
        "input": [f"data:image/png;base64,{image_b64}"]
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": header_auth,
        "Accept": "application/json",
    }
    
    try:
        response = requests.post(invoke_url, headers=headers, json=payload)
        response.raise_for_status()
        result = response.json()
        
        if 'data' in result and len(result['data']) > 0:
            first_result = result['data'][0]
            if 'is_ai_generated' in first_result:
                confidence = first_result['is_ai_generated']
                sources = first_result.get('possible_sources', {})
                
                # Format the result message
                is_ai_generated = "Yes" if confidence >= 0.5 else "No"
                result_message = f"Is this image AI-generated? {is_ai_generated}\n"
                result_message += f"Confidence score: {confidence:.2%}\n\n"
                
                if is_ai_generated == "Yes" and sources:
                    result_message += "Top possible AI models used:\n"
                    sorted_sources = sorted(sources.items(), key=lambda x: x[1], reverse=True)[:3]
                    for source, prob in sorted_sources:
                        if prob > 0.01:
                            result_message += f"- {source}: {prob:.1%}\n"
                
                return image, result_message
                
        return image, "Error: Unable to process image analysis results"
            
    except requests.exceptions.RequestException as e:
        return image, f"Error processing image: {str(e)}"

def create_demo():
    """Create and return the Gradio interface"""
    demo = gr.Interface(
        fn=process_image,
        inputs=gr.Image(type="pil", label="Upload Image"),
        outputs=[
            gr.Image(type="pil", label="Analyzed Image"),
            gr.Textbox(label="Analysis Results", lines=10)
        ],
        title="AI Image Detector",
        description="Upload an image to check if it's AI-generated",
        article="""
        ### How to use:
        1. Upload an image (PNG, JPG, or JPEG)
        2. Click the 'Submit' button
        3. Get a detailed analysis of whether the image is AI-generated
        
        Note: This application requires a valid NVIDIA API key set in environment variables.
        """
    )
    return demo

# Create and launch the application
if __name__ == "__main__":
    demo = create_demo()
    demo.launch()