import os
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
from pdf2image import convert_from_path
from PIL import Image
import easyocr
import uuid
import re
import difflib
import math
from faker import Faker
import datetime
import random
import csv
import tempfile
from io import StringIO
# Initialize Faker
fake = Faker()
# Conditional import for LLM
try:
from llama_cpp import Llama
LLAMA_AVAILABLE = True
except ImportError:
print("Warning: llama_cpp not available. LLM functionality will be disabled.")
LLAMA_AVAILABLE = False
# --- Configuration ---
# Folders for temporary files and results
UPLOAD_FOLDER = 'static/uploads/'
RESULTS_FOLDER = 'static/results/'
# --- Model Paths (Update these paths if necessary) ---
CUSTOM_MODEL_PATH = 'best.pt'
PRETRAINED_MODEL_PATH = 'yolov10s.pt'
SIGNATURE_MODEL_PATH = 'yolov8s.pt'
LLAMA_MODEL_PATH = "unsloth.F16.gguf"
# Detection Parameters
YOLO_CONFIDENCE_THRESHOLD = 0.5
OCR_CONFIDENCE_THRESHOLD = 0.5
# Create directories if they don't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULTS_FOLDER, exist_ok=True)
# --- Global Model Placeholders ---
custom_model, pretrained_model, signature_model, reader, llama_model = None, None, None, None, None
def load_models():
"""
Loads all AI models into the global scope. This function is called on the first
analysis request to avoid startup conflicts. It ensures models are only loaded once.
"""
global custom_model, pretrained_model, signature_model, reader, llama_model
# If models are already loaded, do nothing.
if reader is not None and (llama_model is not None or not LLAMA_AVAILABLE):
print("Models already loaded.")
return
print("=== Loading Models (this may take a moment) ===")
# Helper function to check for model files
def check_model_path(path, name):
if not os.path.exists(path):
print(f"✗ WARNING: {name} model not found at '{path}'. The application may not function correctly.")
return False
return True
# YOLO Models
if check_model_path(CUSTOM_MODEL_PATH, "Custom YOLO"):
try:
custom_model = YOLO(CUSTOM_MODEL_PATH)
print("✓ Custom YOLO model loaded.")
except Exception as e:
print(f"✗ Error loading custom model: {e}")
if check_model_path(PRETRAINED_MODEL_PATH, "Pre-trained YOLO"):
try:
pretrained_model = YOLO(PRETRAINED_MODEL_PATH)
print("✓ Pre-trained YOLO model loaded.")
except Exception as e:
print(f"✗ Error loading pre-trained model: {e}")
if check_model_path(SIGNATURE_MODEL_PATH, "Signature YOLO"):
try:
signature_model = YOLO(SIGNATURE_MODEL_PATH)
print("✓ Signature YOLO model loaded.")
except Exception as e:
print(f"✗ Error loading signature model: {e}")
# OCR Model
try:
reader = easyocr.Reader(['en'], gpu=True)
print("✓ EasyOCR model loaded.")
except Exception as e:
print(f"✗ Error loading EasyOCR: {e}. Text detection will be unavailable.")
# LLM Model - Only load if available
if LLAMA_AVAILABLE and check_model_path(LLAMA_MODEL_PATH, "LLM"):
try:
llama_model = Llama(
model_path=LLAMA_MODEL_PATH,
n_gpu_layers=-1, n_ctx=4096, chat_format="llama-3", verbose=False
)
print("✓ LLM model loaded.")
except Exception as e:
print(f"✗ Error loading LLM model: {e}. Text analysis will be unavailable.")
print("=== All Models Initialized ===")
# YOLO Class Mappings
CUSTOM_CLASS_NAMES = {0: 'face', 1: 'qr', 2: 'signature'}
PRETRAINED_CLASS_MAP = {0: 'face'}
# --- Core Detection & Processing Functions ---
def detect_visual_pii(image_data):
"""Runs the three-stage YOLO detection on a single image."""
all_boxes = []
all_classes = []
if custom_model is None:
print("Custom model not available for visual detection")
return all_boxes, all_classes
# Pass 1: Custom Model
custom_results = custom_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
detected_custom_classes = {CUSTOM_CLASS_NAMES[int(cls)] for cls in custom_results.boxes.cls}
for box, cls in zip(custom_results.boxes.xyxy.cpu().numpy().astype(int), custom_results.boxes.cls):
all_boxes.append(box)
all_classes.append(CUSTOM_CLASS_NAMES[int(cls)])
# Pass 2: Pre-trained Model (Face Fallback)
if 'face' not in detected_custom_classes and pretrained_model is not None:
print(" Custom model missed 'face'. Trying pre-trained model as fallback.")
pretrained_results = pretrained_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
for box in pretrained_results.boxes:
if int(box.cls[0]) in PRETRAINED_CLASS_MAP:
all_boxes.append(box.xyxy.cpu().numpy().astype(int)[0])
all_classes.append("face (fallback)")
# Pass 3: Specialized Model (Signature Fallback)
if 'signature' not in detected_custom_classes and signature_model is not None:
print(" Custom model missed 'signature'. Trying specialized signature model as fallback.")
signature_results = signature_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
for box in signature_results.boxes:
all_boxes.append(box.xyxy.cpu().numpy().astype(int)[0])
all_classes.append("signature (fallback)")
return all_boxes, all_classes
# --- OCR + LLM Functions ---
def calculate_distance(bbox1, bbox2):
"""Calculates the Euclidean distance between the centers of two bounding boxes."""
c1_x = (bbox1[0] + bbox1[2]) / 2
c1_y = (bbox1[1] + bbox1[3]) / 2
c2_x = (bbox2[0] + bbox2[2]) / 2
c2_y = (bbox2[1] + bbox2[3]) / 2
return math.sqrt((c2_x - c1_x)**2 + (c2_y - c1_y)**2)
def refine_pii_flags(ocr_results, isolation_threshold=150):
"""Post-processing step to unmark short, isolated PII detections."""
pii_indices = [i for i, result in enumerate(ocr_results) if result["is_pii"]]
if len(pii_indices) <= 1:
return ocr_results
indices_to_unmark = []
for i in pii_indices:
current_result = ocr_results[i]
normalized_text = re.sub(r'[^a-zA-Z0-9]', '', current_result["text"])
if len(normalized_text) <= 3:
min_dist_to_neighbor = float('inf')
for j in pii_indices:
if i == j:
continue
other_result = ocr_results[j]
dist = calculate_distance(current_result["bbox"], other_result["bbox"])
if dist < min_dist_to_neighbor:
min_dist_to_neighbor = dist
if min_dist_to_neighbor > isolation_threshold:
print(f" - Refining PII: Unmarking short ('{current_result['text']}') and isolated (min_dist: {min_dist_to_neighbor:.2f}px) PII.")
indices_to_unmark.append(i)
for i in indices_to_unmark:
ocr_results[i]["is_pii"] = False
return ocr_results
def parse_pii_output(generated_text):
"""Parse the new curly braces format PII output"""
pii_list = []
try:
match = re.search(r'\{([^}]*)\}', generated_text)
if match:
content = match.group(1)
items = re.findall(r'"([^"]*)"', content)
pii_list = [item.strip() for item in items if item.strip()]
except Exception as e:
print(f"Error parsing PII output: {e}")
pii_list = []
return pii_list
def normalize_text(text):
"""Comprehensive text normalization for better matching"""
if not text:
return ""
normalized = re.sub(r'[.,;:!?()"\'\-_/\\]', '', text)
ocr_corrections = {
'0': 'o', 'O': '0', '1': 'l', 'l': '1', '5': 's', 'S': '5',
'8': 'b', 'B': '8', 'rn': 'm', 'RN': 'M', 'vv': 'w', 'VV': 'W',
'cl': 'd', 'CL': 'D',
}
for wrong, correct in ocr_corrections.items():
normalized = normalized.replace(wrong, correct)
normalized = ' '.join(normalized.split()).lower()
return normalized
def fuzzy_match_score(text1, text2, threshold=0.8):
"""Calculate fuzzy matching score between two strings"""
if not text1 or not text2:
return False
return difflib.SequenceMatcher(None, text1.lower(), text2.lower()).ratio() >= threshold
def levenshtein_distance(s1, s2):
"""Calculate Levenshtein distance between two strings"""
if len(s1) < len(s2):
return levenshtein_distance(s2, s1)
if len(s2) == 0:
return len(s1)
previous_row = list(range(len(s2) + 1))
for i, c1 in enumerate(s1):
current_row = [i + 1]
for j, c2 in enumerate(s2):
insertions = previous_row[j + 1] + 1
deletions = current_row[j] + 1
substitutions = previous_row[j] + (c1 != c2)
current_row.append(min(insertions, deletions, substitutions))
previous_row = current_row
return previous_row[-1]
def is_similar_by_edit_distance(text1, text2, max_distance=2):
"""Check if two texts are similar within edit distance threshold"""
if not text1 or not text2:
return False
distance = levenshtein_distance(text1.lower(), text2.lower())
max_len = max(len(text1), len(text2))
if max_len <= 3:
threshold = 1
elif max_len <= 6:
threshold = 2
else:
threshold = min(max_distance, max_len // 3)
return distance <= threshold
def extract_sentence_text(ocr_results):
"""Extract sentence-based text for LLM input"""
paragraph_text = ""
for item in ocr_results:
if len(item) == 3:
_, text, _ = item
elif len(item) == 2:
_, text = item
else:
print(f"Unexpected OCR result format: {item}")
continue
if text.strip():
paragraph_text += text + " "
return paragraph_text.strip()
def extract_word_bboxes_improved(ocr_results):
"""Improved word extraction with better handling of punctuation and spacing"""
word_bbox_map = []
for item in ocr_results:
if len(item) == 3:
bbox, text, confidence = item
elif len(item) == 2:
bbox, text = item
confidence = 1.0
else:
print(f"Unexpected OCR result format: {item}")
continue
original_text = text.strip()
if not original_text:
continue
if isinstance(bbox[0], (list, tuple)):
x_coords = [point[0] for point in bbox]
y_coords = [point[1] for point in bbox]
line_x1, line_y1 = min(x_coords), min(y_coords)
line_x2, line_y2 = max(x_coords), max(y_coords)
else:
line_x1, line_y1, line_x2, line_y2 = bbox
tokens = re.findall(r'\S+', original_text)
if len(tokens) <= 1:
padding = 1
word_bbox_map.append({
"word": original_text,
"bbox": [
max(0, int(line_x1 - padding)),
max(0, int(line_y1 - padding)),
int(line_x2 + padding),
int(line_y2 + padding)
],
"confidence": confidence,
"original_line": original_text
})
continue
full_width = line_x2 - line_x1
text_without_spaces = original_text.replace(' ', '')
total_chars = len(text_without_spaces)
char_position = 0
for i, token in enumerate(tokens):
token_start_ratio = char_position / total_chars if total_chars > 0 else 0
char_position += len(token)
token_end_ratio = char_position / total_chars if total_chars > 0 else 1
token_x1 = line_x1 + (full_width * token_start_ratio)
token_x2 = line_x1 + (full_width * token_end_ratio)
padding = 1
word_bbox = [
max(0, int(token_x1 - padding)),
max(0, int(line_y1 - padding)),
int(min(token_x2 + padding, line_x2)),
int(line_y2 + padding)
]
word_bbox_map.append({
"word": token,
"bbox": word_bbox,
"confidence": confidence,
"original_line": original_text
})
return sorted(word_bbox_map, key=lambda x: (x['bbox'][1], x['bbox'][0]))
def advanced_match_pii_to_words(pii_list, word_bbox_map):
"""Advanced multi-strategy PII matching with comprehensive fallbacks"""
ocr_results_for_template = []
words = [info['word'] for info in word_bbox_map]
bboxes = [info['bbox'] for info in word_bbox_map]
is_pii_flags = [False] * len(words)
# Pre-process all words with different normalization strategies
normalized_words = [normalize_text(word) for word in words]
print(f"Processing {len(pii_list)} PII items against {len(words)} OCR words")
for pii_idx, pii_item in enumerate(pii_list):
if not pii_item.strip():
continue
print(f"Processing PII item {pii_idx + 1}: '{pii_item}'")
# Normalize the PII item
normalized_pii = normalize_text(pii_item)
pii_words = normalized_pii.split()
if not pii_words:
continue
matched = False
# Strategy 1: Exact matching after normalization
if len(pii_words) == 1:
pii_word = pii_words[0]
for idx, norm_word in enumerate(normalized_words):
if norm_word == pii_word and not is_pii_flags[idx]:
is_pii_flags[idx] = True
matched = True
print(f" ✓ Exact match: '{words[idx]}' -> '{pii_item}'")
else:
# Multi-word exact matching
pii_len = len(pii_words)
start_idx = 0
while start_idx < len(normalized_words) - pii_len + 1:
exact_match = True
for j in range(pii_len):
if normalized_words[start_idx + j] != pii_words[j]:
exact_match = False
break
if exact_match:
# Check spatial proximity
spatial_ok = True
for j in range(1, pii_len):
prev_bbox = bboxes[start_idx + j - 1]
curr_bbox = bboxes[start_idx + j]
horizontal_distance = curr_bbox[0] - prev_bbox[2]
vertical_alignment = (abs(prev_bbox[1] - curr_bbox[1]) < 30 and
abs(prev_bbox[3] - curr_bbox[3]) < 30)
if not (vertical_alignment and horizontal_distance <= 150):
spatial_ok = False
break
if spatial_ok:
for j in range(pii_len):
if not is_pii_flags[start_idx + j]:
is_pii_flags[start_idx + j] = True
matched = True
matched_text = ' '.join(words[start_idx:start_idx + pii_len])
print(f" ✓ Multi-word exact: '{matched_text}' -> '{pii_item}'")
start_idx += pii_len
continue
start_idx += 1
# Strategy 2: Fuzzy matching if exact matching failed
if not matched:
if len(pii_words) == 1:
pii_word = pii_words[0]
for idx, norm_word in enumerate(normalized_words):
if (not is_pii_flags[idx] and
(fuzzy_match_score(norm_word, pii_word, 0.9) or
is_similar_by_edit_distance(norm_word, pii_word, 2))):
is_pii_flags[idx] = True
matched = True
print(f" ✓ Fuzzy match: '{words[idx]}' -> '{pii_item}'")
else:
# Multi-word fuzzy matching
pii_len = len(pii_words)
start_idx = 0
while start_idx < len(normalized_words) - pii_len + 1:
fuzzy_match = True
for j in range(pii_len):
if not (fuzzy_match_score(normalized_words[start_idx + j], pii_words[j], 0.85) or
is_similar_by_edit_distance(normalized_words[start_idx + j], pii_words[j], 2)):
fuzzy_match = False
break
if fuzzy_match:
# Check spatial proximity
spatial_ok = True
for j in range(1, pii_len):
prev_bbox = bboxes[start_idx + j - 1]
curr_bbox = bboxes[start_idx + j]
horizontal_distance = curr_bbox[0] - prev_bbox[2]
vertical_alignment = (abs(prev_bbox[1] - curr_bbox[1]) < 30 and
abs(prev_bbox[3] - curr_bbox[3]) < 30)
if not (vertical_alignment and horizontal_distance <= 150):
spatial_ok = False
break
if spatial_ok:
for j in range(pii_len):
if not is_pii_flags[start_idx + j]:
is_pii_flags[start_idx + j] = True
matched = True
matched_text = ' '.join(words[start_idx:start_idx + pii_len])
print(f" ✓ Multi-word fuzzy: '{matched_text}' -> '{pii_item}'")
start_idx += pii_len
continue
start_idx += 1
# Strategy 3: Substring and partial matching
if not matched:
full_normalized_text = ' '.join(normalized_words)
pos = 0
while True:
start_pos = full_normalized_text.find(normalized_pii, pos)
if start_pos == -1:
break
end_pos = start_pos + len(normalized_pii)
char_count = 0
start_word_idx = None
end_word_idx = None
for idx, norm_word in enumerate(normalized_words):
word_start = char_count
word_end = char_count + len(norm_word)
if start_word_idx is None and word_end > start_pos:
start_word_idx = idx
if word_start < end_pos:
end_word_idx = idx
char_count += len(norm_word) + 1
if start_word_idx is not None and end_word_idx is not None and end_word_idx - start_word_idx + 1 >= len(pii_words):
spatial_ok = True
for j in range(start_word_idx, end_word_idx):
if j + 1 <= end_word_idx:
prev_bbox = bboxes[j]
next_bbox = bboxes[j + 1]
horizontal_distance = next_bbox[0] - prev_bbox[2]
vertical_alignment = (abs(prev_bbox[1] - next_bbox[1]) < 30 and
abs(prev_bbox[3] - next_bbox[3]) < 30)
if not (vertical_alignment and horizontal_distance <= 200):
spatial_ok = False
break
if spatial_ok:
for j in range(start_word_idx, end_word_idx + 1):
if not is_pii_flags[j]:
is_pii_flags[j] = True
matched = True
matched_text = ' '.join(words[start_word_idx:end_word_idx + 1])
print(f" ✓ Substring match: '{matched_text}' -> '{pii_item}'")
pos = end_pos
# Strategy 4: Individual word matching with relaxed criteria
if not matched:
for pii_word in pii_words:
if len(pii_word) < 3:
continue
for idx, norm_word in enumerate(normalized_words):
if not is_pii_flags[idx]:
if (norm_word == pii_word or
fuzzy_match_score(norm_word, pii_word, 0.8) or
is_similar_by_edit_distance(norm_word, pii_word, 2) or
(len(pii_word) > 5 and (pii_word in norm_word or norm_word in pii_word))):
is_pii_flags[idx] = True
print(f" ✓ Individual word match: '{words[idx]}' -> '{pii_word}' from '{pii_item}'")
if not matched:
print(f" ✗ No match found for: '{pii_item}'")
for idx, word_info in enumerate(word_bbox_map):
ocr_results_for_template.append({
"text": word_info["word"],
"bbox": word_info["bbox"],
"is_pii": is_pii_flags[idx],
"confidence": word_info.get("confidence", 1.0)
})
ocr_results_for_template = merge_horizontal_pii_boxes_improved(ocr_results_for_template)
return ocr_results_for_template
def merge_horizontal_pii_boxes_improved(ocr_results, merge_distance=50):
"""Improved merging with better spatial awareness and tighter boxes"""
if not ocr_results:
return ocr_results
merged_results = []
i = 0
while i < len(ocr_results):
current_word = ocr_results[i]
if not current_word["is_pii"]:
merged_results.append(current_word)
i += 1
continue
merge_group = [current_word]
j = i + 1
while j < len(ocr_results):
next_word = ocr_results[j]
if not next_word["is_pii"]:
break
current_bbox = merge_group[-1]["bbox"]
next_bbox = next_word["bbox"]
y_center_current = (current_bbox[1] + current_bbox[3]) / 2
y_center_next = (next_bbox[1] + next_bbox[3]) / 2
y_overlap = abs(y_center_current - y_center_next) < 20
horizontal_distance = next_bbox[0] - current_bbox[2]
if y_overlap and horizontal_distance <= merge_distance:
merge_group.append(next_word)
j += 1
else:
break
if len(merge_group) > 1:
min_x = min(word["bbox"][0] for word in merge_group)
min_y = min(word["bbox"][1] for word in merge_group)
max_x = max(word["bbox"][2] for word in merge_group)
max_y = max(word["bbox"][3] for word in merge_group)
merged_text = " ".join(word["text"] for word in merge_group)
merged_word = {
"text": merged_text,
"bbox": [min_x, min_y, max_x, max_y],
"is_pii": True,
"confidence": max(word.get("confidence", 1.0) for word in merge_group)
}
merged_results.append(merged_word)
print(f" ✓ Merged PII box: '{merged_text}' at [{min_x},{min_y},{max_x},{max_y}]")
else:
merged_results.append(current_word)
i = j
return merged_results
def post_process_pii_detection(ocr_results_for_template, pii_list):
"""Post-process to catch any missed PII using relaxed matching"""
words = [result["text"] for result in ocr_results_for_template]
for pii_item in pii_list:
normalized_pii = normalize_text(pii_item)
pii_words = normalized_pii.split()
if not pii_words:
continue
pii_detected = False
for result in ocr_results_for_template:
if result["is_pii"]:
result_normalized = normalize_text(result["text"])
if (normalized_pii in result_normalized or
result_normalized in normalized_pii or
fuzzy_match_score(result_normalized, normalized_pii, 0.7)):
pii_detected = True
break
if not pii_detected:
print(f" ⚠ PII not detected, trying fallback matching: '{pii_item}'")
for idx, result in enumerate(ocr_results_for_template):
if result["is_pii"]:
continue
word_normalized = normalize_text(result["text"])
for pii_word in pii_words:
if (len(pii_word) > 3 and
(pii_word in word_normalized or
word_normalized in pii_word or
fuzzy_match_score(word_normalized, pii_word, 0.6) or
is_similar_by_edit_distance(word_normalized, pii_word, 3))):
ocr_results_for_template[idx]["is_pii"] = True
print(f" ✓ Fallback match: '{result['text']}' -> '{pii_word}' from '{pii_item}'")
break
return ocr_results_for_template
def detect_pii_from_combined_text(combined_text):
"""Detect PII from combined multi-page text using LLM"""
if llama_model is None:
print("LLM model not available for PII detection")
return [], "LLM model not available"
instruction = (
"Extract all Personally Identifiable Information (PII) of the main subject from the given text. "
"Include data like Name, Date of Birth, Gender, Address, Phone Number, Email, Social Security Number (SSN), Member ID, Group Number, or any other PII data available. "
"Ignore any information about doctors, staff, providers, colleagues, organizations, companies, hospitals, educational institutes, or facilities. "
"Return the results strictly as a flat set of strings enclosed in { } without labels."
)
prompt_content = f"{instruction}\n{combined_text}"
pii_list = []
llama_raw_output = ""
try:
messages = [{"role": "user", "content": prompt_content}]
response = llama_model.create_chat_completion(
messages=messages,
max_tokens=512,
temperature=0.1,
)
llama_raw_output = response['choices'][0]['message']['content']
pii_list = parse_pii_output(llama_raw_output)
print(f"LLM detected {len(pii_list)} PII items from combined text: {pii_list}")
except Exception as e:
print(f"Error during Llama PII detection: {e}")
llama_raw_output = f"Error: {str(e)}"
pii_list = []
return pii_list, llama_raw_output
# --- Combined Processing Function ---
def process_page_combined(img_cv, global_pii_list):
"""Process a single page with both YOLO and OCR+LLM detection"""
all_detections = []
# Step 1: YOLO Visual Detection
print(" Running YOLO visual detection...")
visual_boxes, visual_classes = detect_visual_pii(img_cv)
for box, cls in zip(visual_boxes, visual_classes):
all_detections.append({
"text": cls,
"bbox": box.tolist() if hasattr(box, 'tolist') else box,
"is_pii": True,
"confidence": 1.0,
"detection_type": "visual"
})
print(f" YOLO detected {len(visual_boxes)} visual elements")
# Step 2: OCR + LLM Text Detection
if reader is not None:
print(" Running OCR text extraction...")
word_ocr_results = reader.readtext(img_cv, paragraph=False, width_ths=0.7, height_ths=0.7)
if word_ocr_results:
word_bbox_map = extract_word_bboxes_improved(word_ocr_results)
ocr_results_for_template = advanced_match_pii_to_words(global_pii_list, word_bbox_map)
ocr_results_for_template = post_process_pii_detection(ocr_results_for_template, global_pii_list)
ocr_results_for_template = refine_pii_flags(ocr_results_for_template)
ocr_results_for_template = merge_horizontal_pii_boxes_improved(ocr_results_for_template)
for result in ocr_results_for_template:
if result["is_pii"]:
result["detection_type"] = "text"
all_detections.append(result)
print(f" OCR detected {sum(1 for r in ocr_results_for_template if r['is_pii'])} text PII elements")
return all_detections
def classify_pii(text):
text = text.strip()
clean_text = re.sub(r'\s+', '', text)
if re.match(r'^\d{3}-\d{2}-\d{4}$', text) or re.match(r'^\d{3}-\d{2}-\d{4}$', clean_text):
return 'ssn'
elif re.match(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', text) or re.match(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', clean_text) or '@' in text:
return 'email'
elif re.match(r'^\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}$', text) or re.match(r'^\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}$', clean_text):
return 'phone'
elif re.match(r'^\d{1,2}/\d{1,2}/\d{4}$', text) or re.match(r'^\d{4}-\d{2}-\d{2}$', text) or re.match(r'^\d{1,2}/\d{1,2}/\d{4}$', clean_text) or re.match(r'^\d{1,2}-\d{1,2}-\d{4}$', text) or re.match(r'^\d{1,2}-\d{1,2}-\d{2}$', text) or re.match(r'^\d{2}/\d{2}/\d{4}$', text) or re.match(r'^\d{2}-\d{2}-\d{4}$', text) or re.match(r'^(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{1,2},\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s+\d{1,2},\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^\d{1,2}\s+(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^\d{1,2}\s+(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s+\d{4}$', text, re.IGNORECASE):
return 'dob'
elif re.match(r'^(?=.*\d)[A-Za-z0-9]+$', text) and len(text) > 5:
return 'id'
elif ',' in text or 'St' in text or 'Ave' in text or re.search(r'\d{5}', text):
return 'address'
elif re.match(r'^(male|female|m|f|transgender|nonbinary|non-binary|other|unknown|u|o)$', text.lower()):
return 'gender'
else:
return 'name'
def format_gender(base_gender, original_text):
orig = original_text.strip()
orig_lower = orig.lower()
if orig_lower not in ['male', 'female', 'm', 'f', 'transgender', 'nonbinary', 'non-binary', 'other', 'unknown', 'u', 'o']:
return base_gender.capitalize()
if len(orig) == 1:
char = 'M' if base_gender == 'male' else 'F'
return char.lower() if orig.islower() else char
else:
if orig.isupper():
return base_gender.upper()
elif orig.islower():
return base_gender.lower()
else:
return base_gender.capitalize()
def generate_fake(pii_type, length, original=None):
max_attempts = 100
if pii_type == 'dob' and original:
formats = [
'%m/%d/%Y', '%m/%d/%y', '%d/%m/%Y', '%d/%m/%y', '%Y-%m-%d', '%y-%m-%d', '%m-%d-%Y', '%m-%d-%y',
'%d-%m-%Y', '%d-%m-%y', '%Y/%m/%d', '%y/%m/%d', '%d.%m.%Y', '%m.%d.%Y', '%B %d, %Y', '%b %d, %Y',
'%d %B %Y', '%d %b %Y', '%B %d %Y', '%b %d %Y'
]
for fmt in formats:
try:
datetime.datetime.strptime(original.strip(), fmt)
fake_dt = fake.date_object()
return fake_dt.strftime(fmt)
except ValueError:
pass
return fake.date(pattern='%m/%d/%Y')
elif pii_type == 'gender' and original:
return format_gender(random.choice(['male', 'female']), original)
elif pii_type in ['name', 'email', 'phone', 'address']:
for _ in range(max_attempts):
if pii_type == 'name': f = fake.name()
elif pii_type == 'email': f = fake.email()
elif pii_type == 'phone': f = fake.phone_number()
elif pii_type == 'address': f = fake.address().replace('\n', ', ')
if len(f) == length:
return f
closest = None
min_diff = float('inf')
for _ in range(50):
if pii_type == 'name': f = fake.name()
elif pii_type == 'email': f = fake.email()
elif pii_type == 'phone': f = fake.phone_number()
elif pii_type == 'address': f = fake.address().replace('\n', ', ')
diff = abs(len(f) - length)
if diff < min_diff:
min_diff, closest = diff, f
return closest
elif pii_type == 'ssn':
return fake.ssn()
else: # id and others
return fake.lexify(text='?' * length)
def detect_format(text):
text = text.strip().strip('"')
if text.startswith('ISA*'):
return 'edi_x12'
else:
return 'plain'
def convert_to_readable(text, format_type):
if format_type == 'edi_x12':
parsed_transactions = parse_edi_fallback(text)
output = StringIO()
format_output(parsed_transactions, file=output)
return output.getvalue()
else:
return text
def redact_text(original_text, pii_list):
redacted = original_text
for pii in pii_list:
redacted = re.sub(re.escape(pii), '[REDACTED]', redacted, flags=re.IGNORECASE)
return redacted
def anonymize_text(original_text, pii_list):
anonymized = original_text
pii_map = {}
for pii in pii_list:
normalized = normalize_text(pii)
pii_type = classify_pii(pii)
if normalized not in pii_map:
if pii_type == 'gender':
base_gender = random.choice(['male', 'female'])
fake_val = format_gender(base_gender, pii)
pii_map[normalized] = fake_val
else:
fake_val = generate_fake(pii_type, len(pii), pii)
pii_map[normalized] = fake_val
else:
fake_val = pii_map[normalized]
anonymized = re.sub(re.escape(pii), fake_val, anonymized, flags=re.IGNORECASE)
return anonymized
def parse_edi_fallback(edi_content):
"""
Fallback parser that combines address parts into a single line for easier redaction.
"""
print("Using fallback parser...")
segments = edi_content.replace('~', '\n').split('\n')
segments = [seg.strip() for seg in segments if seg.strip()]
parsed_data = {
'transaction_info': {}, 'patient_info': {}, 'provider_info': {},
'service_info': {}, 'diagnosis_info': {}
}
for segment in segments:
elements = segment.split('*')
segment_id = elements[0]
if segment_id == 'ST':
parsed_data['transaction_info']['transaction_type'] = elements[1]
elif segment_id == 'NM1':
entity_type = elements[1]
if entity_type == 'IL': # Patient
last_name = elements[3] if len(elements) > 3 else ''
first_name = elements[4] if len(elements) > 4 else ''
parsed_data['patient_info']['name'] = f"{first_name} {last_name}".strip()
if len(elements) > 8:
parsed_data['patient_info']['id'] = elements[9]
elif entity_type == 'SJ': # Provider
last_name = elements[3] if len(elements) > 3 else ''
first_name = elements[4] if len(elements) > 4 else ''
parsed_data['provider_info']['name'] = f"{first_name} {last_name}".strip()
if len(elements) > 8:
parsed_data['provider_info']['npi'] = elements[9]
elif segment_id == 'N3':
# Store the first line of the address
parsed_data['patient_info']['address'] = elements[1]
elif segment_id == 'N4':
# Combine City, State, and Zip with the address line
city = elements[1] if len(elements) > 1 else ''
state = elements[2] if len(elements) > 2 else ''
zip_code = elements[3] if len(elements) > 3 else ''
full_address_parts = [city, state, zip_code]
# If an address line already exists from N3, prepend it
if 'address' in parsed_data['patient_info']:
full_address_parts.insert(0, parsed_data['patient_info']['address'])
# Join all parts with ", " and filter out any empty parts
parsed_data['patient_info']['address'] = ", ".join(filter(None, full_address_parts))
elif segment_id == 'DMG':
parsed_data['patient_info']['dob'] = elements[2]
parsed_data['patient_info']['gender'] = 'Female' if elements[3] == 'F' else 'Male'
elif segment_id == 'UM':
parsed_data['service_info']['service_type'] = elements[1]
parsed_data['service_info']['request_category'] = elements[2]
parsed_data['service_info']['service_code'] = elements[3]
if len(elements) > 4:
parsed_data['service_info']['quantity'] = elements[4]
elif segment_id == 'HI':
diagnosis_info = elements[1].split(':')
if len(diagnosis_info) > 1:
parsed_data['diagnosis_info']['code_qualifier'] = diagnosis_info[0]
parsed_data['diagnosis_info']['diagnosis_code'] = diagnosis_info[1]
return [{
'transaction_type': parsed_data['transaction_info'].get('transaction_type', 'Unknown'),
'parsed_data': {
'description': 'Health Care Services Review',
'patient_info': parsed_data['patient_info'],
'provider_info': parsed_data['provider_info'],
'service_info': parsed_data['service_info'],
'diagnosis_info': parsed_data['diagnosis_info']
}
}]
def format_output(parsed_transactions, file=None):
"""Format parsed data for display"""
output_lines = []
for i, transaction in enumerate(parsed_transactions):
output_lines.append(f"\n=== TRANSACTION {i+1} ===")
output_lines.append(f"Transaction Type: {transaction['transaction_type']}")
if 'parsed_data' in transaction:
data = transaction['parsed_data']
if 'description' in data:
output_lines.append(f"Description: {data['description']}")
# Patient Information
if 'patient_info' in data and data['patient_info']:
output_lines.append("\nPATIENT INFORMATION:")
for key, value in data['patient_info'].items():
output_lines.append(f" {key.replace('_', ' ').title()}: {value}")
# Provider Information
if 'provider_info' in data and data['provider_info']:
output_lines.append("\nPROVIDER INFORMATION:")
for key, value in data['provider_info'].items():
output_lines.append(f" {key.replace('_', ' ').title()}: {value}")
# Service Information
if 'service_info' in data and data['service_info']:
output_lines.append("\nSERVICE INFORMATION:")
for key, value in data['service_info'].items():
output_lines.append(f" {key.replace('_', ' ').title()}: {value}")
# Diagnosis Information
if 'diagnosis_info' in data and data['diagnosis_info']:
output_lines.append("\nDIAGNOSIS INFORMATION:")
for key, value in data['diagnosis_info'].items():
output_lines.append(f" {key.replace('_', ' ').title()}: {value}")
output_str = '\n'.join(output_lines)
if file:
file.write(output_str)
else:
print(output_str)
return output_str
# --- Main Gradio Processing Function ---
def analyze_document(file, progress=gr.Progress()):
"""
This function takes an uploaded file, processes it through the PII detection pipeline,
and returns the annotated images, a redacted PDF, and a summary report.
"""
load_models()
if file is None:
return None, None, None, None, None, "Please upload a document to begin."
unique_id = uuid.uuid4().hex
if hasattr(file, 'name'):
filepath = file.name
else:
filepath = str(file)
filename = os.path.basename(filepath)
extension = os.path.splitext(filename)[1]
if extension.lower() == '.csv':
rows = []
with open(filepath, 'r', newline='') as csvfile:
csv_reader = csv.reader(csvfile)
for row in csv_reader:
if row:
rows.append(row[0])
total_rows = len(rows)
print(f"Processing {total_rows} rows for job {unique_id}...")
progress(0.1, desc="Reading CSV rows...")
report = f"## 🔍 Analysis Report for CSV\n**Total Rows:** {total_rows}\n\n---\n"
redacted_rows = []
anonymized_rows = []
for i, text in enumerate(rows):
progress(0.4 + (i / total_rows * 0.5), desc=f"Processing Row {i+1}/{total_rows}...")
format_type = detect_format(text)
readable_text = convert_to_readable(text, format_type)
pii_list, llama_raw_output = detect_pii_from_combined_text(readable_text)
redacted_text = redact_text(readable_text, pii_list)
anonymized_text = anonymize_text(readable_text, pii_list)
redacted_rows.append(redacted_text)
anonymized_rows.append(anonymized_text)
report += f"### 📄 Row {i+1}\n- **Text Detections:** {len(pii_list)}\n- **PII Found:** {', '.join(pii_list) if pii_list else 'None'}\n\n"
progress(0.9, desc="Generating final CSVs...")
redacted_csv_path = os.path.join(RESULTS_FOLDER, f"redacted_{unique_id}.csv")
with open(redacted_csv_path, 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
for txt in redacted_rows:
writer.writerow([txt])
anonymized_csv_path = os.path.join(RESULTS_FOLDER, f"anonymized_{unique_id}.csv")
with open(anonymized_csv_path, 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
for txt in anonymized_rows:
writer.writerow([txt])
progress(1, desc="Complete!")
print("Processing Complete.")
return [], [], [], redacted_csv_path, anonymized_csv_path, report
progress(0, desc="Converting document to images...")
images_to_process = []
try:
if extension.lower() == '.pdf':
try:
images_to_process = [cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR) for page in convert_from_path(filepath, dpi=300)]
except Exception as e:
print(f"PDF conversion error: {e}. Trying fallback method...")
try:
images_to_process = [cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR) for page in convert_from_path(filepath, dpi=150)]
except Exception as e2:
return None, None, None, None, None, f"🔴 **Error:** Could not process PDF. Please ensure Poppler is installed.\nDetails: {e2}"
else:
img = cv2.imread(filepath)
if img is not None:
images_to_process.append(img)
except Exception as e:
return None, None, None, None, None, f"🔴 **Error:** Could not process file. Details: {e}"
if not images_to_process:
return None, None, None, None, None, "🔴 **Error:** No pages could be extracted from the document."
total_pages = len(images_to_process)
print(f"Processing {total_pages} pages for job {unique_id}...")
progress(0.1, desc="Extracting text from all pages (OCR)...")
combined_text, all_pages_data = "", []
for i, img_cv in enumerate(images_to_process):
page_text = ""
if reader:
page_text = extract_sentence_text(reader.readtext(img_cv, paragraph=True))
combined_text += f"\n--- Page {i+1} ---\n{page_text}\n"
all_pages_data.append({"img_cv": img_cv, "page_num": i + 1})
progress(0.4, desc="Analyzing text for PII with LLM...")
global_pii_list, llama_raw_output = detect_pii_from_combined_text(combined_text)
annotated_paths, redacted_paths, anonymized_paths = [], [], []
redacted_pils, anonymized_pils = [], []
report = f"## 🔍 Analysis Report\n**Global PII Found:** `{', '.join(global_pii_list) if global_pii_list else 'None'}`\n\n---\n"
pii_map = {}
for i, page_info in enumerate(all_pages_data):
progress(0.5 + (i / total_pages * 0.4), desc=f"Processing Page {i+1}/{total_pages} (Visual & Text)...")
img_cv, page_num = page_info["img_cv"], page_info["page_num"]
detections = process_page_combined(img_cv, global_pii_list)
annotated_img = img_cv.copy()
redacted_img = img_cv.copy()
anonymized_img = img_cv.copy()
visual_count = sum(1 for d in detections if d["detection_type"] == "visual")
text_count = sum(1 for d in detections if d.get("detection_type") == "text")
for d in detections:
bbox = d.get("bbox", [])
if not bbox: continue
x1, y1, x2, y2 = map(int, bbox)
color = (0, 255, 0) if d.get("detection_type") == "visual" else (0, 0, 255)
cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 3)
if d["detection_type"] == "visual":
cv2.rectangle(redacted_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
cv2.rectangle(anonymized_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
else:
bg_color = (255, 255, 255)
height, width = img_cv.shape[:2]
if x2 + 20 < width:
sample = img_cv[y1:y2, x2:x2+20]
if sample.size > 0:
bg_color = tuple(map(int, np.mean(sample, axis=(0,1))))
else:
if x1 > 20:
sample = img_cv[y1:y2, x1-20:x1]
if sample.size > 0:
bg_color = tuple(map(int, np.mean(sample, axis=(0,1))))
cv2.rectangle(redacted_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
cv2.rectangle(anonymized_img, (x1, y1), (x2, y2), bg_color, -1)
original_text = d["text"]
pii_type = classify_pii(original_text)
normalized = normalize_text(original_text)
key = 'gender' if pii_type == 'gender' else normalized
if key not in pii_map:
if pii_type == 'gender':
base_gender = random.choice(['male', 'female'])
fake_text = format_gender(base_gender, original_text)
pii_map[key] = base_gender
else:
fake_text = generate_fake(pii_type, len(original_text), original_text)
pii_map[key] = fake_text
else:
if pii_type == 'gender':
base_gender = pii_map[key]
fake_text = format_gender(base_gender, original_text)
else:
fake_text = pii_map[key]
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = (y2 - y1) / 40.0
thickness = 2
text_size, _ = cv2.getTextSize(fake_text, font, font_scale, thickness)
box_width = x2 - x1
if text_size[0] > box_width - 10:
font_scale *= (box_width - 10) / text_size[0]
text_size, _ = cv2.getTextSize(fake_text, font, font_scale, thickness)
text_x = x1 + (box_width - text_size[0]) // 2
text_y = y1 + ((y2 - y1) + text_size[1]) // 2
bg_brightness = 0.299 * bg_color[2] + 0.587 * bg_color[1] + 0.114 * bg_color[0]
text_color = (0, 0, 0) if bg_brightness > 128 else (255, 255, 255)
cv2.putText(anonymized_img, fake_text, (text_x, text_y), font, font_scale, text_color, thickness)
# Save all three versions of the image for the galleries
annotated_path = os.path.join(RESULTS_FOLDER, f"annotated_{unique_id}_{page_num}.jpg")
redacted_path = os.path.join(RESULTS_FOLDER, f"redacted_preview_{unique_id}_{page_num}.jpg")
anonymized_path = os.path.join(RESULTS_FOLDER, f"anonymized_preview_{unique_id}_{page_num}.jpg")
cv2.imwrite(annotated_path, annotated_img)
cv2.imwrite(redacted_path, redacted_img)
cv2.imwrite(anonymized_path, anonymized_img)
annotated_paths.append(annotated_path)
redacted_paths.append(redacted_path)
anonymized_paths.append(anonymized_path)
redacted_pils.append(Image.fromarray(cv2.cvtColor(redacted_img, cv2.COLOR_BGR2RGB)))
anonymized_pils.append(Image.fromarray(cv2.cvtColor(anonymized_img, cv2.COLOR_BGR2RGB)))
report += f"### 📄 Page {page_num}\n- **Visual Detections (🟩 Green):** {visual_count}\n- **Text Detections (🟥 Red):** {text_count}\n"
progress(0.9, desc="Generating final PDFs...")
redacted_pdf_path, anonymized_pdf_path = None, None
if redacted_pils:
pdf_path = os.path.join(RESULTS_FOLDER, f"redacted_{unique_id}.pdf")
redacted_pils[0].save(pdf_path, "PDF", resolution=100.0, save_all=True, append_images=redacted_pils[1:])
redacted_pdf_path = pdf_path
if anonymized_pils:
pdf_path = os.path.join(RESULTS_FOLDER, f"anonymized_{unique_id}.pdf")
anonymized_pils[0].save(pdf_path, "PDF", resolution=100.0, save_all=True, append_images=anonymized_pils[1:])
anonymized_pdf_path = pdf_path
progress(1, desc="Complete!")
print("Processing Complete.")
return annotated_paths, redacted_paths, anonymized_paths, redacted_pdf_path, anonymized_pdf_path, report
# --- Gradio Interface Definition ---
title = "🔒 Combined PII Detection System"
description = """
### Advanced Multi-Modal PII Detection
This system uses a combination of visual and textual analysis to detect and redact Personally Identifiable Information from your documents.
- **🖼️ Visual Detection (YOLO):** Detects Faces, QR Codes, and Signatures.
- **📝 Text Detection (OCR + LLM):** Detects Names, Addresses, Phone Numbers, IDs, and other contextual PII.
**How to Use:**
1. Upload a document (PDF or image format).
2. The system will process each page and display annotated previews with colored boxes.
3. A fully redacted PDF with blacked-out PII is generated for you to download.
4. An analysis report summarizes the findings for each page.
"""
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(f"{title}
")
gr.Markdown(
""
"Upload a PDF or image to detect and redact PII using visual detectors (🟩) and text analysis (🟥). "
"Fixed for local environment with proper YOLO support."
"
"
)
with gr.Accordion("About this tool", open=False):
gr.Markdown(description)
with gr.Tabs():
with gr.Tab("Run"):
with gr.Row():
with gr.Column(scale=1):
file_input = gr.File(
label="Upload Document",
file_types=['.pdf', '.jpg', '.jpeg', '.png', '.bmp', '.csv'],
file_count="single",
height=100
)
submit_btn = gr.Button("🚀 Analyze Document", variant="primary")
with gr.Accordion("Tips", open=False):
gr.Markdown(
"- Prefer high-resolution files for better OCR results (300 DPI for PDFs).\n"
"- For PDFs, ensure Poppler is installed on your system.\n"
"- Visual detections are drawn in green; text-based detections are in red.\n"
"- Use the Previews tab to inspect annotated pages and the Report tab to download the redacted PDF."
)
with gr.Column(scale=1):
gr.Markdown("### What happens during analysis")
gr.Markdown(
"- Convert pages to images\n"
"- Run global OCR to build combined text\n"
"- Use LLM to extract possible PII strings\n"
"- Match PII back to words and merge boxes\n"
"- Render annotated previews and build a redacted PDF"
)
clear_btn = gr.Button("🧹 Clear Results", variant="secondary")
with gr.Tab("Annotated Preview (Detection)"):
gr.Markdown("### Annotated Previews (🟩 Visual, 🟥 Text)")
annotated_gallery_output = gr.Gallery(
label="Annotated Pages", show_label=False, elem_id="gallery_annotated",
columns=[2], rows=[1], object_fit="contain", height=480
)
with gr.Tab("Redacted Preview"):
gr.Markdown("### Redacted Previews (Blacked Out)")
redacted_gallery_output = gr.Gallery(
label="Redacted Pages", show_label=False, elem_id="gallery_redacted",
columns=[2], rows=[1], object_fit="contain", height=480
)
with gr.Tab("Anonymized Preview"):
gr.Markdown("### Anonymized Previews (Fake Data)")
anonymized_gallery_output = gr.Gallery(
label="Anonymized Pages", show_label=False, elem_id="gallery_anonymized",
columns=[2], rows=[1], object_fit="contain", height=480
)
with gr.Tab("Report & Download"):
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Downloads")
redacted_file_output = gr.File(label="Redacted PDF (Blacked Out)")
anonymized_file_output = gr.File(label="Anonymized PDF (Fake Data)")
with gr.Column(scale=2):
gr.Markdown("### Analysis Report")
report_output = gr.Markdown(label="Analysis Report")
outputs_list = [
annotated_gallery_output,
redacted_gallery_output,
anonymized_gallery_output,
redacted_file_output,
anonymized_file_output,
report_output
]
submit_btn.click(
fn=analyze_document,
inputs=file_input,
outputs=outputs_list
)
clear_btn.click(
fn=lambda: ([], [], [], None, None, "Ready. Upload a document and click Analyze."),
inputs=None,
outputs=outputs_list
)
if __name__ == "__main__":
demo.queue().launch(server_port=8000)