app.py

import torch
from PIL import Image
import requests
from openai import OpenAI
from transformers import (Owlv2Processor, Owlv2ForObjectDetection,
                          AutoProcessor, AutoModelForMaskGeneration)
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import base64
import io
import numpy as np
import gradio as gr
import json
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')


def resize_and_compress(image, max_width=800, max_height=800, quality=50):
    """Resize (if > max_width/height) and compress the image to keep Base64 under ~1MB."""
    if not isinstance(image, Image.Image):
        raise ValueError("Input must be a PIL Image")

    width, height = image.size
    if width > max_width or height > max_height:
        aspect_ratio = width / height
        if aspect_ratio > 1:
            new_width = max_width
            new_height = int(new_width / aspect_ratio)
        else:
            new_height = max_height
            new_width = int(new_height * aspect_ratio)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

    buffered = io.BytesIO()
    # Save as JPEG with reduced quality
    image.save(buffered, format="JPEG", quality=quality)
    buffered.seek(0)
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def analyze_image(image):
    client = OpenAI(api_key=OPENAI_API_KEY)

    # Step 1: Resize + compress to keep the Base64 string under 1 MB
    base64_image = resize_and_compress(image, max_width=800, max_height=800, quality=50)

    # Build the list-of-dicts prompt
    prompt_dict = [
        {
            "type": "text",
            "text": """Your task is to determine if the image is surprising or not. 
            If the image is surprising, which element is surprising (max 6 words). 
            Otherwise, 'NA'. Also rate how surprising (1-5). 
            Return JSON like:
            {
              "label": "[surprising or not surprising]",
              "element": "[element]",
              "rating": [1-5]
            }
            """
        },
        {
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{base64_image}"
            }
        }
    ]

    # JSON-encode to ensure content is a string
    json_prompt = json.dumps(prompt_dict)

    # Send request
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": json_prompt
            }
        ],
        max_tokens=100,
        temperature=0.1,
        response_format={"type": "json_object"}
    )

    return response.choices[0].message.content



def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([1.0, 0.0, 0.0, 0.5])

    if len(mask.shape) == 4:
        mask = mask[0, 0]

    mask_image = np.zeros((*mask.shape, 4), dtype=np.float32)
    mask_image[mask > 0] = color

    ax.imshow(mask_image)


def process_image_detection(image, target_label, surprise_rating):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Get original image DPI and size
    original_dpi = image.info.get('dpi', (72, 72))
    original_size = image.size

    # Calculate relative font size based on image dimensions
    base_fontsize = min(original_size) / 40  # Adjust this divisor as needed

    owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
    owlv2_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").to(device)

    sam_processor = AutoProcessor.from_pretrained("facebook/sam-vit-base")
    sam_model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-base").to(device)

    image_np = np.array(image)

    inputs = owlv2_processor(text=[target_label], images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = owlv2_model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]]).to(device)
    results = owlv2_processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]

    dpi = 300  # Increased DPI for better text rendering
    figsize = (original_size[0] / dpi, original_size[1] / dpi)
    fig = plt.figure(figsize=figsize, dpi=dpi)

    ax = plt.Axes(fig, [0., 0., 1., 1.])
    fig.add_axes(ax)
    plt.imshow(image)

    scores = results["scores"]
    if len(scores) > 0:
        max_score_idx = scores.argmax().item()
        max_score = scores[max_score_idx].item()

        if max_score > 0.2:
            box = results["boxes"][max_score_idx].cpu().numpy()

            sam_inputs = sam_processor(
                image,
                input_boxes=[[[box[0], box[1], box[2], box[3]]]],
                return_tensors="pt"
            ).to(device)

            with torch.no_grad():
                sam_outputs = sam_model(**sam_inputs)

            masks = sam_processor.image_processor.post_process_masks(
                sam_outputs.pred_masks.cpu(),
                sam_inputs["original_sizes"].cpu(),
                sam_inputs["reshaped_input_sizes"].cpu()
            )

            mask = masks[0].numpy() if isinstance(masks[0], torch.Tensor) else masks[0]
            show_mask(mask, ax=ax)

            # Draw rectangle around the detected area
            rect = patches.Rectangle(
                (box[0], box[1]),
                box[2] - box[0],
                box[3] - box[1],
                linewidth=max(2, min(original_size) / 500),
                edgecolor='red',
                facecolor='none'
            )
            ax.add_patch(rect)

            # Confidence score
            plt.text(
                box[0], box[1] - base_fontsize,
                f'{max_score:.2f}',
                color='red',
                fontsize=base_fontsize,
                fontweight='bold',
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2)
            )

            # Label + rating
            plt.text(
                box[2] + base_fontsize / 2, box[1],
                f'Unexpected (Rating: {surprise_rating}/5)\n{target_label}',
                color='red',
                fontsize=base_fontsize,
                fontweight='bold',
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2),
                verticalalignment='bottom'
            )

    plt.axis('off')

    # Save figure to buffer
    buf = io.BytesIO()
    plt.savefig(
        buf,
        format='png',
        dpi=dpi,
        bbox_inches='tight',
        pad_inches=0,
        metadata={'dpi': original_dpi}
    )
    buf.seek(0)
    plt.close()

    # Convert buffer back to PIL
    output_image = Image.open(buf)
    output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)

    final_buf = io.BytesIO()
    output_image.save(final_buf, format='PNG', dpi=original_dpi)
    final_buf.seek(0)

    return final_buf


def process_and_analyze(image):
    if image is None:
        return None, "Please upload an image first."

    if OPENAI_API_KEY is None:
        return None, "OpenAI API key not found in environment variables."

    try:
        # Handle different input types
        if isinstance(image, tuple):
            image = image[0]
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        if not isinstance(image, Image.Image):
            raise ValueError("Invalid image format")

        # Analyze image with GPT
        gpt_response = analyze_image(image)
        response_data = json.loads(gpt_response)

        # If surprising, try to detect the element
        if response_data["label"].lower() == "surprising" and response_data["element"].lower() != "na":
            result_buf = process_image_detection(image, response_data["element"], response_data["rating"])
            result_image = Image.open(result_buf)
            analysis_text = (
                f"Label: {response_data['label']}\n"
                f"Element: {response_data['element']}\n"
                f"Rating: {response_data['rating']}/5"
            )
            return result_image, analysis_text
        else:
            # If not surprising or element=NA
            return image, "Not Surprising"

    except Exception as e:
        return None, f"Error processing image: {str(e)}"


def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Image Surprise Analysis")

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Upload Image")
                analyze_btn = gr.Button("Analyze Image")

            with gr.Column():
                output_image = gr.Image(label="Processed Image")
                output_text = gr.Textbox(label="Analysis Results")

        analyze_btn.click(
            fn=process_and_analyze,
            inputs=[input_image],
            outputs=[output_image, output_text]
        )

    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch()