segment anything原来可以这么玩

小王不头秃

发布于 2024-06-19 09:32:49

26700

代码可运行

文章被收录于专栏：后端/图像融合/机器学习/爬虫后端/图像融合/机器学习/爬虫

运行总次数：0

代码可运行

前言

最近，大模型的热度确实是非常非常的高，从chatgpt到segment anything，这些东西整的我这刚入门的小白确实有点懵逼。最近实在是不知道干啥，想想能不能用大模型整点花活。

这时候突然想起来了去年的极乐净土，确实做的有点粗糙，去年就是提取了下边缘加个rgb就完事了，这让我有点不大尽兴，直到我看到了小鸡舞。 segment anything是干啥的，不就是图像分割吗，那我是不是可以把这位舞者分割出来，然后换到其他背景里，说干就干。

内容

大致的思路其实就是我们用segment anything把图像分割出来，然后把人的掩码单独拿出来，然后再随便整个背景，把我们提取的人物放进去就完事了，整个过程看着还是蛮简单的哈，我们具体看看咋做。

首先第一步就是要配置下segment anything的环境，这里我们不使用segment anything官方提供的服务跑了，因为确实要跑好多张，所以我们自己配一下。其实很简单就是安装一下pytorch gpu版本（显卡稍微差点的话可以用cpu的），然后把项目从github拉下来，在安装下他要求的库就ok了，网上有很多配置segment anything 的博客，这里我就不展开说了。

这时候我们把视频转换成图片，就直接用opencv转换就可以，详细代码在下一节。然后我们使用segment anything对其进行分割，就可以拿到这样的掩码

我们可以看到，分割的效果非常nice，但是有一个问题，我怎么把人的掩码单独提出来呢？？？

博主找了一下，好像是已经有人在segment anything的基础上进行分类了，但确实有点麻烦，博主本着能坐着就不能站着，能躺着就不坐着的心态思考了下，发现还真有个个省事。

要知道，我们可是有yolo的呀，这目标检测能力嘎嘎的

那我们就可以直接提取检测框里的最大掩码就可以了，这样我们就拿到了人物的掩码

到此为止，人物就提取出来了，接下来我们就是随便找个背景拼接下就完事了，看看具体实现吧。

具体实现

因为segment anything和yolov5都有现成的代码，我们就不介绍了，我们只介绍这个拼接的部分

拼接代码

# -*- codeing = utf-8 -*-
# @Time : 2023/7/5 19:30
# @Author : xiaow
# @File : test.py
# @Software : PyCharm
import os

import cv2

import numpy as np


def mix():
    # 背景图片存放的位置
    backs = os.listdir('../video2img2')
    # 人物掩码的位置
    masks = os.listdir('mask')
    # 0 1 掩码的位置
    mask2s = os.listdir('mask2')
    # 设置输出视频的帧数
    fps = 15
    videopath = 'test10.avi'  # 图片保存地址及格式
    size = (1000, 666)
    out1 = cv2.VideoWriter(videopath, cv2.VideoWriter_fourcc(*'DIVX'), fps, size)

    for i in range(len(backs)):
        print(masks[i])

        back = cv2.imread('../video2img2/' + backs[i])
        # 修改背景尺寸大小
        back = cv2.resize(back, size, interpolation=cv2.INTER_AREA)
        human = cv2.imread('mask/' + masks[i])
        mask2 = cv2.imread('mask2/' + mask2s[i], 0)


        # 设置mask,human大小和背景相同大小  start
        width = back.shape[0]
        height = back.shape[1]
        width_diff = width - human.shape[0]
        height_diff = height - human.shape[1]
        human = np.pad(human, ((width_diff // 2, width_diff // 2), (height_diff // 2, height_diff // 2), (0, 0)))
        mask2 = np.pad(mask2, ((width_diff // 2, width_diff // 2), (height_diff // 2, height_diff // 2)))
        # 设置mask,human大小和背景相同大小  end

        mask2 = np.expand_dims(mask2, 2)
        out1.write((1 - mask2) * back + human)


if __name__ == '__main__':
    mix()