Deleted 'gpt4v.py' and moved its functionalities to 'utils.py' and 'app.py'.

Dockerfile

@@ -32,7 +32,7 @@ RUN pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 -f https://download
 
 # Install dependencies
 RUN pip install --no-cache-dir gradio==3.50.2 opencv-python supervision==0.17.0rc4 \
-    pillow requests
+    pillow requests setofmark==0.1.0rc3
 
 # Install SAM and Detectron2
 RUN pip install 'git+https://github.com/facebookresearch/segment-anything.git'
@@ -44,7 +44,6 @@ RUN wget -c -O $HOME/app/weights/sam_vit_h_4b8939.pth https://dl.fbaipublicfiles
 
 COPY app.py .
 COPY utils.py .
-COPY gpt4v.py .
 COPY sam_utils.py .
 
 RUN find $HOME/app

app.py

@@ -4,13 +4,13 @@ from typing import List, Dict, Tuple, Any, Optional
 import cv2
 import gradio as gr
 import numpy as np
+import som
 import supervision as sv
 import torch
 from segment_anything import sam_model_registry
 
-from gpt4v import prompt_image
 from sam_utils import sam_interactive_inference, sam_inference
-from utils import postprocess_masks, Visualizer, extract_numbers_in_brackets
+from utils import postprocess_masks, Visualizer
 
 HOME = os.getenv("HOME")
 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
@@ -21,17 +21,21 @@ SAM_MODEL_TYPE = "vit_h"
 
 ANNOTATED_IMAGE_KEY = "annotated_image"
 DETECTIONS_KEY = "detections"
