Zero-shot object detection on videos using PaliGemma.

app.py

@@ -0,0 +1,112 @@
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+import numpy as np
+from transformers import AutoTokenizer, PaliGemmaForConditionalGeneration, PaliGemmaProcessor
+import torch
+import gradio as gr
+
+# Load PaliGemma
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_id = "google/paligemma-3b-mix-224"
+model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
+processor = PaliGemmaProcessor.from_pretrained(model_id)
+
+# Function to draw bounding boxes (your original code)
+def draw_bounding_box(draw, coordinates, label, width, height):
+    y1, x1, y2, x2 = coordinates
+    y1, x1, y2, x2 = map(round, (y1*height, x1*width, y2*height, x2*width))
+
+    text_width, text_height = draw.textsize(label)
+    draw.rectangle([(x1, y1 - text_height - 2), (x1 + text_width + 4, y1)], fill="red")
+
+    # Draw label text
+    draw.text((x1 + 2, y1 - text_height - 2), label, fill="white")
+
+    # Draw bounding box
+    draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=2)
+
+def process_video(video_path, input_text):
+    cap = cv2.VideoCapture(video_path)
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    out = cv2.VideoWriter('output_paligemma_keras.avi', fourcc, 20.0, (int(cap.get(3)), int(cap.get(4))))
+
+    while(True):
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Convert the frame to a PIL Image
+        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+        # Send text prompt and image as input.
+        inputs = processor(text=input_text, images=img,
+                            padding="longest", do_convert_rgb=True, return_tensors="pt").to("cuda")
+        inputs = inputs.to(dtype=model.dtype)
+
+        # Get output.
+        with torch.no_grad():
+            output = model.generate(**inputs, max_length=496)
+
+        paligemma_response = processor.decode(output[0], skip_special_tokens=True)[len(input_text):].lstrip("\n")
+        # print(paligemma_response)  # For debugging
+
+        detections = paligemma_response.split(" ; ")
+
+        # Parse the output bounding box coordinates
+        parsed_coordinates = []
+        labels = []
+
+        for item in detections:
+            # Remove '<loc>' tags and split the string
+            # print(item)
+            detection = item.replace("<loc", "").split()
+
+            if len(detection) >= 2:
+              coordinates_str = detection[0]
+              label = detection[1]
+              labels.append(label)
+            else:
+              # No label detected, skip the iteration.
+              continue
+
+            # Split the coordinates string by '>' to get individual coordinates
+            coordinates = coordinates_str.split(">")
+            coordinates = coordinates[:4]  # Slicing to ensure only 4 values
+
+            if coordinates[-1] == '':
+                coordinates = coordinates[:-1]
+                # print(coordinates)
+
+            coordinates = [int(coord)/1024 for coord in coordinates]
+            # location_values = [int(loc) for loc in re.findall(r'\d{4}', coordinates)]
+            # y1, x1, y2, x2 = [value / 1024 for value in location_values]
+            parsed_coordinates.append(coordinates)
+
+        width = img.size[0]
+        height = img.size[1]
+
+        # Draw bounding boxes on the frame using PIL
+        draw = ImageDraw.Draw(img)
+        for coordinates, label in zip(parsed_coordinates, labels):
+            draw_bounding_box(draw, coordinates, label, width=width, height=height)
+
+        # Convert the PIL Image back to OpenCV format
+        frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+        # Write the frame to the output video
+        out.write(frame)
+
+    cap.release()
+    out.release()
+
+    return "output_paligemma_keras.avi"
+
+demo = gr.Interface(
+    fn=process_video,
+    inputs=[gr.Video(label="Input Video"), gr.Textbox(label="detect <class-name>")],
+    outputs=[gr.Video(label="Output Video")],
+    title="PaliGemma Object Detection",
+    description="Upload a video and specify the object you want to detect."
+)
+
+demo.launch(share=True)