LPX55 commited on
Commit
d7d030d
·
1 Parent(s): 91032a2

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

Files changed (1) hide show
  1. app.py +5 -3
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
- print(result_5output)
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['AI Image'] >= 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,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: </span>
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
  >