LearningnRunning commited on
Commit
043f3b8
·
1 Parent(s): 83413d2

Update data_processing.py

Browse files
Files changed (1) hide show
  1. utils/data_processing.py +6 -1
utils/data_processing.py CHANGED
@@ -148,7 +148,12 @@ def detect_nsfw(input_image, conf_threshold=0.3, iou_threshold=0.45, label_mode=
148
  )
149
 
150
  # Construct detailed result text
151
- if label_counts:
 
 
 
 
 
152
  result_text = "Detected content:\n"
153
  for label, count in label_counts.items():
154
  result_text += f" - {label}: {count} instance(s)\n"
 
148
  )
149
 
150
  # Construct detailed result text
151
+ if "head" in label_counts:
152
+ head_count = label_counts["head"]
153
+ result_text = (
154
+ f"Detected content:\n - head: {head_count} instance(s)\nThis image appears to be safe."
155
+ )
156
+ elif label_counts:
157
  result_text = "Detected content:\n"
158
  for label, count in label_counts.items():
159
  result_text += f" - {label}: {count} instance(s)\n"