Create app.py

app.py

@@ -0,0 +1,113 @@
+import streamlit as st
+import requests
+import base64
+from PIL import Image
+import io
+
+st.set_page_config(page_title="AI Image Detector", page_icon="🔍")
+
+st.title("AI Image Detector")
+st.write("Upload an image to check if it's AI-generated")
+
+# Initialize session state for API key
+if 'api_key' not in st.session_state:
+    st.session_state.api_key = ''
+
+# API key input
+api_key = st.text_input("Enter NVIDIA API Key:", value=st.session_state.api_key, type="password")
+st.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')
+                }
+        
+        st.error("Unexpected response format from API")
+        return None
+            
+    except requests.exceptions.RequestException as e:
+        st.error(f"Error processing image: {str(e)}")
+        return None
+
+# File uploader
+uploaded_file = st.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)
+    st.image(image, caption="Uploaded Image", use_column_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 st.spinner("Analyzing image..."):
+        result = process_image(img_byte_arr, api_key)
+        
+        if result and result['status'] == 'SUCCESS':
+            confidence = result['confidence']
+            sources = result['sources']
+            
+            st.write("---")
+            st.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":
+                st.error(f"Is this image AI-generated? **{is_ai_generated}**")
+                
+                # Show top 3 possible sources if AI-generated
+                if sources:
+                    st.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
+                            st.write(f"- {source}: {prob:.1%}")
+            else:
+                st.success(f"Is this image AI-generated? **{is_ai_generated}**")
+            
+            # Show confidence score in smaller text
+            st.caption(f"Confidence score: {confidence:.2%}")
+
+elif not api_key and uploaded_file is not None:
+    st.warning("Please enter your NVIDIA API key first")
+
+# Add footer with instructions
+st.markdown("---")
+st.markdown("""
+### How to use:
+1. Enter your NVIDIA API key
+2. Upload an image (PNG, JPG, or JPEG)
+3. Wait for the analysis result
+4. Get a Yes/No answer based on whether the image is AI-generated
+""")