Spaces:
Paused
Paused
File size: 5,745 Bytes
e2e727d e796478 e2e727d 589f453 e2e727d 873583b e2e727d 18cec49 e2e727d ed51f1c e2e727d 47d4c3b 46bace6 e2e727d 47d4c3b e2e727d 46bace6 e2e727d e003fb9 e2e727d 46bace6 e2e727d 589f453 71e295f e2e727d 589f453 47d4c3b 589f453 e2e727d 589f453 e2e727d e83030d 589f453 e2e727d 8b87815 e2e727d 8b87815 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
#Script added by SPDraptor
import spaces
import copy
import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from typing import Any
import supervision as sv
from sam2.build_sam import build_sam2, build_sam2_video_predictor
from sam2.sam2_image_predictor import SAM2ImagePredictor
import time
device = torch.device('cuda')
model_id = 'microsoft/Florence-2-large'
models_dict = {
'Florence_model':AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval(),
'Florence_processor':AutoProcessor.from_pretrained(model_id, trust_remote_code=True),
}
SAM_CHECKPOINT = "/home/user/app/sam2_hiera_large.pt"
SAM_CONFIG = "sam2_hiera_l.yaml"
@spaces.GPU(duration=20)
def load_sam_image_model(
device: torch.device,
config: str = SAM_CONFIG,
checkpoint: str = SAM_CHECKPOINT
) -> SAM2ImagePredictor:
model = build_sam2(config, checkpoint)
return SAM2ImagePredictor(sam_model=model)
@spaces.GPU(duration=20)
def run_sam_inference(
model: Any,
image: Image,
detections: sv.Detections
) -> sv.Detections:
image = np.array(image.convert("RGB"))
model.set_image(image)
print(type(detections.xyxy),detections.xyxy)
if detections.xyxy.size == 0:
return {
'code': 400,
'data':'null',
'message':'The AI couldn’t detect the object you want to mask.'
}
mask, score, _ = model.predict(box=detections.xyxy, multimask_output=False)
# dirty fix; remove this later
if len(mask.shape) == 4:
mask = np.squeeze(mask)
detections.mask = mask.astype(bool)
return {
'code': 200,
'data':detections,
'message':'The AI couldn’t detect the object you want to mask.'
}
@spaces.GPU(duration=20)
def florence2(image,task_prompt, text_input=None):
"""
Calling the Microsoft Florence2 model
"""
model = models_dict['Florence_model']
processor = models_dict['Florence_processor']
# print(image)
if text_input is None:
prompt = task_prompt
else:
prompt = task_prompt + text_input
input_florence = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
print(input_florence)
generated_ids = model.generate(
input_ids=input_florence["input_ids"],
pixel_values=input_florence["pixel_values"],
max_new_tokens=1024,
early_stopping=False,
do_sample=False,
num_beams=3,
)
generated_text = processor.batch_decode(generated_ids,
skip_special_tokens=False)[0]
parsed_answer = processor.post_process_generation(
generated_text,
task=task_prompt,
image_size=(image.width, image.height))
return parsed_answer
def draw_MASK(image, prediction, fill_mask=False):
"""
Draws segmentation masks with polygons on an image.
Parameters:
- image_path: Path to the image file.
- prediction: Dictionary containing 'polygons' and 'labels' keys.
'polygons' is a list of lists, each containing vertices of a polygon.
'labels' is a list of labels corresponding to each polygon.
- fill_mask: Boolean indicating whether to fill the polygons with color.
"""
width=image.width
height=image.height
new_image = Image.new("RGB", (width, height), color="black")
draw = ImageDraw.Draw(new_image)
scale = 1
for polygons, label in zip(prediction['polygons'], prediction['labels']):
color = "white"
fill_color = "white" if fill_mask else None
for _polygon in polygons:
_polygon = np.array(_polygon).reshape(-1, 2)
if len(_polygon) < 3:
print('Invalid polygon:', _polygon)
continue
_polygon = (_polygon * scale).reshape(-1).tolist()
if fill_mask:
draw.polygon(_polygon, outline=color, fill=fill_color)
else:
draw.polygon(_polygon, outline=color)
draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)
return new_image
@spaces.GPU(duration=20)
def masking_process(image,obj):
# task_prompt = '<REGION_TO_SEGMENTATION>'
# # task_prompt = '<OPEN_VOCABULARY_DETECTION>'
# print(type(task_prompt),type(obj))
# print('1')
start_time = time.time()
image = Image.fromarray(image).convert("RGB")
# results = florence2(image,task_prompt, text_input=obj)
# output_image = copy.deepcopy(image)
# img=draw_MASK(output_image,
# results['<REGION_TO_SEGMENTATION>'],
# fill_mask=True)
# mask=img.convert('1')
task_prompt = '<OPEN_VOCABULARY_DETECTION>'
# image = Image.open("/content/tiger.jpeg").convert("RGB")
# obj = "Tiger"
Florence_results = florence2(image,task_prompt, text_input=obj)
# print('2')
SAM_IMAGE_MODEL = load_sam_image_model(device=device)
# print('3')
detections = sv.Detections.from_lmm(
lmm=sv.LMM.FLORENCE_2,
result=Florence_results,
resolution_wh=image.size
)
# print('4')
response = run_sam_inference(SAM_IMAGE_MODEL, image, detections)
print(f'Time taken by masking model: {time.time() - start_time}')
# print('5')
if response['code'] == 400:
print("no object found")
return "no object found"
else:
detections2=response['data']
mask = Image.fromarray(detections2.mask[0])
# response['data']=mask
torch.cuda.empty_cache()
return mask |