Lattice / app.py
dosimeter's picture
Update app.py
fa081b0 verified
Raw
History Blame Contribute Delete
5.49 kB
import gradio as gr
import re
import pandas as pd
import os
import cv2
import numpy as np
from PIL import Image
import pytesseract
# Disable the DecompressionBomb limit to allow processing of extremely long screenshots
Image.MAX_IMAGE_PIXELS = None
def clean_and_format_number(phone_str):
"""
Cleans the extracted string to ensure it's a valid phone number.
Removes spaces, dashes, brackets, and names.
"""
# Keep only digits and the plus sign
cleaned = re.sub(r'[^\d+]', '', str(phone_str))
# Extract just the digits to check length
digits_only = cleaned.replace('+', '')
# Most valid international numbers are between 10 and 15 digits long
if 10 <= len(digits_only) <= 15:
# Standardize format: ensure it starts with a '+' for Excel safety
if not cleaned.startswith('+'):
cleaned = '+' + cleaned
return cleaned
return None
def extract_numbers_from_images(image_files):
"""Main processing function triggered by Gradio."""
if not image_files:
return None, "⚠️ Please upload at least one image."
all_unique_numbers = set()
logs =[]
for i, file_obj in enumerate(image_files):
# Gradio 'file' type returns a string path or an object depending on the version.
file_path = file_obj if isinstance(file_obj, str) else file_obj.name
file_basename = os.path.basename(file_path)
logs.append(f"⚡ Processing image {i+1}/{len(image_files)}: {file_basename}...")
try:
# 1. Read image incredibly fast using OpenCV
# We use imdecode to safely handle file paths with special characters
img_array = np.fromfile(file_path, np.uint8)
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
if img is None:
logs.append(f"❌ Error: Could not read image {file_basename}.")
continue
# 2. Convert to Grayscale (Speeds up OCR by 3x and improves contrast)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
height, width = gray.shape
# If the image is extremely wide, resize it down to 1080px width proportionally
if width > 1080:
ratio = 1080.0 / width
new_height = int(height * ratio)
gray = cv2.resize(gray, (1080, new_height), interpolation=cv2.INTER_AREA)
height, width = gray.shape
# 3. Slice image into safe chunks to prevent memory overload on massive screenshots
MAX_HEIGHT = 4000
chunk_count = 0
for y in range(0, height, MAX_HEIGHT):
chunk_count += 1
# Crop the chunk using numpy slicing (instantaneous)
chunk = gray[y:min(y + MAX_HEIGHT, height), :]
# Run Tesseract OCR on the chunk
text = pytesseract.image_to_string(chunk)
# 4. Pure Logic Regex Extraction
# Looks for a sequence of 9-25 chars containing digits, spaces, dashes, or brackets
pattern = r'\+?[\d\s\-\(\)]{9,25}'
potential_matches = re.findall(pattern, text)
valid_count = 0
for match in potential_matches:
clean_num = clean_and_format_number(match)
if clean_num:
all_unique_numbers.add(clean_num)
valid_count += 1
logs.append(f" -> Scanned {chunk_count} chunk(s). Current total unique numbers: {len(all_unique_numbers)}")
except Exception as e:
logs.append(f"❌ Error processing {file_basename}: {str(e)}")
# Check if we got any numbers
if not all_unique_numbers:
return None, "\n".join(logs) + "\n\n⚠️ No phone numbers were extracted from the uploaded images."
# Convert set to sorted list
sorted_numbers = sorted(list(all_unique_numbers))
# Output only the phone numbers, as requested
df = pd.DataFrame({"Phone Number": sorted_numbers})
# Outputting to Excel to permanently prevent the Microsoft CSV Scientific Notation bug
excel_filename = "extracted_whatsapp_numbers.xlsx"
df.to_excel(excel_filename, index=False)
final_status = "\n".join(logs) + f"\n\n🎉 SUCCESS! Instantly extracted {len(sorted_numbers)} total unique numbers."
return excel_filename, final_status
# --- GRADIO INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
"""
)
with gr.Row():
with gr.Column(scale=1):
image_input = gr.File(
file_count="multiple",
type="filepath",
label="Upload WhatsApp Screenshots",
file_types=["image"]
)
submit_btn = gr.Button("Extract Numbers Instantly", variant="primary")
with gr.Column(scale=1):
excel_output = gr.File(label="Download Extracted Numbers (Excel)")
status_output = gr.Textbox(label="Processing Logs & Status", lines=12, interactive=False)
submit_btn.click(
fn=extract_numbers_from_images,
inputs=[image_input],
outputs=[excel_output, status_output]
)
if __name__ == "__main__":
demo.launch()