ADD IMAGE DETECTION

app/app_ui.py

@@ -1,8 +1,10 @@
 import gradio as gr
-from app.safety_check import is_prompt_safe
+from app.safety_check import is_prompt_safe, is_image_safe
 from app.gen_ai import generate_response
 from app.mlops_logger import log_prompt
+from PIL import Image
 
+# === Kiểm duyệt prompt ===
 def handle_prompt(prompt):
     safe, info = is_prompt_safe(prompt)
     if not safe:
@@ -13,13 +15,29 @@ def handle_prompt(prompt):
     log_prompt(prompt, "OK", True, response)
     return "✅ Prompt an toàn", response
 
-demo = gr.Interface(
-    fn=handle_prompt,
-    inputs=gr.Textbox(label="Nhập Prompt"),
-    outputs=[
-        gr.Textbox(label="Trạng thái kiểm duyệt"),
-        gr.Textbox(label="Kết quả GenAI")
-    ],
-    title="SAIFGuard: GenAI Prompt Checker",
-    description="Kiểm tra prompt an toàn và tạo văn bản bằng GPT2"
-)
+# === Kiểm duyệt ảnh ===
+def check_image_safety(image: Image.Image):
+    safe, reasons = is_image_safe(image)
+    if safe:
+        return f"✅ Ảnh an toàn: {', '.join(reasons)}"
+    else:
+        return f"🚨 Ảnh KHÔNG an toàn: {', '.join(reasons)}"
+
+# === Giao diện ===
+with gr.Blocks(title="SAIFGuard") as demo:
+    gr.Markdown("## 🛡️ SAIFGuard")
+    
+    with gr.Tab("📝 Kiểm duyệt Prompt"):
+        prompt_input = gr.Textbox(label="Nhập Prompt")
+        prompt_status = gr.Textbox(label="Trạng thái kiểm duyệt")
+        prompt_output = gr.Textbox(label="Kết quả GenAI")
+        prompt_button = gr.Button("Kiểm tra Prompt")
+        prompt_button.click(handle_prompt, inputs=prompt_input, outputs=[prompt_status, prompt_output])
+
+    with gr.Tab("🖼️ Kiểm duyệt Hình ảnh"):
+        image_input = gr.Image(type="pil", label="Tải ảnh lên")
+        image_output = gr.Textbox(label="Trạng thái kiểm duyệt hình ảnh")
+        image_button = gr.Button("Kiểm tra Hình ảnh")
+        image_button.click(fn=check_image_safety, inputs=image_input, outputs=image_output)
+
+demo.launch()

app/safety_check.py

@@ -1,9 +1,21 @@
 from detoxify import Detoxify
+from transformers import AutoProcessor, AutoModelForImageClassification , ViTForImageClassification, ViTFeatureExtractor
+from PIL import Image
+import torch
 
-# Load model một lần
+# ==== Load models ====
+# Load mô hình kiểm duyệt ảnh bạo lực
+violence_model = ViTForImageClassification.from_pretrained('jaranohaal/vit-base-violence-detection')
+violence_processor = ViTFeatureExtractor.from_pretrained('jaranohaal/vit-base-violence-detection')
+
+# Load mô hình kiểm duyệt văn bản
 detox_model = Detoxify('original')
 
-# Hàm kiểm tra prompt
+# Load mô hình kiểm duyệt ảnh
+image_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection")
+image_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+
+# ==== Hàm kiểm duyệt prompt ====
 def is_prompt_safe(prompt: str):
     results = detox_model.predict(prompt)
     
@@ -14,3 +26,32 @@ def is_prompt_safe(prompt: str):
     if flagged:
         return False, list(flagged.keys())
     return True, []
+
+# ==== Hàm kiểm duyệt Hình ảnh ====
+def is_image_safe(image: Image.Image):
+    reasons = []
+
+    # --- Kiểm tra NSFW ---
+    nsfw_inputs = image_processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        nsfw_outputs = image_model(**nsfw_inputs)
+    nsfw_logits = nsfw_outputs.logits
+    nsfw_pred = nsfw_logits.argmax(-1).item()
+    nsfw_label = image_model.config.id2label[nsfw_pred]
+    if nsfw_label.lower() in ["porn", "hentai", "sexy"]:
+        reasons.append(f"Khiêu dâm ({nsfw_label})")
+
+    # --- Kiểm tra Bạo lực ---
+    violence_inputs = violence_processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        violence_outputs = violence_model(**violence_inputs)
+    violence_logits = violence_outputs.logits
+    violence_pred = violence_logits.argmax(-1).item()
+    violence_label = violence_model.config.id2label[violence_pred]
+    if violence_label.lower() in ["violence", "bloody", "weapon", "fight"]:
+        reasons.append(f"Bạo lực ({violence_label})")
+
+    # --- Kết quả ---
+    if reasons:
+        return False, reasons
+    return True, [f"An toàn ({nsfw_label}, {violence_label})"]

requirements.txt

@@ -4,3 +4,4 @@ torch
 protobuf
 presidio-analyzer
 detoxify
+Pillow