pytorch 如何将分割结果与另一幅图像叠加

jhdbpxl9  于 2023-05-22  发布在  其他
关注(0)|答案(1)|浏览(172)

这是我的输入图像:

我有这样的代码来覆盖输入图像的分割颜色:

# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
input_image = Image.open("samping.JPG")
input_image = input_image.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)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

# move the input and model to GPU for speed if available
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    model.to('cuda')

with torch.no_grad():
    output = model(input_batch)['out'][0]
output_predictions = output.argmax(0)

# create a color pallette, selecting a color for each class
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
colors = (colors % 255).numpy().astype("uint8")

# plot the semantic segmentation predictions of 21 classes in each color
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
r.putpalette(colors)

然后我得到这个结果:这是我的分割图像叠加颜色:

但是,我想用另一个图像覆盖/改变颜色,例如这个图像:

我想要的结果图像:


改变分割板(颜色)到我的标志,我怎么才能实现呢?

mwkjh3gx

mwkjh3gx1#

欢迎来到我的世界,糊涂虫!
OpenCV是一个计算机视觉库,非常适合这类任务。您可以使用pip install opencv-python安装它。简而言之,以下是所需的步骤:
1.找出黑色像素的位置。

1.找出最大的黑色像素块存在的地方,通过将它们转换为轮廓,然后计算它们各自的大小。(注意右侧背面悬挂的几个黑色像素,用绿色突出显示)

1.找到与最大轮廓最匹配的旋转矩形。

1.调整覆盖图像的大小并旋转它。在覆盖图像周围创建一个衬垫,以便在旋转它时它不会被切断。

(使用的图像是Gary Larson's "Cow Tools",我使用了一个高大的图像来展示缩放属性)
1.将生成的叠加图像粘贴到原始图像的顶部。在计算位置时一定要考虑到垫子,否则它会走到你打算去的东南方。

瞧!

import cv2
import numpy as np

# Load images
img = cv2.imread('stack_car.jpg')
# Replace the overlay image with your own,
# for now I will use Gary Larson's "Cow Tools" because it showcases the
# scaling property, what a fantastic image though
overlay_img = cv2.imread('cow_tools.png', -1) 

# Convert image to grayscale for easier thresholding
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# Threshold the image
_, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY_INV)

# Show the threshold image
cv2.imshow('Threshold Image', thresh)
cv2.waitKey(0)

# Find contours
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# Draw contours on a copy of the image for visualization
img_contours = img.copy()
cv2.drawContours(img_contours, contours, -1, (0, 255, 0), 3)
cv2.imshow('Contours', img_contours)
cv2.waitKey(0)

# Find the largest chunk of black pixels
max_contour = max(contours, key=cv2.contourArea)

# Get a rotated box around the largest contour
rotated_rect = cv2.minAreaRect(max_contour)

# draw the rotated rectangle on a copy of the image for visualization
img_rotated_rect = img.copy()
box = cv2.boxPoints(rotated_rect)
box = np.int0(box)
cv2.drawContours(img_rotated_rect, [box], 0, (0, 0, 255), 2)
cv2.imshow('Rotated Rectangle', img_rotated_rect)
cv2.waitKey(0)

# Get the rotation angle, width, and height from the rotated rectangle
center, (width, height), angle = rotated_rect

# Resize the overlay image to match the rotated rectangle
overlay_img_resized = cv2.resize(overlay_img, (int(width), int(height)))

# Add some transparent padding to the image so that it doesn't get cut off when it rotates
# (I'm not sure of the best math for pad, but overshooting is better than undershooting)
pad = int(max(width, height))
overlay_img_resized = cv2.copyMakeBorder(overlay_img_resized, pad, pad, pad, pad, cv2.BORDER_CONSTANT, value=(0, 0, 0, 0))

# Rotate the image by the angle
M = cv2.getRotationMatrix2D((overlay_img_resized.shape[1] / 2, overlay_img_resized.shape[0] / 2), -angle, 1)
overlay_img_resized = cv2.warpAffine(overlay_img_resized, M, (overlay_img_resized.shape[1], overlay_img_resized.shape[0]))

# Show the resized overlay image
cv2.imshow('Resized Overlay Image', overlay_img_resized)
cv2.waitKey(0)

# Overlay the rotated and resized image onto the original image,
# make sure to take pad into account, otherwise it will drift southeast
for i in range(overlay_img_resized.shape[0]):
    for j in range(overlay_img_resized.shape[1]):
        if overlay_img_resized[i, j][3] != 0:
            img[int(center[1]-height/2)-pad+i, int(center[0]-width/2)-pad+j] = overlay_img_resized[i, j][:3]

# Show the final image
cv2.imshow('Final Image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

相关问题