Spaces:
Build error
Build error
from transformers import pipeline | |
from ultralytics import YOLO | |
from transformers import SamModel, SamProcessor | |
import torch | |
from PIL import Image | |
from diffusers.utils import load_image | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
def generate_mask_from_bbox(image_path): | |
""" | |
Generates a mask from the bounding box of an image using YOLO and SAM-ViT models. | |
Args: | |
image_path (str): The path to the input image. | |
Returns: | |
numpy.ndarray: The generated mask as a NumPy array. | |
""" | |
# Initialize YOLO and Processor | |
yolo = YOLO('yolov8m.pt') | |
processor = SamProcessor.from_pretrained('facebook/sam-vit-large') | |
model = SamModel.from_pretrained("facebook/sam-vit-large").to(device) | |
# Generate bounding boxes | |
results = yolo(image_path) | |
bboxes = results[0].boxes.xyxy.tolist() | |
input_boxes = [[[bboxes[0]]]] | |
# Process inputs | |
inputs = processor(load_image(image_path), input_boxes=input_boxes, return_tensors="pt").to("cuda") | |
# Perform inference | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
# Post-process masks | |
mask = processor.image_processor.post_process_masks( | |
outputs.pred_masks.cpu(), | |
inputs["original_sizes"].cpu(), | |
inputs["reshaped_input_sizes"].cpu() | |
)[0][0][0].numpy() | |
print(mask) | |
return mask | |