Spaces:

nnibras
/

LAB4

Build error

LAB4 / src /segmentation.py

Upload 17 files

6aab31e verified almost 2 years ago

963 Bytes

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)

	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)