Update app.py

app.py

@@ -2,9 +2,7 @@ import subprocess
 import sys
 from os.path import abspath, dirname,join
 import spaces
-
-sys.path.append(join(dirname(abspath(__file__)),'GroundingDINO'))
-
+@spaces.GPU()
 def run_commands():
     commands = [
         "git clone https://github.com/IDEA-Research/GroundingDINO.git",
@@ -22,211 +20,216 @@ def run_commands():
             print(result.stdout.decode())
         except subprocess.CalledProcessError as e:
             print(f"Command '{command}' failed with error: {e.stderr.decode()}")
+            
+run_commands()
 
-@spaces.GPU(30)
-def video_app_setup_and_run_pipeline():
-
-        #Setup needs gpu available to install properly
-        run_commands()
-        from typing import List
-        from Utils import get_video_properties
-        from GroundingDINO.groundingdino.util.inference import load_model, predict
-        import cv2
-        import numpy as np
-        import torch
-        from PIL import Image
-        import GroundingDINO.groundingdino.datasets.transforms as T
-        from torchvision.ops import box_convert
-        from torchvision import transforms
-        from torch import nn
-        from os.path import dirname, abspath
-        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
+from typing import List
+from Utils import get_video_properties
+from GroundingDINO.groundingdino.util.inference import load_model, predict
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+import GroundingDINO.groundingdino.datasets.transforms as T
+from torchvision.ops import box_convert
+from torchvision import transforms
+from torch import nn
+from os.path import dirname, abspath
+import yaml
+import supervision as sv
+import gradio as gr
+sys.path.append(join(dirname(abspath(__file__)),'GroundingDINO'))
 
-        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()
-        
-        class VideoObjectDetection:
-        
-            def __init__(self,
-                         text_prompt: str
-                         ):
-        
-                self.text_prompt = text_prompt
-        
-            def crop(self, frame, boxes):
+class DinoVisionTransformerClassifier(nn.Module):
+    def __init__(self):
         
-                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 = []
+        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):
         
-                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)
+        x = self.transformer(x)
+        x = self.transformer.norm(x)
+        x = self.classifier(x)
+        return x
+
+class ImageClassifier:
+
+    def __init__(self):
         
-                labels = [
-                    f"{phrase} {logit:.2f}"
-                    for phrase, logit
-                    in zip(custom_labels, custom_logits)
+        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]),
                 ]
-        
-                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')
+            )
+        }
+        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()
+
+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
+
+    @spaces.GPU()
+    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
 
-        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)
+            # 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)
 
-        iface.launch(share=False, debug=True)
+            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
+    
             
 if __name__ == "__main__":
-    video_app_setup_and_run_pipeline()
+    
+    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)