Ensemble Kosmos2

app.py

@@ -1,105 +1,384 @@
+import gradio as gr
+import random
+import numpy as np
+import os
+import requests
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+import cv2
+import ast
 import torch
 from efficientnet_pytorch import EfficientNet
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
 from super_gradients.training import models
-import cv2
-import numpy as np
 
-device = torch.device("cpu")
-
-# Load the YOLO-NAS model
-yolo_nas_l = models.get("yolo_nas_l", pretrained_weights="coco")
-
-def bounding_boxes_overlap(box1, box2):
-    """Check if two bounding boxes overlap or touch."""
-    x1, y1, x2, y2 = box1
-    x3, y3, x4, y4 = box2
-    return not (x3 > x2 or x4 < x1 or y3 > y2 or y4 < y1)
-
-def merge_boxes(box1, box2):
-    """Return the encompassing bounding box of two boxes."""
-    x1, y1, x2, y2 = box1
-    x3, y3, x4, y4 = box2
-    x = min(x1, x3)
-    y = min(y1, y3)
-    w = max(x2, x4)
-    h = max(y2, y4)
-    return (x, y, w, h)
-
-def save_merged_boxes(predictions, image_np):
-    """Save merged bounding boxes as separate images."""
-    processed_boxes = set()
-    roi = None  # Initialize roi to None
-
-    for image_prediction in predictions:
-        bboxes = image_prediction.prediction.bboxes_xyxy  
-        for box1 in bboxes:
-            for box2 in bboxes:
-                if np.array_equal(box1, box2):
-                    continue
-                if bounding_boxes_overlap(box1, box2) and tuple(box1) not in processed_boxes and tuple(box2) not in processed_boxes:
-                    merged_box = merge_boxes(box1, box2)
-                    roi = image_np[int(merged_box[1]):int(merged_box[3]), int(merged_box[0]):int(merged_box[2])]
-                    processed_boxes.add(tuple(box1))
-                    processed_boxes.add(tuple(box2))
-                    break  # Exit the inner loop once a match is found
-            if roi is not None:
-                break  # Exit the outer loop once a match is found
-
-    return roi
-
-# Load the EfficientNet model
-def load_model(model_path):
-    model = torch.load(model_path)
-    model = model.to(device)
-    model.eval()  # Set the model to evaluation mode
-    return model
-
-# Perform inference on an image
-def predict_image(image, model):
-    # First, get the ROI using YOLO-NAS
-    image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    predictions = yolo_nas_l.predict(image_np, iou=0.3, conf=0.35)
-    roi_new = save_merged_boxes(predictions, image_np)
+
+class Kosmos2:
+    def __init__(self):
+        self.colors = [
+            (0, 255, 0),
+            (0, 0, 255),
+            (255, 255, 0),
+            (255, 0, 255),
+            (0, 255, 255),
+            (114, 128, 250),
+            (0, 165, 255),
+            (0, 128, 0),
+            (144, 238, 144),
+            (238, 238, 175),
+            (255, 191, 0),
+            (0, 128, 0),
+            (226, 43, 138),
+            (255, 0, 255),
+            (0, 215, 255),
+            (255, 0, 0),    
+        ]
+
+        self.color_map = {
+            f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for color_id, color in enumerate(self.colors)
+        }
+
+        self.ckpt = "ydshieh/kosmos-2-patch14-224"
+        self.model = AutoModelForVision2Seq.from_pretrained(self.ckpt, trust_remote_code=True).to("cuda")
+        self.processor = AutoProcessor.from_pretrained(self.ckpt, trust_remote_code=True)
+
+    def is_overlapping(self, rect1, rect2):
+        x1, y1, x2, y2 = rect1
+        x3, y3, x4, y4 = rect2
+        return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)
+
+    def draw_entity_boxes_on_image(self, image, entities, show=False, save_path=None, entity_index=-1):
+        """_summary_
+        Args:
+            image (_type_): image or image path
+            collect_entity_location (_type_): _description_
+        """
+        if isinstance(image, Image.Image):
+            image_h = image.height
+            image_w = image.width
+            image = np.array(image)[:, :, [2, 1, 0]]
+        elif isinstance(image, str):
+            if os.path.exists(image):
+                pil_img = Image.open(image).convert("RGB")
+                image = np.array(pil_img)[:, :, [2, 1, 0]]
+                image_h = pil_img.height
+                image_w = pil_img.width
+            else:
+                raise ValueError(f"invaild image path, {image}")
+        elif isinstance(image, torch.Tensor):
+            # pdb.set_trace()
+            image_tensor = image.cpu()
+            reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
+            reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
+            image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
+            pil_img = T.ToPILImage()(image_tensor)
+            image_h = pil_img.height
+            image_w = pil_img.width
+            image = np.array(pil_img)[:, :, [2, 1, 0]]
+        else:
+            raise ValueError(f"invaild image format, {type(image)} for {image}")
+        
+        if len(entities) == 0:
+            return image
+
+        indices = list(range(len(entities)))
+        if entity_index >= 0:
+            indices = [entity_index]
+
+        # Not to show too many bboxes
+        entities = entities[:len(self.color_map)]
+        
+        new_image = image.copy()
+        previous_bboxes = []
+        # size of text
+        text_size = 1
+        # thickness of text
+        text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))
+        box_line = 3
+        (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+        base_height = int(text_height * 0.675)
+        text_offset_original = text_height - base_height
+        text_spaces = 3
+
+        # num_bboxes = sum(len(x[-1]) for x in entities)
+        used_colors = self.colors  # random.sample(colors, k=num_bboxes)
+
+        color_id = -1
+        for entity_idx, (entity_name, (start, end), bboxes) in enumerate(entities):
+            color_id += 1
+            if entity_idx not in indices:
+                continue
+            for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):
+                # if start is None and bbox_id > 0:
+                #     color_id += 1
+                orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)
+
+                # draw bbox
+                # random color
+                color = used_colors[color_id]  # tuple(np.random.randint(0, 255, size=3).tolist())
+                new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)
+
+                l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1
+
+                x1 = orig_x1 - l_o
+                y1 = orig_y1 - l_o
+
+                if y1 < text_height + text_offset_original + 2 * text_spaces:
+                    y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
+                    x1 = orig_x1 + r_o
+
+                # add text background
+                (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+                text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1
+
+                for prev_bbox in previous_bboxes:
+                    while self.is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):
+                        text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
+                        text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
+                        y1 += (text_height + text_offset_original + 2 * text_spaces)
+
+                        if text_bg_y2 >= image_h:
+                            text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
+                            text_bg_y2 = image_h
+                            y1 = image_h
+                            break
+
+                alpha = 0.5
+                for i in range(text_bg_y1, text_bg_y2):
+                    for j in range(text_bg_x1, text_bg_x2):
+                        if i < image_h and j < image_w:
+                            if j < text_bg_x1 + 1.35 * c_width:
+                                # original color
+                                bg_color = color
+                            else:
+                                # white
+                                bg_color = [255, 255, 255]
+                            new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)
+
+                cv2.putText(
+                    new_image, f"  {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces), cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA
+                )
+                # previous_locations.append((x1, y1))
+                previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))
+
+        pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])
+        if save_path:
+            pil_image.save(save_path)
+        if show:
+            pil_image.show()
+
+        return pil_image
+
+    def generate_predictions(self, image_input, text_input):
+
+        # Save the image and load it again to match the original Kosmos-2 demo.
+        # (https://github.com/microsoft/unilm/blob/f4695ed0244a275201fff00bee495f76670fbe70/kosmos-2/demo/gradio_app.py#L345-L346)
+        user_image_path = "/tmp/user_input_test_image.jpg"
+        image_input.save(user_image_path)
+        # This might give different results from the original argument `image_input`
+        image_input = Image.open(user_image_path)
+
+        if text_input == "Brief":
+            text_input = "<grounding>An image of"
+        elif text_input == "Detailed":
+            text_input = "<grounding>Describe this image in detail:"
+        else:
+            text_input = f"<grounding>{text_input}"
+
+        inputs = self.processor(text=text_input, images=image_input, return_tensors="pt")
+
+        generated_ids = self.model.generate(
+            pixel_values=inputs["pixel_values"].to("cuda"),
+            input_ids=inputs["input_ids"][:, :-1].to("cuda"),
+            attention_mask=inputs["attention_mask"][:, :-1].to("cuda"),
+            img_features=None,
+            img_attn_mask=inputs["img_attn_mask"][:, :-1].to("cuda"),
+            use_cache=True,
+            max_new_tokens=128,
+        )
+        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        # By default, the generated  text is cleanup and the entities are extracted.
+        processed_text, entities = self.processor.post_process_generation(generated_text)
+
+        annotated_image = self.draw_entity_boxes_on_image(image_input, entities, show=False)
+
+        color_id = -1
+        entity_info = []
+        filtered_entities = []
+        for entity in entities:
+            entity_name, (start, end), bboxes = entity
+            if start == end:
+                # skip bounding bbox without a `phrase` associated
+                continue
+            color_id += 1
+            # for bbox_id, _ in enumerate(bboxes):
+                # if start is None and bbox_id > 0:
+                #     color_id += 1
+            entity_info.append(((start, end), color_id))
+            filtered_entities.append(entity)
+
+        colored_text = []
+        prev_start = 0
+        end = 0
+        for idx, ((start, end), color_id) in enumerate(entity_info):
+            if start > prev_start:
+                colored_text.append((processed_text[prev_start:start], None))
+            colored_text.append((processed_text[start:end], f"{color_id}"))
+            prev_start = end
+
+        if end < len(processed_text):
+            colored_text.append((processed_text[end:len(processed_text)], None))
+
+        return annotated_image, colored_text, str(filtered_entities)
+
+class VehiclePredictor:
+    def __init__(self, model_path):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.yolo_nas_l = models.get("yolo_nas_l", pretrained_weights="coco")
+        self.classifier_model = torch.load(model_path)
+        self.classifier_model = self.classifier_model.to(self.device)
+        self.classifier_model.eval()  # Set the model to evaluation mode
+
+    def bounding_boxes_overlap(self, box1, box2):
+        """Check if two bounding boxes overlap or touch."""
+        x1, y1, x2, y2 = box1
+        x3, y3, x4, y4 = box2
+        return not (x3 > x2 or x4 < x1 or y3 > y2 or y4 < y1)
     
