noumanjavaid commited on
Commit
8959c8c
·
verified ·
1 Parent(s): 4b17346

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +113 -0
app.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import requests
3
+ import base64
4
+ from PIL import Image
5
+ import io
6
+
7
+ st.set_page_config(page_title="AI Image Detector", page_icon="🔍")
8
+
9
+ st.title("AI Image Detector")
10
+ st.write("Upload an image to check if it's AI-generated")
11
+
12
+ # Initialize session state for API key
13
+ if 'api_key' not in st.session_state:
14
+ st.session_state.api_key = ''
15
+
16
+ # API key input
17
+ api_key = st.text_input("Enter NVIDIA API Key:", value=st.session_state.api_key, type="password")
18
+ st.session_state.api_key = api_key
19
+
20
+ def process_image(image_bytes, api_key):
21
+ header_auth = f"Bearer {api_key}"
22
+ invoke_url = "https://ai.api.nvidia.com/v1/cv/hive/ai-generated-image-detection"
23
+
24
+ # Convert image bytes to base64
25
+ image_b64 = base64.b64encode(image_bytes).decode()
26
+
27
+ payload = {
28
+ "input": [f"data:image/png;base64,{image_b64}"]
29
+ }
30
+ headers = {
31
+ "Content-Type": "application/json",
32
+ "Authorization": header_auth,
33
+ "Accept": "application/json",
34
+ }
35
+
36
+ try:
37
+ response = requests.post(invoke_url, headers=headers, json=payload)
38
+ response.raise_for_status()
39
+ result = response.json()
40
+
41
+ # Check if response contains the expected structure
42
+ if 'data' in result and len(result['data']) > 0:
43
+ first_result = result['data'][0]
44
+ if 'is_ai_generated' in first_result:
45
+ return {
46
+ 'confidence': first_result['is_ai_generated'],
47
+ 'sources': first_result.get('possible_sources', {}),
48
+ 'status': first_result.get('status', 'UNKNOWN')
49
+ }
50
+
51
+ st.error("Unexpected response format from API")
52
+ return None
53
+
54
+ except requests.exceptions.RequestException as e:
55
+ st.error(f"Error processing image: {str(e)}")
56
+ return None
57
+
58
+ # File uploader
59
+ uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])
60
+
61
+ if uploaded_file is not None and api_key:
62
+ # Display the uploaded image
63
+ image = Image.open(uploaded_file)
64
+ st.image(image, caption="Uploaded Image", use_column_width=True)
65
+
66
+ # Convert image to bytes
67
+ img_byte_arr = io.BytesIO()
68
+ image.save(img_byte_arr, format=image.format)
69
+ img_byte_arr = img_byte_arr.getvalue()
70
+
71
+ # Process the image
72
+ with st.spinner("Analyzing image..."):
73
+ result = process_image(img_byte_arr, api_key)
74
+
75
+ if result and result['status'] == 'SUCCESS':
76
+ confidence = result['confidence']
77
+ sources = result['sources']
78
+
79
+ st.write("---")
80
+ st.write("### Result")
81
+
82
+ # Determine if image is AI-generated (using 50% threshold)
83
+ is_ai_generated = "Yes" if confidence >= 0.5 else "No"
84
+
85
+ # Display result with appropriate styling
86
+ if is_ai_generated == "Yes":
87
+ st.error(f"Is this image AI-generated? **{is_ai_generated}**")
88
+
89
+ # Show top 3 possible sources if AI-generated
90
+ if sources:
91
+ st.write("Top possible AI models used:")
92
+ sorted_sources = sorted(sources.items(), key=lambda x: x[1], reverse=True)[:3]
93
+ for source, prob in sorted_sources:
94
+ if prob > 0.01: # Only show sources with >1% probability
95
+ st.write(f"- {source}: {prob:.1%}")
96
+ else:
97
+ st.success(f"Is this image AI-generated? **{is_ai_generated}**")
98
+
99
+ # Show confidence score in smaller text
100
+ st.caption(f"Confidence score: {confidence:.2%}")
101
+
102
+ elif not api_key and uploaded_file is not None:
103
+ st.warning("Please enter your NVIDIA API key first")
104
+
105
+ # Add footer with instructions
106
+ st.markdown("---")
107
+ st.markdown("""
108
+ ### How to use:
109
+ 1. Enter your NVIDIA API key
110
+ 2. Upload an image (PNG, JPG, or JPEG)
111
+ 3. Wait for the analysis result
112
+ 4. Get a Yes/No answer based on whether the image is AI-generated
113
+ """)