Spaces:
Running
on
Zero
Running
on
Zero
| from transformers import Owlv2Processor, Owlv2ForObjectDetection | |
| from typing import List | |
| import os | |
| import numpy as np | |
| import supervision as sv | |
| import uuid | |
| import torch | |
| from tqdm import tqdm | |
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import spaces | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble") | |
| model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device) | |
| BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator() | |
| MASK_ANNOTATOR = sv.MaskAnnotator() | |
| LABEL_ANNOTATOR = sv.LabelAnnotator() | |
| def calculate_end_frame_index(source_video_path): | |
| video_info = sv.VideoInfo.from_video_path(source_video_path) | |
| return min( | |
| video_info.total_frames, | |
| video_info.fps * 2 | |
| ) | |
| def annotate_image( | |
| input_image, | |
| detections, | |
| labels | |
| ) -> np.ndarray: | |
| output_image = MASK_ANNOTATOR.annotate(input_image, detections) | |
| output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections) | |
| output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels) | |
| return output_image | |
| def process_video( | |
| input_video, | |
| labels, | |
| progress=gr.Progress(track_tqdm=True) | |
| ): | |
| labels = labels.split(",") | |
| video_info = sv.VideoInfo.from_video_path(input_video) | |
| total = calculate_end_frame_index(input_video) | |
| frame_generator = sv.get_video_frames_generator( | |
| source_path=input_video, | |
| end=total | |
| ) | |
| result_file_name = f"{uuid.uuid4()}.mp4" | |
| result_file_path = os.path.join("./", result_file_name) | |
| with sv.VideoSink(result_file_path, video_info=video_info) as sink: | |
| for _ in tqdm(range(total), desc="Processing video.."): | |
| frame = next(frame_generator) | |
| # list of dict of {"box": box, "mask":mask, "score":score, "label":label} | |
| results = query(frame, labels) | |
| print("results", results) | |
| detections = sv.Detections.from_transformers(results[0]) | |
| final_labels = [] | |
| for id in results[0]["labels"]: | |
| final_labels.append(labels[id]) | |
| frame = annotate_image( | |
| input_image=frame, | |
| detections=detections, | |
| labels=final_labels, | |
| ) | |
| sink.write_frame(frame) | |
| return result_file_path | |
| def query(image, texts): | |
| inputs = processor(text=texts, images=image, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| target_sizes = torch.Tensor([image.shape[:-1]]) | |
| results = processor.post_process_object_detection(outputs=outputs, threshold=0.3, target_sizes=target_sizes) | |
| return results | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## Zero-shot Object Tracking with OWLv2 🦉") | |
| gr.Markdown("This is a demo for zero-shot object tracking using [OWLv2](https://huggingface.co/google/owlv2-base-patch16-ensemble) model by Google.") | |
| gr.Markdown("Simply upload a video and enter the candidate labels, or try the example below. 👇") | |
| 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", "dog,cat"]], | |
| 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) | |