19 lines
888 B
Python
19 lines
888 B
Python
from PIL import Image
|
|
import requests
|
|
from transformers import SamModel, SamProcessor
|
|
import os
|
|
os.environ['TRANSFORMERS_OFFLINE']="1"
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
|
|
|
|
model = SamModel.from_pretrained("weight/segment/SamModel",local_files_only=True).to("cuda")
|
|
processor = SamProcessor.from_pretrained("weight/segment/SamModel",local_files_only=True)
|
|
|
|
img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
|
|
raw_image = Image.open("/home/ykn/algorithm_system/flask_web/11.jpg").convert("RGB")
|
|
input_points = [[[450, 600]]] # 2D localization of a window
|
|
|
|
inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to("cuda")
|
|
outputs = model(**inputs)
|
|
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
|
|
scores = outputs.iou_scores
|