-
 MARKDOWN = """
-[![arXiv](https://img.shields.io/badge/arXiv-1703.06870v3-b31b1b.svg)](https://arxiv.org/pdf/2310.11441.pdf)
-
-<h1 style='text-align: center'>
-    <img 
-        src='https://som-gpt4v.github.io/website/img/som_logo.png' 
-        style='height:50px; display:inline-block'
-    />  
-    Set-of-Mark (SoM) Prompting Unleashes Extraordinary Visual Grounding in GPT-4V
-</h1>
+<div align='center'>
+    <h1>
+        <img 
+            src='https://som-gpt4v.github.io/website/img/som_logo.png' 
+            style='height:50px; display:inline-block'
+        />  
+        Set-of-Mark (SoM) Prompting Unleashes Extraordinary Visual Grounding in GPT-4V
+    </h1>
+    <br>
+    [<a href="https://arxiv.org/abs/2109.07529"> arXiv paper </a>] 
+    [<a href="https://som-gpt4v.github.io"> project page </a>]
+    [<a href="https://github.com/roboflow/set-of-mark"> python package </a>]
+    [<a href="https://github.com/microsoft/SoM"> code </a>]
+</div>
 
 ## 🚧 Roadmap
 
@@ -90,7 +94,7 @@ def prompt(
         return "⚠️ Please set your OpenAI API key first"
     if state is None or ANNOTATED_IMAGE_KEY not in state:
         return "⚠️ Please generate SoM visual prompt first"
-    return prompt_image(
+    return som.prompt_image(
         api_key=api_key,
         image=cv2.cvtColor(state[ANNOTATED_IMAGE_KEY], cv2.COLOR_BGR2RGB),
         prompt=message
@@ -114,15 +118,17 @@ def highlight(
     if len(history) == 0:
         return None
 
-    response = history[-1][-1]
-    detections_ids = extract_numbers_in_brackets(text=response)
-    highlighted_detections = [
-        (detections.mask[detection_id], str(detection_id))
-        for detection_id
-        in detections_ids
+    text = history[-1][-1]
+    relevant_masks = som.extract_relevant_masks(
+        text=text,
+        detections=detections
+    )
+    relevant_masks = [
+        (mask, mark)
+        for mark, mask
+        in relevant_masks.items()
     ]
-
-    return annotated_image, highlighted_detections
+    return annotated_image, relevant_masks
 
 
 image_input = gr.Image(
@@ -131,7 +137,8 @@ image_input = gr.Image(
     tool="sketch",
     interactive=True,
     brush_radius=20.0,
-    brush_color="#FFFFFF"
+    brush_color="#FFFFFF",
+    height=512
 )
 checkbox_annotation_mode = gr.CheckboxGroup(
     choices=["Mark", "Polygon", "Mask", "Box"],
@@ -147,7 +154,8 @@ image_output = gr.AnnotatedImage(
     color_map={
         str(i): sv.ColorPalette.default().by_idx(i).as_hex()
         for i in range(64)
-    }
+    },
+    height=512
 )
 openai_api_key = gr.Textbox(
     show_label=False,

gpt4v.py

@@ -1,81 +0,0 @@
-import cv2
-import base64
-import requests
-
-import numpy as np
-
-
-META_PROMPT = '''
-For any labels or markings on an image that you reference in your response, please 
-enclose them in square brackets ([]) and list them explicitly. Do not use ranges; for 
-example, instead of '1 - 4', list as '[1], [2], [3], [4]'. These labels could be 
-numbers or letters and typically correspond to specific segments or parts of the image.
-'''
-API_URL = "https://api.openai.com/v1/chat/completions"
-
-
-def encode_image_to_base64(image: np.ndarray) -> str:
-    """
-    Encodes an image into a base64-encoded string in JPEG format.
-
-    Parameters:
-        image (np.ndarray): The image to be encoded. This should be a numpy array as
-            typically used in OpenCV.
-
-    Returns:
-        str: A base64-encoded string representing the image in JPEG format.
-    """
-    success, buffer = cv2.imencode('.jpg', image)
-    if not success:
-        raise ValueError("Could not encode image to JPEG format.")
-
-    encoded_image = base64.b64encode(buffer).decode('utf-8')
-    return encoded_image
-
-
-def compose_headers(api_key: str) -> dict:
-    return {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-
-
-def compose_payload(image: np.ndarray, prompt: str) -> dict:
-    base64_image = encode_image_to_base64(image)
-    return {
-        "model": "gpt-4-vision-preview",
-        "messages": [
-            {
-                "role": "system",
-                "content": [
-                    META_PROMPT
-                ]
-            },
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt
-                    },
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/jpeg;base64,{base64_image}"
-                        }
-                    }
-                ]
-            }
-        ],
-        "max_tokens": 800
-    }
-
-
-def prompt_image(api_key: str, image: np.ndarray, prompt: str) -> str:
-    headers = compose_headers(api_key=api_key)
-    payload = compose_payload(image=image, prompt=prompt)
-    response = requests.post(url=API_URL, headers=headers, json=payload).json()
-
-    if 'error' in response:
-        raise ValueError(response['error']['message'])
-    return response['choices'][0]['message']['content']

utils.py

@@ -1,7 +1,5 @@
-import re
-from typing import List
-
 import cv2
+import som
 
 import numpy as np
 import supervision as sv
@@ -13,7 +11,7 @@ class Visualizer:
         self,
         line_thickness: int = 2,
         mask_opacity: float = 0.1,
-        text_scale: float = 0.5
+        text_scale: float = 0.6
     ) -> None:
         self.box_annotator = sv.BoundingBoxAnnotator(
             color_lookup=sv.ColorLookup.INDEX,
@@ -25,6 +23,8 @@ class Visualizer:
             color_lookup=sv.ColorLookup.INDEX,
             thickness=line_thickness)
         self.label_annotator = sv.LabelAnnotator(
+            color=sv.Color.black(),
+            text_color=sv.Color.white(),
             color_lookup=sv.ColorLookup.INDEX,
             text_position=sv.Position.CENTER_OF_MASS,
             text_scale=text_scale)
@@ -85,7 +85,11 @@ def refine_mask(
         relative_area = area / total_area
         if relative_area < area_threshold:
             cv2.drawContours(
-                mask, [contour], -1, (0 if mode == 'islands' else 255), -1
+                image=mask,
+                contours=[contour],
+                contourIdx=-1,
+                color=(0 if mode == 'islands' else 255),
+                thickness=-1
             )
 
     return np.where(mask > 0, 1, 0).astype(bool)
@@ -116,52 +120,6 @@ def filter_masks_by_relative_area(
     return masks[min_area_filter & max_area_filter]
 
 
-def compute_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
-    """
-    Computes the Intersection over Union (IoU) of two masks.
-
-    Parameters:
-        mask1, mask2 (np.ndarray): Two mask arrays.
-
-    Returns:
-        float: The IoU of the two masks.
-    """
-    intersection = np.logical_and(mask1, mask2).sum()
-    union = np.logical_or(mask1, mask2).sum()
-    return intersection / union if union != 0 else 0
-
-
-def filter_highly_overlapping_masks(
-    masks: np.ndarray,
-    iou_threshold: float
-) -> np.ndarray:
-    """
-    Removes masks with high overlap from a set of masks.
-
-    Parameters:
-        masks (np.ndarray): A 3D numpy array with shape (N, H, W), where N is the
-            number of masks, and H and W are the height and width of the masks.
-        iou_threshold (float): The IoU threshold above which masks will be considered as
-            overlapping.
-
-    Returns:
-        np.ndarray: A 3D numpy array of masks with highly overlapping masks removed.
-    """
-    num_masks = masks.shape[0]
-    keep_mask = np.ones(num_masks, dtype=bool)
-
-    for i in range(num_masks):
-        for j in range(i + 1, num_masks):
-            if not keep_mask[i] or not keep_mask[j]:
-                continue
-
-            iou = compute_iou(masks[i, :, :], masks[j, :, :])
-            if iou > iou_threshold:
-                keep_mask[j] = False
-
-    return masks[keep_mask]
-
-
 def postprocess_masks(
     detections: sv.Detections,
     area_threshold: float = 0.01,
@@ -200,7 +158,7 @@ def postprocess_masks(
         masks=masks,
         min_relative_area=min_relative_area,
         max_relative_area=max_relative_area)
-    masks = filter_highly_overlapping_masks(
+    masks = som.mask_non_max_suppression(
         masks=masks,
         iou_threshold=iou_threshold)
 
@@ -208,18 +166,3 @@ def postprocess_masks(
         xyxy=sv.mask_to_xyxy(masks),
         mask=masks
     )
-
-
-def extract_numbers_in_brackets(text: str) -> List[int]:
-    """
-    Extracts all numbers enclosed in square brackets from a given string.
-
-    Args:
-        text (str): The string to be searched.
-
-    Returns:
-        List[int]: A list of integers found within square brackets.
-    """
-    pattern = r'\[(\d+)\]'
-    numbers = [int(num) for num in re.findall(pattern, text)]
-    return numbers