"""
Gradio UI for Forensic Image Analysis Agent
Clean, full-page chat interface with streaming support
"""
import os
import sys
import json
import logging
from pathlib import Path
from typing import Optional, Tuple
import gradio as gr
# Add src to path
sys.path.insert(0, str(Path(__file__).parent))
from src.agents import ForensicAgent
# Set up logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)
# Initialize agent (will be created on first use)
agent: Optional[ForensicAgent] = None
agent_config: dict = {}
def _ensure_weights():
"""Ensure TruFor and DRUNet weights are downloaded on startup."""
try:
from src.utils.weight_downloader import ensure_trufor_weights
workspace_root = Path(__file__).parent
success, message = ensure_trufor_weights(workspace_root=workspace_root, auto_download=True)
print(f"[Weights] {message}")
except Exception as e:
print(f"[Weights] TruFor: {e}")
try:
from src.utils.weight_downloader import ensure_drunet_weights
workspace_root = Path(__file__).parent
drunet_weights_path = workspace_root / "src" / "tools" / "forensic" / "drunet" / "weights" / "drunet_gray.pth"
success, message = ensure_drunet_weights(weights_path=drunet_weights_path, auto_download=True)
print(f"[Weights] {message}")
except Exception as e:
print(f"[Weights] DRUNet: {e}")
def initialize_agent(
model_name: str = "gpt-5.1",
temperature: float = 0.2,
reasoning_effort: Optional[str] = None,
use_tools: bool = True
) -> Tuple[str, Optional[ForensicAgent]]:
"""Initialize the forensic agent."""
global agent, agent_config
try:
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
return "OPENAI_API_KEY not found. Please set it in Hugging Face Space Secrets.", None
agent = ForensicAgent(
llm_model=model_name,
temperature=temperature,
reasoning_effort=reasoning_effort,
api_key=api_key,
max_iterations=50
)
agent_config = {
'model_name': model_name,
'temperature': temperature,
'reasoning_effort': reasoning_effort,
'use_tools': use_tools
}
return f"Agent ready: {model_name}", agent
except Exception as e:
return f"Error: {str(e)}", None
def format_tool_result(tool_name: str, tool_result) -> str:
"""Format tool result for display."""
if isinstance(tool_result, dict):
# Extract key findings for cleaner display
formatted_lines = []
# Common keys to highlight
if 'manipulation_detected' in tool_result:
status = "Yes" if tool_result['manipulation_detected'] else "No"
formatted_lines.append(f"**Manipulation Detected:** {status}")
if 'confidence' in tool_result:
formatted_lines.append(f"**Confidence:** {tool_result['confidence']:.1%}")
if 'ai_generated_probability' in tool_result:
formatted_lines.append(f"**AI Generated Probability:** {tool_result['ai_generated_probability']:.1%}")
if 'anomaly_score' in tool_result:
formatted_lines.append(f"**Anomaly Score:** {tool_result['anomaly_score']:.3f}")
if 'findings' in tool_result and isinstance(tool_result['findings'], list):
if tool_result['findings']:
formatted_lines.append("**Findings:**")
for finding in tool_result['findings'][:5]: # Limit to 5 findings
formatted_lines.append(f" - {finding}")
if 'error' in tool_result:
formatted_lines.append(f"**Error:** {tool_result['error']}")
# If we extracted useful info, return that; otherwise show JSON
if formatted_lines:
return "\n".join(formatted_lines)
else:
try:
result_str = json.dumps(tool_result, indent=2)
if len(result_str) > 1500:
result_str = result_str[:1500] + "\n..."
return f"```json\n{result_str}\n```"
except:
return str(tool_result)[:1500]
else:
result_str = str(tool_result)
if len(result_str) > 1500:
result_str = result_str[:1500] + "..."
return result_str
def respond_chat(
message,
history: list,
use_tools: bool = True,
model_name: str = "gpt-5.1",
temperature: float = 0.2,
reasoning_effort: Optional[str] = None,
):
"""
Chat function with streaming support.
Args:
message: Dict with "text" and "files" keys (multimodal=True)
history: List of previous messages
use_tools: Whether to use forensic tools
model_name: LLM model name
temperature: Temperature setting
reasoning_effort: Reasoning effort level
"""
global agent
# Extract image and text from message
image_path = None
query_text = ""
if isinstance(message, dict):
query_text = message.get("text", "")
files = message.get("files", [])
if files:
image_path = files[0]
elif isinstance(message, str):
query_text = message
if image_path is None:
yield "Please upload an image to analyze. You can drag & drop or paste an image."
return
if not os.path.exists(image_path):
yield f"Image file not found: {image_path}"
return
try:
# Convert reasoning_effort string to None if needed
reasoning_effort_clean = None if reasoning_effort == "None" else reasoning_effort
# Initialize agent if needed or if settings changed
needs_reinit = (
agent is None or
agent_config.get('model_name') != model_name or
agent_config.get('temperature') != temperature or
agent_config.get('reasoning_effort') != reasoning_effort_clean or
agent_config.get('use_tools') != use_tools
)
if needs_reinit:
status_msg, agent_obj = initialize_agent(
model_name=model_name,
temperature=temperature,
reasoning_effort=reasoning_effort_clean,
use_tools=use_tools
)
if agent_obj is None:
yield f"Failed to initialize agent: {status_msg}"
return
agent = agent_obj
# Streaming state
output_parts = []
current_section = None
tools_used = []
llm_content = []
# Stream analysis
logger.info(f"Starting analysis stream for image: {image_path}, query: {query_text[:50] if query_text else 'None'}...")
for event in agent.analyze_stream(
image_path=image_path,
user_query=query_text if query_text else None,
use_tools=use_tools
):
event_type = event.get('type')
logger.debug(f"[UI] Event received: {event_type}")
if event_type == 'status':
# Don't show transient status messages in output
pass
elif event_type == 'tool_call':
tool_name = event.get('tool_name', 'unknown')
tool_args = event.get('tool_args', {})
logger.info(f"[UI] Tool call received: {tool_name}")
logger.debug(f"[UI] Tool arguments: {tool_args}")
if tool_name not in tools_used:
tools_used.append(tool_name)
# Add tool section
output_parts.append(f"\n\n---\n### Tool: {tool_name}\n*Running...*")
yield ''.join(output_parts)
elif event_type == 'tool_result':
tool_name = event.get('tool_name', 'unknown')
tool_result = event.get('tool_result', '')
logger.info(f"[UI] Tool result received: {tool_name}")
logger.debug(f"[UI] Tool result preview: {str(tool_result)[:200]}...")
# Format the result
formatted_result = format_tool_result(tool_name, tool_result)
# Replace "Running..." with result
current_output = ''.join(output_parts)
search_str = f"### Tool: {tool_name}\n*Running...*"
if search_str in current_output:
replacement = f"### Tool: {tool_name}\n{formatted_result}"
updated_output = current_output.replace(search_str, replacement, 1)
output_parts = [updated_output]
else:
# Add as new section if not found
output_parts.append(f"\n\n---\n### Tool: {tool_name}\n{formatted_result}")
yield ''.join(output_parts)
elif event_type == 'llm_chunk':
chunk = event.get('content', '')
if chunk:
# If this is the first LLM content, add a separator
if not llm_content:
output_parts.append("\n\n---\n### Analysis\n\n")
llm_content.append(chunk)
output_parts.append(chunk)
yield ''.join(output_parts)
elif event_type == 'final':
# Final yield
yield ''.join(output_parts)
# Ensure we have final output
final_output = ''.join(output_parts)
if not final_output.strip():
final_output = "Analysis complete. No detailed output was generated."
yield final_output
except Exception as e:
import traceback
logger.error(f"Error during analysis: {str(e)}", exc_info=True)
yield f"Error during analysis:\n\n```\n{str(e)}\n\n{traceback.format_exc()}\n```"
def create_interface():
"""Create the Gradio interface."""
# Custom CSS for full-page layout and clean styling
custom_css = """
/* Full height layout */
.gradio-container {
max-width: 100% !important;
padding: 0 !important;
}
/* Chat container fills available space */
.chat-container {
height: calc(100vh - 120px) !important;
}
/* Chatbot area */
.chatbot {
height: 100% !important;
overflow-y: auto !important;
}
/* Message styling */
.message {
max-width: 85% !important;
}
/* Sidebar styling */
.sidebar {
background: #f8f9fa !important;
border-right: 1px solid #e0e0e0 !important;
}
/* Tool output code blocks */
.message pre {
max-height: 300px;
overflow-y: auto;
}
/* Header */
.header {
padding: 12px 20px;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
text-align: center;
}
.header h1 {
margin: 0;
font-size: 1.4em;
font-weight: 600;
}
.header p {
margin: 4px 0 0 0;
font-size: 0.9em;
opacity: 0.9;
}
/* Hide default Gradio footer */
footer {
display: none !important;
}
/* Multimodal textbox styling */
.multimodal-textbox {
border-radius: 12px !important;
}
"""
with gr.Blocks() as demo:
# Inject custom CSS (placed at top to ensure it loads early)
gr.HTML(f"", elem_id="custom-css")
# Header
gr.HTML("""
""")
with gr.Row(equal_height=True):
# Sidebar with settings
with gr.Column(scale=1, min_width=280):
gr.Markdown("### Settings")
use_tools = gr.Checkbox(
label="Use Forensic Tools",
value=True,
info="Run JPEG, ELA, frequency analysis, etc."
)
model = gr.Dropdown(
label="Model",
choices=["gpt-5", "gpt-5.1", "gpt-5-mini", "gpt-5-nano"],
value="gpt-5.1"
)
temperature = gr.Slider(
label="Temperature",
minimum=0.0,
maximum=1.0,
value=0.2,
step=0.1
)
reasoning = gr.Dropdown(
label="Reasoning Effort",
choices=["None", "low", "medium", "high"],
value="None"
)
gr.Markdown("---")
gr.Markdown("""
**How to use:**
- Upload or paste an image
- Optionally add a question
- Results stream in real-time
**Tips:**
- Scroll up during analysis to read - it won't force scroll down
- Tool results show key findings
""")
# Main chat area
with gr.Column(scale=4):
chatbot = gr.Chatbot(
height="70vh",
placeholder="Upload an image to begin forensic analysis...",
layout="panel",
autoscroll=True, # Respects user scroll position
)
with gr.Row():
msg = gr.MultimodalTextbox(
placeholder="Drop an image here or type a question about the image...",
show_label=False,
scale=4,
file_count="single",
file_types=["image"],
)
# State to store the last message for bot response
last_message = gr.State(value=None)
def user_message(message, history):
"""Add user message to history and store for bot."""
if not message:
return history, gr.MultimodalTextbox(value=None), None
# Store the message for bot_response
stored_message = message.copy() if isinstance(message, dict) else message
# Format user message for display
user_content = []
if isinstance(message, dict):
if message.get("files"):
for f in message["files"]:
user_content.append({"type": "file", "path": f})
if message.get("text"):
user_content.append({"type": "text", "text": message["text"]})
if not user_content:
return history, gr.MultimodalTextbox(value=None), None
# Add to history
if len(user_content) == 1 and user_content[0].get("type") == "text":
history.append({"role": "user", "content": user_content[0]["text"]})
else:
history.append({"role": "user", "content": user_content})
return history, gr.MultimodalTextbox(value=None), stored_message
def bot_response(stored_message, history, use_tools, model, temperature, reasoning):
"""Generate bot response with streaming."""
if not history or stored_message is None:
return history
# Add empty assistant message
history.append({"role": "assistant", "content": ""})
# Stream response
for chunk in respond_chat(stored_message, history[:-1], use_tools, model, temperature, reasoning):
history[-1]["content"] = chunk
yield history
# Event chain: user submits -> add to history -> generate response
msg.submit(
user_message,
[msg, chatbot],
[chatbot, msg, last_message],
queue=False
).then(
bot_response,
[last_message, chatbot, use_tools, model, temperature, reasoning],
chatbot
)
return demo
if __name__ == "__main__":
_ensure_weights()
status_msg, _ = initialize_agent()
print(status_msg)
demo = create_interface()
# In Gradio 6.0+, theme is passed to launch()
app_theme = gr.themes.Soft(
primary_hue="indigo",
secondary_hue="purple",
neutral_hue="slate",
font=gr.themes.GoogleFont("Inter"),
)
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=False,
theme=app_theme,
)