File size: 963 Bytes
6aab31e |
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 |
import torch
from torchvision import models, transforms
from PIL import Image
import numpy as np
def segment_person(image_path):
# Load the pre-trained DeepLabV3 model
model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
# Load and preprocess the input image
input_image = Image.open(image_path).convert("RGB")
preprocess = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
input_tensor = preprocess(input_image).unsqueeze(0)
with torch.no_grad():
output = model(input_tensor)["out"][0]
mask = output.argmax(0).byte().numpy()
# Convert mask to an image with transparency
segmented_image = np.array(input_image)
segmented_image = np.dstack([segmented_image, mask * 255]) # Add alpha channel
return Image.fromarray(segmented_image)
|