-    if roi_new is None:
-        roi_new = image_np  # Use the original image if no ROI is found
-
-    # Convert ROI back to PIL Image for EfficientNet
-    roi_image = Image.fromarray(cv2.cvtColor(roi_new, cv2.COLOR_BGR2RGB))
-
-    # Define the image transformations
-    transform = transforms.Compose([
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-
-    # Convert PIL Image to Tensor
-    roi_image_tensor = transform(roi_image).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        outputs = model(roi_image_tensor)
-        _, predicted = outputs.max(1)
-        prediction_text = 'Accident' if predicted.item() == 0 else 'No accident'
+    def merge_boxes(self, box1, box2):
+        """Return the encompassing bounding box of two boxes."""
+        x1, y1, x2, y2 = box1
+        x3, y3, x4, y4 = box2
+        x = min(x1, x3)
+        y = min(y1, y3)
+        w = max(x2, x4)
+        h = max(y2, y4)
+        return (x, y, w, h)
     
-    return roi_image, prediction_text  # Return both the roi_image and the prediction text
-
-# Load the EfficientNet model outside the function to avoid loading it multiple times
-model_path = 'vehicle.pt'
-model = load_model(model_path)
-
-# Gradio UI
-title = "Vehicle Collision Classification"
-description = "Upload an image to determine if it depicts a vehicle accident. Powered by EfficientNet."
-examples = [["roi_none.png"], ["test2.jpeg"]]  
-
-gr.Interface(fn=lambda img: predict_image(img, model), 
-             inputs=gr.inputs.Image(type="pil"), 
-             outputs=[gr.outputs.Image(type="pil"), "text"],  # Updated outputs to handle both image and text
-             title=title, 
-             description=description, 
-             examples=examples).launch()
+    def save_merged_boxes(self, predictions, image_np):
+        """Save merged bounding boxes as separate images."""
+        processed_boxes = set()
+        roi = None  # Initialize roi to None
+
+        for image_prediction in predictions:
+            bboxes = image_prediction.prediction.bboxes_xyxy  
+            for box1 in bboxes:
+                for box2 in bboxes:
+                    if np.array_equal(box1, box2):
+                        continue
+                    if self.bounding_boxes_overlap(box1, box2) and tuple(box1) not in processed_boxes and tuple(box2) not in processed_boxes:
+                        merged_box = self.merge_boxes(box1, box2)
+                        roi = image_np[int(merged_box[1]):int(merged_box[3]), int(merged_box[0]):int(merged_box[2])]
+                        processed_boxes.add(tuple(box1))
+                        processed_boxes.add(tuple(box2))
+                        break  # Exit the inner loop once a match is found
+                if roi is not None:
+                    break  # Exit the outer loop once a match is found
+
+        return roi
+
+    # Perform inference on an image
+    def predict_image(self, image, model):
+        # First, get the ROI using YOLO-NAS
+        image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        predictions = self.yolo_nas_l.predict(image_np, iou=0.3, conf=0.35)
+        roi_new = self.save_merged_boxes(predictions, image_np)
+        
+        if roi_new is None:
+            roi_new = image_np  # Use the original image if no ROI is found
+
+        # Convert ROI back to PIL Image for EfficientNet
+        roi_image = Image.fromarray(cv2.cvtColor(roi_new, cv2.COLOR_BGR2RGB))
+
+        # Define the image transformations
+        transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        # Convert PIL Image to Tensor
+        roi_image_tensor = transform(roi_image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.classifier_model(roi_image_tensor)
+            _, predicted = outputs.max(1)
+            prediction_text = 'Accident' if predicted.item() == 0 else 'No accident'
+        
+        return roi_image, prediction_text  # Return both the roi_image and the prediction text
+
+
+def main():
+    kosmos2 = Kosmos2()
+    vehicle_predictor = VehiclePredictor('vehicle.pt')
+
+    with gr.Blocks(title="Advanced Vehicle Contextualization & Collision Prediction", theme=gr.themes.Base()).queue() as demo:
+        gr.Markdown(("""
+            # Models used - 
+            Kosmos-2: Grounding Multimodal Large Language Models to the World
+            [[Paper]](https://arxiv.org/abs/2306.14824) [[Code]](https://github.com/microsoft/unilm/blob/master/kosmos-2)
+            YOLO-NAS [[Code]](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md)
+            EfficientNet-b0
+            """))
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Test Image")
+                text_input = gr.Radio(["Brief", "Detailed"], label="Description Type", value="Brief")
+                run_button = gr.Button(label="Run", visible=True)
+
+            with gr.Column():
+                image_output_kosmos = gr.Image(type="pil", label="Kosmos-2 Output Image")
+                text_output_kosmos = gr.HighlightedText(
+                                    label="Generated Description by Kosmos-2",
+                                    combine_adjacent=False,
+                                    show_legend=True,
+                                ).style(color_map=kosmos2.color_map)
+                
+                image_output_vehicle = gr.Image(type="pil", label="Collision Predictor Output Image", size=(112, 112))
+                text_output_vehicle = gr.Textbox(label="Collision Predictor Result")
+
+        # record which text span (label) is selected
+        selected = gr.Number(-1, show_label=False, placeholder="Selected", visible=False)
+
+        # record the current `entities`
+        entity_output = gr.Textbox(visible=False)
+
+        # get the current selected span label
+        def get_text_span_label(evt: gr.SelectData):
+            if evt.value[-1] is None:
+                return -1
+            return int(evt.value[-1])
+        # and set this information to `selected`
+        text_output_kosmos.select(get_text_span_label, None, selected)
+        
+        # update output image when we change the span (enity) selection
+        def update_output_image(img_input, image_output, entities, idx):
+            entities = ast.literal_eval(entities)
+            updated_image = kosmos2.draw_entity_boxes_on_image(img_input, entities, entity_index=idx)
+            return updated_image
+        selected.change(update_output_image, [image_input, image_output_kosmos, entity_output, selected], [image_output_kosmos])
+
+        def combined_predictions(img, description_type):
+            # Kosmos2 predictions
+            kosmos_image, kosmos_text, entities = kosmos2.generate_predictions(img, description_type)
+            
+            # VehiclePredictor predictions
+            vehicle_image, vehicle_text = vehicle_predictor.predict_image(img, vehicle_predictor.classifier_model)
+            
+            return kosmos_image, kosmos_text, entities, vehicle_image, vehicle_text
+
+        run_button.click(fn=combined_predictions,
+                         inputs=[image_input, text_input],
+                         outputs=[image_output_kosmos, text_output_kosmos, entity_output, image_output_vehicle, text_output_vehicle],
+                         show_progress=True, queue=True)
+
+    demo.launch(share=True)
+
+if __name__ == "__main__":
+    main()