Create app.py

app.py

@@ -0,0 +1,165 @@
+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()