|
import os |
|
|
|
if not os.path.isdir("weights"): |
|
os.mkdir("weights") |
|
|
|
os.system("python -m pip install --upgrade pip") |
|
os.system( |
|
"wget https://raw.githubusercontent.com/asharma381/cs291I/main/backend/original_images/000749.png" |
|
) |
|
os.system( |
|
"wget -q -O weights/sam_vit_h_4b8939.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" |
|
) |
|
os.system( |
|
"wget -q -O weights/ram_plus_swin_large_14m.pth https://huggingface.co/xinyu1205/recognize-anything-plus-model/resolve/main/ram_plus_swin_large_14m.pth" |
|
) |
|
os.system( |
|
"wget -q -O weights/groundingdino_swint_ogc.pth https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth" |
|
) |
|
os.system("pip install git+https://github.com/xinyu1205/recognize-anything.git") |
|
os.system("pip install git+https://github.com/IDEA-Research/GroundingDINO.git") |
|
os.system("pip install git+https://github.com/facebookresearch/segment-anything.git") |
|
os.system("pip install openai==0.27.4") |
|
os.system("pip install tenacity") |
|
|
|
|
|
from typing import List, Tuple |
|
|
|
import cv2 |
|
import gradio as gr |
|
import groundingdino.config.GroundingDINO_SwinT_OGC |
|
import numpy as np |
|
import openai |
|
import torch |
|
from groundingdino.util.inference import Model |
|
from PIL import Image, ImageDraw |
|
from ram import get_transform |
|
from ram import inference_ram as inference |
|
from ram.models import ram_plus |
|
from scipy.spatial.distance import cdist |
|
from segment_anything import SamPredictor, sam_model_registry |
|
from supervision import Detections |
|
from tenacity import retry, wait_fixed |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
ram_model = None |
|
ram_threshold_multiplier = 1 |
|
gdino_model = None |
|
sam_model = None |
|
sam_predictor = None |
|
|
|
print("CUDA Available:", torch.cuda.is_available()) |
|
|
|
|
|
def get_tags_ram( |
|
image: Image.Image, threshold_multiplier=0.8, weights_folder="weights" |
|
) -> List[str]: |
|
global ram_model, ram_threshold_multiplier |
|
if ram_model is None: |
|
print("Loading RAM++ Model...") |
|
ram_model = ram_plus( |
|
pretrained=f"{weights_folder}/ram_plus_swin_large_14m.pth", |
|
vit="swin_l", |
|
image_size=384, |
|
) |
|
ram_model.eval() |
|
ram_model = ram_model.to(device) |
|
|
|
ram_model.class_threshold *= threshold_multiplier / ram_threshold_multiplier |
|
ram_threshold_multiplier = threshold_multiplier |
|
transform = get_transform() |
|
|
|
image = transform(image).unsqueeze(0).to(device) |
|
res = inference(image, ram_model) |
|
return [s.strip() for s in res[0].split("|")] |
|
|
|
|
|
def get_gdino_result( |
|
image: Image.Image, |
|
classes: List[str], |
|
box_threshold: float = 0.25, |
|
weights_folder="weights", |
|
) -> Tuple[Detections, List[str]]: |
|
global gdino_model |
|
|
|
if gdino_model is None: |
|
print("Loading GroundingDINO Model...") |
|
config_path = groundingdino.config.GroundingDINO_SwinT_OGC.__file__ |
|
gdino_model = Model( |
|
model_config_path=config_path, |
|
model_checkpoint_path=f"{weights_folder}/groundingdino_swint_ogc.pth", |
|
device=device, |
|
) |
|
|
|
detections, phrases = gdino_model.predict_with_caption( |
|
image=np.array(image), |
|
caption=", ".join(classes), |
|
box_threshold=box_threshold, |
|
text_threshold=0.25, |
|
) |
|
|
|
return detections, phrases |
|
|
|
|
|
def get_sam_model(weights_folder="weights"): |
|
global sam_model |
|
if sam_model is None: |
|
sam_checkpoint = f"{weights_folder}/sam_vit_h_4b8939.pth" |
|
sam_model = sam_model_registry["vit_h"](checkpoint=sam_checkpoint) |
|
sam_model.to(device=device) |
|
return sam_model |
|
|
|
|
|
def filter_tags_gdino(image: Image.Image, tags: List[str]) -> List[str]: |
|
detections, phrases = get_gdino_result(image, tags) |
|
filtered_tags = [] |
|
for tag in tags: |
|
for ( |
|
phrase, |
|
area, |
|
) in zip(phrases, detections.area): |
|
if area < 0.9 * image.size[0] * image.size[1] and tag in phrase: |
|
filtered_tags.append(tag) |
|
break |
|
return filtered_tags |
|
|
|
|
|
def read_file_to_string(file_path: str) -> str: |
|
content = "" |
|
|
|
try: |
|
with open(file_path, "r", encoding="utf8") as file: |
|
content = file.read() |
|
except FileNotFoundError: |
|
print(f"The file {file_path} was not found.") |
|
except Exception as e: |
|
print(f"An error occurred while reading {file_path}: {e}") |
|
|
|
return content |
|
|
|
|
|
@retry(wait=wait_fixed(2)) |
|
def completion_with_backoff(**kwargs): |
|
return openai.ChatCompletion.create(**kwargs) |
|
|
|
|
|
def gpt4( |
|
usr_prompt: str, sys_prompt: str = "", api_key: str = "", model: str = "gpt-4" |
|
) -> str: |
|
openai.api_key = api_key |
|
|
|
message = [ |
|
{"role": "system", "content": sys_prompt}, |
|
{"role": "user", "content": usr_prompt}, |
|
] |
|
|
|
response = completion_with_backoff( |
|
model=model, |
|
messages=message, |
|
temperature=0.2, |
|
max_tokens=1000, |
|
frequency_penalty=0.0, |
|
) |
|
|
|
return response["choices"][0]["message"]["content"] |
|
|
|
|
|
def select_best_tag( |
|
filtered_tags: List[str], object_to_place: str, api_key: str = "" |
|
) -> str: |
|
user_template = read_file_to_string("user_template.txt").format(object=object_to_place) |
|
user_prompt = user_template + "\n".join(filtered_tags) |
|
system_prompt = read_file_to_string("system_template.txt") |
|
return gpt4(user_prompt, system_prompt, api_key=api_key) |
|
|
|
|
|
def get_location_gsam( |
|
image: Image.Image, prompt: str, weights_folder="weights" |
|
) -> Tuple[int, int]: |
|
global sam_predictor |
|
|
|
BOX_TRESHOLD = 0.25 |
|
RESIZE_RATIO = 3 |
|
|
|
detections, phrases = get_gdino_result( |
|
image=image, |
|
classes=[prompt], |
|
box_threshold=BOX_TRESHOLD, |
|
) |
|
|
|
while len(detections.xyxy) == 0: |
|
BOX_TRESHOLD -= 0.02 |
|
detections, phrases = get_gdino_result( |
|
image=image, |
|
classes=[prompt], |
|
box_threshold=BOX_TRESHOLD, |
|
) |
|
|
|
sam_model = get_sam_model(weights_folder) |
|
|
|
if sam_predictor is None: |
|
print("Loading SAM Model...") |
|
sam_predictor = SamPredictor(sam_model) |
|
|
|
sam_predictor.set_image(np.array(image)) |
|
result_masks = [] |
|
for box in detections.xyxy: |
|
masks, scores, logits = sam_predictor.predict(box=box, multimask_output=True) |
|
index = np.argmax(scores) |
|
result_masks.append(masks[index]) |
|
detections.mask = np.array(result_masks) |
|
|
|
combined_mask = detections.mask[0] |
|
for mask in detections.mask[1:]: |
|
combined_mask += mask |
|
combined_mask[combined_mask > 1] = 1 |
|
mask = cv2.resize( |
|
combined_mask.astype("uint8"), |
|
( |
|
combined_mask.shape[1] // RESIZE_RATIO, |
|
combined_mask.shape[0] // RESIZE_RATIO, |
|
), |
|
) |
|
|
|
mask_2_pad = np.pad(mask, pad_width=2, mode="constant", constant_values=0) |
|
mask_1_pad = np.pad(mask, pad_width=1, mode="constant", constant_values=0) |
|
|
|
windows = np.lib.stride_tricks.sliding_window_view(mask_2_pad, (3, 3)) |
|
windows_all_zero = (windows == 0).all(axis=(2, 3)) |
|
|
|
result = np.where(windows_all_zero, 2, mask_1_pad) |
|
mask_0_coordinates = np.argwhere(result == 0) |
|
mask_1_coordinates = np.argwhere(result == 1) |
|
distances = cdist(mask_1_coordinates, mask_0_coordinates, "euclidean") |
|
max_min_distance_index = np.argmax(np.min(distances, axis=1)) |
|
y, x = mask_1_coordinates[max_min_distance_index] |
|
|
|
return int(x) * RESIZE_RATIO, int(y) * RESIZE_RATIO |
|
|
|
|
|
def run_octo_pipeline(input_image, object, api_key): |
|
print("Inside run_octo_pipeline with input_image=", input_image, "object=", object) |
|
|
|
print("Loading Image...") |
|
image = input_image.convert("RGB") |
|
|
|
print("Stage 1...") |
|
tags = get_tags_ram(image, threshold_multiplier=0.8) |
|
print("RAM++ Tags", tags) |
|
filtered_tags = filter_tags_gdino(image, tags) |
|
print("Filtered Tags", filtered_tags) |
|
|
|
print("Stage 2...") |
|
selected_tag = select_best_tag(filtered_tags, object, api_key=api_key) |
|
print("GPT-4 Selected Tag", selected_tag) |
|
|
|
print("Stage 3...") |
|
x, y = get_location_gsam(image, selected_tag) |
|
print("G-SAM Location", "(" + str(x) + "," + str(y) + ")") |
|
|
|
draw = ImageDraw.Draw(image) |
|
radius = 10 |
|
bbox = (x - radius, y - radius, x + radius, y + radius) |
|
draw.ellipse(bbox, fill="red") |
|
return [image] |
|
|
|
|
|
block = gr.Blocks() |
|
|
|
with block: |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image(type="pil", value="000749.png") |
|
object = gr.Textbox(label="Object", placeholder="Enter an object") |
|
api_key = gr.Textbox(label="OpenAI API Key", placeholder="Enter OpenAI API Key") |
|
|
|
with gr.Column(): |
|
gallery = gr.Gallery( |
|
label="Output", |
|
show_label=False, |
|
elem_id="gallery", |
|
preview=True, |
|
object_fit="scale-down", |
|
) |
|
|
|
iface = gr.Interface( |
|
fn=run_octo_pipeline, inputs=[input_image, object, api_key], outputs=gallery |
|
) |
|
iface.launch() |
|
|
