Update app.py

app.py

@@ -1,134 +1,134 @@
-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"
-
-with gr.Blocks() as demo:
-    gr.Markdown("## Zero-shot Object Tracking with PaliGemma")
-    gr.Markdown("This is a demo for zero-shot object tracking using [PaliGemma](https://huggingface.co/google/paligemma-3b-mix-448) vision language model by Google.")
-    gr.Markdown("Simply upload a video and enter the candidate labels, or try the example below. Text input should be ; separated. 👇")
-    with gr.Tab(label="Video"):
-        with gr.Row():
-            input_video = gr.Video(label='Input Video')
-            output_video = gr.Video(label='Output Video')
-        with gr.Row():
-            candidate_labels = gr.Textbox(
-                label='Labels',
-                placeholder='Labels separated by a comma',
-            )
-            submit = gr.Button()
-        gr.Examples(
-            fn=process_video,
-            examples=[["./cats.mp4", "bird ; cat"]],
-            inputs=[
-                input_video,
-                candidate_labels,
-                
-            ],
-            outputs=output_video
-        )
-
-    submit.click(fn=process_video,
-        inputs=[input_video, candidate_labels],
-        outputs=output_video
-        )
-
+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"
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Zero-shot Object Tracking with PaliGemma")
+    gr.Markdown("This is a demo for zero-shot object tracking using [PaliGemma](https://huggingface.co/google/paligemma-3b-mix-448) vision language model by Google.")
+    gr.Markdown("Simply upload a video and enter the candidate labels, or try the example below. Text input should be ; separated. 👇")
+    with gr.Tab(label="Video"):
+        with gr.Row():
+            input_video = gr.Video(label='Input Video')
+            output_video = gr.Video(label='Output Video')
+        with gr.Row():
+            candidate_labels = gr.Textbox(
+                label='Labels',
+                placeholder='Labels separated by a comma',
+            )
+            submit = gr.Button()
+        gr.Examples(
+            fn=process_video,
+            examples=[["./input.mp4", "detect person"]],
+            inputs=[
+                input_video,
+                candidate_labels,
+                
+            ],
+            outputs=output_video
+        )
+
+    submit.click(fn=process_video,
+        inputs=[input_video, candidate_labels],
+        outputs=output_video
+        )
+
 demo.launch(debug=False, show_error=True)