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 encode_image_to_base64(image):
    print(f"Encode image type: {type(image)}")  # Debug print
    
    try:
        # If image is a tuple (as sometimes provided by Gradio), take the first element
        if isinstance(image, tuple):
            print(f"Image is tuple with length: {len(image)}")  # Debug print
            if len(image) > 0 and image[0] is not None:
                if isinstance(image[0], np.ndarray):
                    image = Image.fromarray(image[0])
                else:
                    image = image[0]
            else:
                raise ValueError("Invalid image tuple provided")

        # If image is a numpy array, convert to PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # If image is a path string, open it
        elif isinstance(image, str):
            image = Image.open(image)

        print(f"Image type after conversion: {type(image)}")  # Debug print

        # Ensure image is in PIL Image format
        if not isinstance(image, Image.Image):
            raise ValueError(f"Input must be a PIL Image, numpy array, or valid image path. Got {type(image)}")

        # Convert image to RGB if it's in RGBA mode
        if image.mode == 'RGBA':
            image = image.convert('RGB')

        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    except Exception as e:
        print(f"Encode error details: {str(e)}")  # Debug print
        raise
        

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

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": """Your task is to determine if the image is surprising or not surprising.    
                    if the image is surprising, determine which element, figure or object in the image is making the image surprising and write it only in one sentence with no more then 6 words, otherwise, write 'NA'.    
                    Also rate how surprising the image is on a scale of 1-5, where 1 is not surprising at all and 5 is highly surprising.
                    Additionally, write one sentence about what would be expected in this scene, and one sentence about what is unexpected.
                    Provide the response as a JSON with the following structure:    
                    {
                        "label": "[surprising OR not surprising]",
                        "element": "[element]",
                        "rating": [1-5],
                        "expected": "[one sentence about what would be expected]",
                        "unexpected": "[one sentence about what is unexpected]"
                    }"""
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                }
            ]
        }
    ]

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        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):
    try:
        # Debug print to understand mask type
        print(f"show_mask input type: {type(mask)}")
        
        # Convert mask if it's a tuple
        if isinstance(mask, tuple):
            if len(mask) > 0 and mask[0] is not None:
                mask = mask[0]
            else:
                raise ValueError("Invalid mask tuple")

        # Convert torch tensor to numpy if needed
        if torch.is_tensor(mask):
            mask = mask.cpu().numpy()
            
        # Handle 4D tensor/array case
        if len(mask.shape) == 4:
            mask = mask[0, 0]
        # Handle 3D tensor/array case
        elif len(mask.shape) == 3:
            mask = mask[0]

        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])

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

        ax.imshow(mask_image)
        
    except Exception as e:
        print(f"show_mask error: {str(e)}")
        print(f"mask shape: {getattr(mask, 'shape', 'no shape')}")
        raise


def process_image_detection(image, target_label, surprise_rating):
   try:
       # Handle different image input types
       if isinstance(image, tuple):
           if len(image) > 0 and image[0] is not None:
               if isinstance(image[0], np.ndarray):
                   image = Image.fromarray(image[0])
               else:
                   image = image[0]
           else:
               raise ValueError("Invalid image tuple provided")
       elif isinstance(image, np.ndarray):
           image = Image.fromarray(image)
       elif isinstance(image, str):
           image = Image.open(image)

       # Ensure image is in PIL Image format
       if not isinstance(image, Image.Image):
           raise ValueError(f"Input must be a PIL Image, got {type(image)}")

       # Ensure image is in RGB mode
       if image.mode != 'RGB':
           image = image.convert('RGB')

       device = "cuda" if torch.cuda.is_available() else "cpu"
       print(f"Using device: {device}")

       # Get original image DPI and size
       original_dpi = image.info.get('dpi', (72, 72))
       original_size = image.size
       print(f"Image size: {original_size}")

       # Calculate relative font size
       base_fontsize = min(original_size) / 80

       print("Loading models...")
       owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-large-patch14")
       owlv2_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-large-patch14").to(device)
       sam_processor = AutoProcessor.from_pretrained("facebook/sam-vit-large")
       sam_model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-large").to(device)

       print("Running object detection...")
       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]

       # Use original image dimensions for figure size
       dpi = 300
       width, height = image.size
       figsize = (width / dpi, height / dpi)
       
       fig = plt.figure(figsize=figsize, dpi=dpi)
       ax = plt.Axes(fig, [0., 0., 1., 1.])
       fig.add_axes(ax)
       ax.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:
               print("Processing detection results...")
               box = results["boxes"][max_score_idx].cpu().numpy()

               print("Running SAM model...")
               # Convert image to numpy array if needed for SAM
               if isinstance(image, Image.Image):
                   image_np = np.array(image)
               else:
                   image_np = image

               sam_inputs = sam_processor(
                   image_np,
                   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()
               )

               print(f"Mask type: {type(masks)}, Mask shape: {len(masks)}")
               mask = masks[0]
               if isinstance(mask, torch.Tensor):
                   mask = mask.numpy()
               
               show_mask(mask, ax=ax)

               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)

               # Only add the probability score
               #plt.text(
               #    box[0], box[1] - base_fontsize,
               #    f'{max_score:.2f}',
               #    color='red',
               #    fontsize=base_fontsize,
               #    fontweight='bold'
               #)

       plt.axis('off')

       print("Saving final image...")
       try:
           buf = io.BytesIO()
           fig.savefig(buf, 
                     format='png',
                     dpi=dpi,
                     bbox_inches='tight',
                     pad_inches=0)
           buf.seek(0)
           
           # Open as PIL Image
           output_image = Image.open(buf)
           
           # Convert to RGB if needed
           if output_image.mode != 'RGB':
               output_image = output_image.convert('RGB')
           
           # Save to final buffer
           final_buf = io.BytesIO()
           output_image.save(final_buf, format='PNG', dpi=original_dpi)
           final_buf.seek(0)
           
           plt.close(fig)
           buf.close()
           
           return final_buf

       except Exception as e:
           print(f"Save error details: {str(e)}")
           raise

   except Exception as e:
       print(f"Process image detection error: {str(e)}")
       print(f"Error occurred at line {e.__traceback__.tb_lineno}")
       raise

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

    print(f"Initial image type: {type(image)}")  # Debug print

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

    try:
        # Convert the image to PIL format if needed
        if isinstance(image, tuple):
            print(f"Image is tuple, length: {len(image)}")  # Debug print
            if len(image) > 0 and image[0] is not None:
                if isinstance(image[0], np.ndarray):
                    image = Image.fromarray(image[0])
                else:
                    print(f"First element type: {type(image[0])}")  # Debug print
                    image = image[0]
            else:
                return None, "Invalid image format provided"
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif isinstance(image, str):
            image = Image.open(image)
        
        print(f"Image type after conversion: {type(image)}")  # Debug print

        if not isinstance(image, Image.Image):
            return None, f"Invalid image format: {type(image)}"

        # Ensure image is in RGB mode
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Analyze image
        print("Starting GPT analysis...")  # Debug print
        gpt_response = analyze_image(image)
        print(f"GPT response: {gpt_response}")  # Debug print
        
        try:
            response_data = json.loads(gpt_response)
        except json.JSONDecodeError:
            return None, "Error: Invalid response format from GPT"

        if not all(key in response_data for key in ["label", "element", "rating"]):
            return None, "Error: Missing required fields in analysis response"

        print(f"Response data: {response_data}")  # Debug print

        if response_data["label"].lower() == "surprising" and response_data["element"].lower() != "na":
            try:
                print("Starting image detection...")  # Debug print
                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\n"
                    f"Expected: {response_data['expected']}\n"
                    f"Unexpected: {response_data['unexpected']}"
            )
                return result_image, analysis_text
            except Exception as detection_error:
                print(f"Detection error details: {str(detection_error)}")  # Debug print
                return None, f"Error in image detection processing: {str(detection_error)}"
        else:
            return image, "Not Surprising"

    except Exception as e:
        error_type = type(e).__name__
        error_msg = str(e)
        detailed_error = f"Error ({error_type}): {error_msg}"
        print(detailed_error)  # Debug print
        return None, f"Error processing image: {error_msg}"

# Create Gradio interface
def create_interface():
    with gr.Blocks() as demo:
        with gr.Row():  # Horizontal layout for the left side alignment
            with gr.Column(scale=1):  # Adjust the scale for the left section
                gr.Image(
                    value="appendix/icon.webp", 
                    width=65, 
                    interactive=False, 
                    show_label=False,
                    show_download_button=False,
                    elem_id="icon"
                )
                with gr.Column(scale=3):
                    gr.Markdown("## Image Anomaly Detection")

        gr.Markdown(
            "This project offers a tool that identifies surprising elements in images, "
            "pinpointing what violates our expectations. It analyzes images for unexpected objects, "
            "locations, social scenarios, settings, and roles."
        )
 
        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]
        )

           
        # Display example images in a row using Gradio Image components
        with gr.Row():
            gr.Image(value="appendix/gradio_example.png", width=250, show_label=False, interactive=False, show_download_button=False)
            gr.Image(value="appendix/gradio_example2.png", width=250, show_label=False, interactive=False, show_download_button=False)
            gr.Image(value="appendix/gradio_example3.png", width=250, show_label=False, interactive=False, show_download_button=False)
            gr.Image(value="appendix/gradio_example4.png", width=250, show_label=False, interactive=False, show_download_button=False)

    
    return demo





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