app.py

import gradio as gr
import os
import spaces
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from PIL import Image
import numpy as np
import traceback
from typing import Any, Optional

import utils
from utils import BoundingBox
import blurnonymize

os.environ['HF_HOME'] = '/data/.huggingface'

MODEL_NAME = "cborg/qwen2.5VL-3b-privacydetector"
MAX_NEW_TOKENS = 2048
TEMPERATURE = 1.0
MIN_P = 0.1
SYSTEM_PROMPT = """You are a helpful assistant for privacy analysis of images. Please always answer in English. Please obey the users instructions and follow the provided format."""
DEFAULT_PROMPT = """You are an expert at pixel perfect image analysis and in privacy. Your task is to find all private data in the image and report its position, as well as explanations as to why it is private data. Private data is all data that relates to a unique person and can be used to identify them.

First write down your thoughts within a <think> block.
Please go through all objects in the image and consider whether they are private data or not.
End this with a </think> block.

After going through everything, output your findings in an <output></output> block as a json list with the following keys:
{"label": <|object_ref_start|>str<|object_ref_end|>, "description": str, "explanation": str, "bounding_box": <|box_start|>[x_min, y_min, x_max, y_max]<|box_end|>, "severity": int}

Some things to remember:

- private data is only data thats linked to a human person, common examples being a persons face, name, address, license plate
- whenever something can be used to identify a unique human person, it is private data
- report sensitive data as well, such as a nude person
- Severity is a number between 0 and 10, with 0 being not private data and 10 being extremely sensitive private data.
- don't report items which dont contain private data in the final output, you may mention them in your thoughts
- animals and animal faces are not personal data, so a giraffe or a dog is not private data
- you can use whatever format you want within the <think> </think> blocks
- only output valid JSON in between the <output> </output> blocks, adhering to the schema provided
- output the bounding box always as an array of form [x_min, y_min, x_max, y_max]
- private data have a severity greater than 0, so a human face would have severity 6
- go through the image step by step and report the private data, its better to be a bit too sensitive than to miss anything
- put the bounding boxes around the human's face and not the entire person when reporting people as personal data
- if something has been blurred out, or is very blurry and therefore not recognizable, do not report it as private data
- Think step by step, take your time.

Here is the image to analyse, start your analysis directly after:
"""


def build_messages(image, history: Optional[list[dict[str, Any]]] = None, prompt: Optional[str] = None, system_prompt_text: Optional[str] = None):
    if not prompt:
        prompt = DEFAULT_PROMPT
    if not system_prompt_text:
        system_prompt_text = SYSTEM_PROMPT # Fallback if not provided

    if history:
        return [
            *history,
            {"role": "user", "content": [{"type": "text", "text": prompt}]},
        ]

    return [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": system_prompt_text, # Use the passed system prompt
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image", "image": image},
            ],
        },
    ]


# --- Model Loading ---
# Load model using unsloth for 4-bit quantization
try:
    # default: Load the model on the available device(s)
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True
    ).to("cuda").eval()
    tokenizer = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)

    model.to("cuda").eval()  # Ensure model is on GPU and in eval mode

    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    print(traceback.format_exc())
    # Optionally raise or handle the error to prevent app launch if model fails
    raise gr.Error(f"Failed to load model {MODEL_NAME}. Check logs. Error: {e}")


# --- Core Processing Function ---
@spaces.GPU(duration=30)
def anonymise_image(input_image_np: np.ndarray, boxes: list[BoundingBox]):
    # --- Blurnonymizer Instance ---
    try:
        blurnonymizer_instance = blurnonymize.ImageBlurnonymizer()
        return blurnonymizer_instance.censor_image_blur_easy(
            input_image_np, boxes, method="segmentation", verbose=False # Set verbose as needed
        )
    except Exception as e:
        print(f"Error initializing Blurnonymizer: {e}")
        print(traceback.format_exc())
        raise gr.Error(f"Failed to initialize Blurnonymizer. Check logs. Error: {e}")


def run_model_inference(input_image_pil: Image.Image, prompt_text: str, system_prompt_text: str):
    """
    Runs model inference on the input image and prompt.
    """

    # 1. Run Model Inference
    print("Running model inference...")
    messages = build_messages(
        input_image_pil,
        prompt=prompt_text,
        system_prompt_text=system_prompt_text # Pass system prompt here
    )

    input_text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = tokenizer(
        text=[input_text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    out_tokens = model.generate(
        **inputs,
        max_new_tokens=MAX_NEW_TOKENS,
        use_cache=True,
        temperature=TEMPERATURE,
        min_p=MIN_P,
    )
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, out_tokens)
    ]
    raw_model_output = tokenizer.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )[0]

    input_height = inputs['image_grid_thw'][0][1]*14
    input_width = inputs['image_grid_thw'][0][2]*14

    if input_height != input_image_pil.height:
        print("[!] tokenized image height differs from actual height:")
        print(f"Actual: {input_image_pil.height}, processed: {input_height}")

    if input_width != input_image_pil.width:
        print("[!] tokenized image width differs from actual width:")
        print(f"Actual: {input_image_pil.width}, processed: {input_width}")

    print("[+] Model inference completed.")
    print("[*] Raw output:")
    print(raw_model_output)

    return raw_model_output, input_height, input_width


@spaces.GPU(duration=90)  # Request GPU for this function, allow up to 120 seconds
def analyze_image(input_image_pil: Image.Image, prompt_text: Optional[str] = None, system_prompt_text: Optional[str] = None):
    """
    Analyzes the input image using the VLM, visualizes findings, and anonymizes.
    """
    if input_image_pil is None:
        raise gr.Error("Please upload an image.")
    
    # Use default prompts if none are provided
    final_prompt_text = prompt_text if prompt_text else DEFAULT_PROMPT
    final_system_prompt_text = system_prompt_text if system_prompt_text else SYSTEM_PROMPT

    try:
        raw_model_output, image_height, image_width = run_model_inference(input_image_pil, final_prompt_text, final_system_prompt_text)
    except Exception as e:
        print(f"Error during model inference: {e}")
        print(traceback.format_exc())
        raise gr.Error(f"Model inference failed: {e}")

    visualized_image_np, anonymized_image_np = perform_anonymisation(input_image_pil, raw_model_output)

    return raw_model_output, visualized_image_np, anonymized_image_np


@spaces.GPU(duration=90)
def perform_anonymisation(input_image_pil: Image.Image, raw_model_output: str) -> tuple[np.ndarray, np.ndarray]:
    original_image_np = np.array(input_image_pil)

    try:
        print("Parsing findings...")
        # Use the provided utility functions
        parsed_findings = utils.parse_into_models(
            utils.parse_json_response(raw_model_output), strict=False
        )
        print(f"[+] Parsed {len(parsed_findings)} findings.")
        if not parsed_findings:
            print("[*] No findings were parsed from the model output.")

    except Exception as e:
        print(f"Error parsing model output: {e}")
        print(traceback.format_exc())
        # Don't raise error here, allow visualization/anonymization steps to proceed if possible
        # or return early with only original image if parsing is critical
        gr.Warning(
            f"Could not parse findings from model output: {e}. Visualization and anonymization might be incomplete."
        )
        # Fallback: visualize/anonymize based on empty findings list if needed
        parsed_findings = []  # Ensure it's an empty list for downstream steps

    # Initialize boxes_for_viz before the try block
    boxes_for_viz = []
    try:
        # 3. Visualize Findings
        print("Visualizing findings...")
        if parsed_findings:
            # Convert Findings to BoundingBox for visualization function
            boxes_for_viz = [BoundingBox.from_finding(f) for f in parsed_findings]
            # Ensure image is in the correct format (np array) for visualize_boxes_annotated
            visualized_image_np = utils.visualize_boxes_annotated(
                original_image_np, boxes_for_viz
            )
            print("Visualization generated.")
        else:
            print("No findings to visualize, using original image.")
            visualized_image_np = (
                original_image_np.copy()
            )  # Show original if no findings

    except Exception as e:
        print(f"Error during visualization: {e}")
        print(traceback.format_exc())
        gr.Warning(f"Failed to visualize findings: {e}")
        visualized_image_np = original_image_np.copy()  # Fallback to original

    try:
        # 4. Anonymize Image
        print("Anonymizing image...")
        # Use the blurnonymize function with the raw output (as it might contain info needed by the func)
        # Ensure image is numpy array
        # Check if boxes_for_viz is populated before calling anonymise_image
        if boxes_for_viz:
            anonymized_image_np = anonymise_image(original_image_np, boxes_for_viz)
            print("Anonymization generated.")
        else:
            print("No boxes found for anonymization, using original image.")
            anonymized_image_np = original_image_np.copy()

    except Exception as e:
        print(f"Error during anonymization: {e}")
        print(traceback.format_exc())
        gr.Warning(f"Failed to anonymize image: {e}")
        anonymized_image_np = original_image_np.copy()  # Fallback to original

    return visualized_image_np, anonymized_image_np


# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Private Data Detection & Anonymization UI")
    gr.Markdown(f"Using model: `{MODEL_NAME}` on ZeroGPU.")

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Upload Image")
            system_prompt_input = gr.Textbox(label="System Prompt", value=SYSTEM_PROMPT, lines=5, interactive=True) # New system prompt input
            prompt_textbox = gr.Textbox(
                label="Analysis Prompt", value=DEFAULT_PROMPT, lines=4,
            )
            analyze_button = gr.Button("Analyze Image")
        with gr.Column(scale=2):
            with gr.Column():
                raw_output = gr.TextArea(
                    label="Raw Model Output", interactive=True, show_copy_button=True,
                )
                output_visualized = gr.Image(
                    label="Detected Privacy Findings", type="numpy", interactive=False
                )
                output_anonymized = gr.Image(
                    label="Anonymized", type="numpy", interactive=False
                )
                re_anonymise = gr.Button("Anonymise based off manual edits")

    analyze_button.click(
        fn=analyze_image,
        inputs=[input_image, prompt_textbox, system_prompt_input], # Add system_prompt_input here
        outputs=[raw_output, output_visualized, output_anonymized],
    )

    re_anonymise.click(
        fn = perform_anonymisation,
        inputs=[input_image, raw_output],
        outputs=[output_visualized, output_anonymized],
    )

# --- Launch App ---
if __name__ == "__main__":
    demo.queue().launch(
        debug=True
    )  # Enable queue for handling multiple requests, debug mode for logs