import subprocess import sys from os.path import abspath, dirname,join import spaces def run_commands(): commands = [ "git clone https://github.com/IDEA-Research/GroundingDINO.git", "pip install -e ./GroundingDINO", "cd GroundingDINO", "mkdir GroundingDINO/weights", "cd GroundingDINO/weights && wget -q https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth", "ls" ] for command in commands: try: print(f"Running command: {command}") result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) print(result.stdout.decode()) except subprocess.CalledProcessError as e: print(f"Command '{command}' failed with error: {e.stderr.decode()}") from typing import List from Utils import get_video_properties import cv2 import numpy as np import torch from PIL import Image from torchvision.ops import box_convert from torchvision import transforms from torch import nn import yaml import supervision as sv import gradio as gr class DinoVisionTransformerClassifier(nn.Module): def __init__(self): super(DinoVisionTransformerClassifier, self).__init__() self.transformer = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14") self.classifier = nn.Sequential(nn.Linear(384, 256), nn.ReLU(), nn.Linear(256, 2)) def forward(self, x): x = self.transformer(x) x = self.transformer.norm(x) x = self.classifier(x) return x class ImageClassifier: def __init__(self): with open(f"{dirname(abspath(__file__))}/config.yaml", 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) labels = config["labels"] self.labels = labels self.dino = DinoVisionTransformerClassifier() model_path = f"{dirname(abspath(__file__))}/model.pth" state_dict = torch.load(model_path) self.dino.load_state_dict(state_dict) def preprocess(self, image: np.ndarray) -> torch.Tensor: data_transforms = { "test": transforms.Compose( [ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]), ] ) } image_pillow = Image.fromarray(image) img_transformed = data_transforms['test'](image_pillow) return img_transformed def predict(self, image): image = self.preprocess(image) image = image.unsqueeze(0) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.dino.to(device) self.dino.eval() with torch.no_grad(): output = self.dino(image.to(device)) logit, predicted = torch.max(output.data, 1) return self.labels[predicted[0].item()], logit[0].item() def pipeline_to_setup_with_gpu(): run_commands() sys.path.append(join(dirname(abspath(__file__)),'GroundingDINO')) from GroundingDINO.groundingdino.util.inference import load_model, predict import GroundingDINO.groundingdino.datasets.transforms as T class VideoObjectDetection: def __init__(self, text_prompt: str ): self.text_prompt = text_prompt def crop(self, frame, boxes): h, w, _ = frame.shape boxes = boxes * torch.Tensor([w, h, w, h]) xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() min_col, min_row, max_col, max_row = map(int, xyxy[0]) crop_image = frame[min_row:max_row, min_col:max_col, :] return crop_image def annotate(self, image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], frame_rgb: np.ndarray, classifier) -> np.ndarray: h, w, _ = image_source.shape boxes = boxes * torch.Tensor([w, h, w, h]) xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() detections = sv.Detections(xyxy=xyxy) print(xyxy.shape) custom_labels = [] custom_logits = [] for box in xyxy: min_col, min_row, max_col, max_row = map(int, box) crop_image = frame_rgb[min_row:max_row, min_col:max_col, :] label, logit = classifier.predict(crop_image) print() if logit >= 1: custom_labels.append(label) custom_logits.append(logit) else: custom_labels.append('unknown human face') custom_logits.append(logit) labels = [ f"{phrase} {logit:.2f}" for phrase, logit in zip(custom_labels, custom_logits) ] box_annotator = sv.BoxAnnotator() annotated_frame = box_annotator.annotate(scene=image_source, detections=detections, labels=labels) return annotated_frame def preprocess_image(self, image: np.ndarray) -> torch.Tensor: transform = T.Compose( [ T.RandomResize([800], max_size=1333), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) image_pillow = Image.fromarray(image) image_transformed, _ = transform(image_pillow, None) return image_transformed def generate_video(self, video_path) -> None: # Load model, set up variables and get video properties cap, fps, width, height, fourcc = get_video_properties(video_path) model = load_model("GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py", "GroundingDINO/weights/groundingdino_swint_ogc.pth") predictor = ImageClassifier() TEXT_PROMPT = self.text_prompt BOX_TRESHOLD = 0.6 TEXT_TRESHOLD = 0.6 # Read video frames, crop image based on text prompt object detection and generate dataset_train import time frame_count = 0 delay = 1 / fps # Delay in seconds between frames while cap.isOpened(): start_time = time.time() ret, frame = cap.read() if not ret: break if cv2.waitKey(1) & 0xff == ord('q'): break # Convert bgr frame to rgb frame to image to torch tensor transformed frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image_transformed = self.preprocess_image(frame_rgb) boxes, logits, phrases = predict( model=model, image=image_transformed, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD ) # Get boxes if boxes.size()[0] > 0: annotated_frame = self.annotate(image_source=frame, boxes=boxes, logits=logits, phrases=phrases, frame_rgb=frame_rgb, classifier=predictor) # cv2.imshow('Object detection', annotated_frame) frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB) yield frame_rgb elapsed_time = time.time() - start_time time_to_wait = max(delay - elapsed_time, 0) time.sleep(time_to_wait) frame_count += 1 video_annotator = VideoObjectDetection( text_prompt='human face') with gr.Blocks() as iface: video_input = gr.Video(label="Upload Video") run_button = gr.Button("Start Processing") output_image = gr.Image(label="Classified video") run_button.click(fn=video_annotator.generate_video, inputs=video_input, outputs=output_image) iface.launch(share=False, debug=True) if __name__ == "__main__": pipeline_to_setup_with_gpu()