Fix ViT-5 model prediction logic and output formatting
Browse files- Modify confidence threshold check to use 'Deepfake' instead of 'AI Image'
- Remove unnecessary print statement before final output print
- Adjust print statement placement for better logging
- Minor UI text spacing adjustment in results HTML template
app.py
CHANGED
@@ -203,13 +203,13 @@ def predict_image(img, confidence_threshold):
|
|
203 |
prediction_5 = clf_5(img_pilvits)
|
204 |
result_5 = {pred['label']: pred['score'] for pred in prediction_5}
|
205 |
result_5output = [5, 'ViT-base Newcomer', result_5['Realism'], result_5['Deepfake']]
|
206 |
-
|
207 |
# Ensure the result dictionary contains all class names
|
208 |
for class_name in class_names_5:
|
209 |
if class_name not in result_5:
|
210 |
result_5[class_name] = 0.0
|
211 |
# Check if either class meets the confidence threshold
|
212 |
-
if result_5['
|
213 |
label_5 = f"AI, Confidence: {result_5['Deepfake']:.4f}"
|
214 |
result_5output += ['AI']
|
215 |
elif result_5['Real Image'] >= confidence_threshold:
|
@@ -220,6 +220,8 @@ def predict_image(img, confidence_threshold):
|
|
220 |
result_5output += ['UNCERTAIN']
|
221 |
except Exception as e:
|
222 |
label_5 = f"Error: {str(e)}"
|
|
|
|
|
223 |
|
224 |
|
225 |
# try:
|
@@ -452,7 +454,7 @@ def generate_results_html(results):
|
|
452 |
class="flex flex-col bg-gray-800 rounded-sm p-4 m-1 border border-gray-800 shadow-xs transition hover:shadow-lg dark:shadow-gray-700/25">
|
453 |
<div
|
454 |
class="-m-4 h-24 {get_header_color(results[4][-1])[0]} rounded-sm rounded-b-none transition border group-hover:border-gray-100 group-hover:shadow-lg group-hover:{get_header_color(results[4][-1])[4]}">
|
455 |
-
<span class="text-gray-300 font-mono tracking-widest p-4 pb-3 block text-xs text-center">MODEL 5
|
456 |
<span
|
457 |
class="flex w-24 mx-auto tracking-wide items-center justify-center rounded-full {get_header_color(results[4][-1])[2]} px-1 py-0.5 {get_header_color(results[4][-1])[3]}"
|
458 |
>
|
|
|
203 |
prediction_5 = clf_5(img_pilvits)
|
204 |
result_5 = {pred['label']: pred['score'] for pred in prediction_5}
|
205 |
result_5output = [5, 'ViT-base Newcomer', result_5['Realism'], result_5['Deepfake']]
|
206 |
+
|
207 |
# Ensure the result dictionary contains all class names
|
208 |
for class_name in class_names_5:
|
209 |
if class_name not in result_5:
|
210 |
result_5[class_name] = 0.0
|
211 |
# Check if either class meets the confidence threshold
|
212 |
+
if result_5['Deepfake'] >= confidence_threshold:
|
213 |
label_5 = f"AI, Confidence: {result_5['Deepfake']:.4f}"
|
214 |
result_5output += ['AI']
|
215 |
elif result_5['Real Image'] >= confidence_threshold:
|
|
|
220 |
result_5output += ['UNCERTAIN']
|
221 |
except Exception as e:
|
222 |
label_5 = f"Error: {str(e)}"
|
223 |
+
|
224 |
+
print(result_5output)
|
225 |
|
226 |
|
227 |
# try:
|
|
|
454 |
class="flex flex-col bg-gray-800 rounded-sm p-4 m-1 border border-gray-800 shadow-xs transition hover:shadow-lg dark:shadow-gray-700/25">
|
455 |
<div
|
456 |
class="-m-4 h-24 {get_header_color(results[4][-1])[0]} rounded-sm rounded-b-none transition border group-hover:border-gray-100 group-hover:shadow-lg group-hover:{get_header_color(results[4][-1])[4]}">
|
457 |
+
<span class="text-gray-300 font-mono tracking-widest p-4 pb-3 block text-xs text-center">MODEL 5:</span>
|
458 |
<span
|
459 |
class="flex w-24 mx-auto tracking-wide items-center justify-center rounded-full {get_header_color(results[4][-1])[2]} px-1 py-0.5 {get_header_color(results[4][-1])[3]}"
|
460 |
>
|