import gradio as gr
import torch
import numpy as np
from transformers import OwlViTProcessor, OwlViTForObjectDetection, ResNetModel
from torchvision import transforms
from PIL import Image
import cv2
import torch.nn.functional as F
import tempfile
import os

# Load models
resnet = ResNetModel.from_pretrained("Microsoft/resnet-50")
resnet.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
resnet = resnet.to(device)

mixin = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
model = mixin.to(device)

# Preprocess the image
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return transform(image).unsqueeze(0)

def extract_embedding(image):
    image_tensor = preprocess_image(image).to(device)
    with torch.no_grad():
        output = resnet(image_tensor)
        embedding = output.pooler_output
    return embedding

def cosine_similarity(embedding1, embedding2):
    return F.cosine_similarity(embedding1, embedding2)

def l2_distance(embedding1, embedding2):
    return torch.norm(embedding1 - embedding2, p=2)

def save_array_to_temp_image(arr):
    rgb_arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(rgb_arr)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
    temp_file_name = temp_file.name
    temp_file.close()
    img.save(temp_file_name)
    return temp_file_name

def detect_and_crop(target_image, query_image, threshold=0.6, nms_threshold=0.3):
    target_sizes = torch.Tensor([target_image.size[::-1]])
    inputs = processor(images=target_image, query_images=query_image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model.image_guided_detection(**inputs)
    
    img = cv2.cvtColor(np.array(target_image), cv2.COLOR_BGR2RGB)
    outputs.logits = outputs.logits.cpu()
    outputs.target_pred_boxes = outputs.target_pred_boxes.cpu()
    
    results = processor.post_process_image_guided_detection(outputs=outputs, threshold=threshold, nms_threshold=nms_threshold, target_sizes=target_sizes)
    boxes, scores = results[0]["boxes"], results[0]["scores"]

    if len(boxes) == 0:
        return []

    filtered_boxes = []
    for box in boxes:
        x1, y1, x2, y2 = [int(i) for i in box.tolist()]
        cropped_img = img[y1:y2, x1:x2]
        if cropped_img.size != 0:
            filtered_boxes.append(cropped_img)

    return filtered_boxes

def process_video(video_path, query_image, skipframes=0):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return

    frame_count = 0
    all_results = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % (skipframes + 1) == 0:
            frame_file = save_array_to_temp_image(frame)
            result_frames = detect_and_crop(Image.open(frame_file), query_image)
            for res in result_frames:
                saved_res = save_array_to_temp_image(res)
                embedding1 = extract_embedding(query_image)
                embedding2 = extract_embedding(Image.open(saved_res))
                dist = l2_distance(embedding1, embedding2).item()
                cos = cosine_similarity(embedding1, embedding2).item()
                all_results.append({'l2_dist': dist, 'cos': cos})
        frame_count += 1
    cap.release()
    return all_results

def process_videos_and_compare(image, video, skipframes=5, threshold=0.47):
    def median(values):
        n = len(values)
        return (values[n // 2 - 1] + values[n // 2]) / 2 if n % 2 == 0 else values[n // 2]

    results = process_video(video, image, skipframes)
    if results:
        l2_dists = [item['l2_dist'] for item in results]
        cosines = [item['cos'] for item in results]
        avg_l2_dist = sum(l2_dists) / len(l2_dists)
        avg_cos = sum(cosines) / len(cosines)
        median_l2_dist = median(sorted(l2_dists))
        median_cos = median(sorted(cosines))
        result = {
            "avg_l2_dist": avg_l2_dist,
            "avg_cos": avg_cos,
            "median_l2_dist": median_l2_dist,
            "median_cos": median_cos,
            "avg_cos_dist": 1 - avg_cos,
            "median_cos_dist": 1 - median_cos,
            "is_present": avg_cos >= threshold
        }
    else:
        result = {
            "avg_l2_dist": float('inf'),
            "avg_cos": 0,
            "median_l2_dist": float('inf'),
            "median_cos": 0,
            "avg_cos_dist": float('inf'),
            "median_cos_dist": float('inf'),
            "is_present": False
        }
    return result

def interface(video, image, skipframes, threshold):
    result = process_videos_and_compare(image, video, skipframes, threshold)
    return result

iface = gr.Interface(
    fn=interface,
    inputs=[
        gr.Video(label="Upload a Video"),
        gr.Image(type="pil", label="Upload a Query Image"),
        gr.Slider(minimum=0, maximum=10, step=1, default=5, label="Skip Frames"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.47, label="Threshold")
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="Object Detection in Video",
    description="""
    **Instructions:**

    1. **Upload a Video**: Select a video file to upload. 
    2. **Upload a Query Image**: Select an image file that contains the object you want to detect in the video.
    3. **Set Skip Frames**: Adjust the slider to set the number of frames to skip between each processing.
    4. **Set Threshold**: Adjust the slider to set the threshold for cosine similarity to determine if the object is present in the video.
    5. **View Results**: The result will show the average and median distances and similarities, and whether the object is present in the video based on the threshold.
    """
)

if __name__ == "__main__":
    iface.launch()