利用前景和背景的mask图来得到边缘平滑的图像（graphcut实现）

github上的代码都是交互式的，我这个是通过mask来生成标签，无需交互

原图：

边界是不清晰的，我首先得到两种区域的mask图，然后腐蚀一下

左图是背景的mask图，右图是copper的mask图

腐蚀是针对图片中的白色部分

main.py代码如下：(除此之外我写了一个GraphMaker类放在graphcut.py脚本上)

import cv2
import numpy as np
from matplotlib import pyplot as plt
import os.path as osp
import os
from tqdm import tqdm
import numpy as np

import cv2
import networkx as nx
import math
from  graphcut import GraphMaker


def erode_dilate(img_path):
    """对copper进行腐蚀
    copper区域腐蚀一下，就是前景的hint
    背景区域腐蚀一下就是背景的hint
    """
    img_ori = cv2.imread(img_path)
    #img_ori =  img_ori*255
    # 开始进行腐蚀操作
    retVal, image = cv2.threshold(img_ori, 20, 255, cv2.THRESH_BINARY)
    corrosion_img = cv2.getStructuringElement(cv2.MORPH_CROSS, (1, 1))  ##腐蚀预处理，确定处理核的大小,矩阵操作
    pic_matrix = cv2.erode(image, corrosion_img, iterations=10)  # 进行腐蚀操作

    return pic_matrix


def get_new_img(img_end, save_p):
    cv2.imwrite(save_p, img_end)


def get_labels(background_hint,copper_hint):
    # 得到标签
    size = background_hint.shape[:2]
    labels = np.zeros(size)
    for i in tqdm(range((background_hint.shape[0]))):
        for j in range(background_hint.shape[1]):
            if background_hint[i,j][0] ==255:
                labels[i][j] = 1
                # background 通过mask定义标签
                
                graphcut_class.add_seed(i,j,0)
    for i in tqdm(range(copper_hint.shape[0])):
        for j in range(copper_hint.shape[1]):
            if copper_hint[i, j][0] == 255:
                labels[i][j] = -1
                # foreground 通过mask定义标签
                graphcut_class.add_seed(i, j, 1)
    return labels



if __name__ == '__main__':
    
    
    ori_folder = '/cloud_disk/users/huh/pcb/seg_test_0218/ddrslim_pcbv1'
    img_list = os.listdir(ori_folder)
    ori_list = []
    for img_name in img_list:
        if 'img' in img_name:
            ori_list.append(img_name[:-7])
    end_folder = '/cloud_disk/users/huh/pcb/seg_test_0218/test'
    for img_name in tqdm(ori_list):
        # 学姐搞反了
        fore_img = osp.join(ori_folder,img_name+'dust.png')
        background_img = osp.join(ori_folder,img_name+'copper.png')
        ori_img = osp.join(ori_folder,img_name+'img.png')
        os.system('cp {} {}'.format(ori_img,end_folder))
        fore_matrix = erode_dilate(fore_img)
        background_matrix = erode_dilate(background_img)
        # 最后生成图片的地址
        end_path = osp.join(end_folder,img_name+'end.png')
        # # 初始化graphcut类
        graphcut_class = GraphMaker(ori_img,end_path)
        get_labels(background_matrix,fore_matrix)

        graphcut_class.create_graph()
        graphcut_class.cut_graph()

        graphcut_class.save_image(end_path)






    

graphcut.py脚本代码如下：

import cv2
import numpy as np
import maxflow


class GraphMaker:

    foreground = 1
    background = 0

    seeds = 0
    segmented = 1

    # default = 0.5
    default = 0.3
    MAXIMUM = 1000000000

    def __init__(self,filename,outfilename):
        self.image = None
        self.graph = None
        self.overlay = None
        self.seed_overlay = None
        self.segment_overlay = None
        self.mask = None
        self.load_image(filename)
        self.background_seeds = []
        self.foreground_seeds = []
        self.background_average = np.array(3)
        self.foreground_average = np.array(3)
        self.nodes = []
        self.edges = []
        self.current_overlay = self.seeds
        self.outfilename = outfilename

    def load_image(self, filename):
        self.image = cv2.imread(filename)
        self.graph = np.zeros_like(self.image)
        self.seed_overlay = np.zeros_like(self.image)
        self.segment_overlay = np.zeros_like(self.image)
        self.mask = None







    def add_seed(self, x, y, type):
        if self.image is None:
            print('Please load an image before adding seeds.')
        if type == self.background:
            if not self.background_seeds.__contains__((x, y)):
                self.background_seeds.append((x, y))
                cv2.rectangle(self.seed_overlay, (x-1, y-1), (x+1, y+1), (0, 0, 255), -1)
        elif type == self.foreground:
            if not self.foreground_seeds.__contains__((x, y)):
                self.foreground_seeds.append((x, y))
                cv2.rectangle(self.seed_overlay, (x-1, y-1), (x+1, y+1), (0, 255, 0), -1)


    def clear_seeds(self):
        self.background_seeds = []
        self.foreground_seeds = []
        self.seed_overlay = np.zeros_like(self.seed_overlay)

    def get_overlay(self):
        if self.current_overlay == self.seeds:
            return self.seed_overlay
        else:
            return self.segment_overlay

    def get_image_with_overlay(self, overlayNumber):
        if overlayNumber == self.seeds:
            return cv2.addWeighted(self.image, 0.9, self.seed_overlay, 0.4, 0.1)
        else:
            return cv2.addWeighted(self.image, 0.9, self.segment_overlay, 0.4, 0.1)

    def create_graph(self):
        if len(self.background_seeds) == 0 or len(self.foreground_seeds) == 0:
            print("Please enter at least one foreground and background seed.")
            return

        print("Making graph")
        print("Finding foreground and background averages")
        self.find_averages()

        print("Populating nodes and edges")
        self.populate_graph()

        print("Cutting graph")
        self.cut_graph()

    def find_averages(self):
        self.graph = np.zeros((self.image.shape[0], self.image.shape[1]))
        print(self.graph.shape)
        self.graph.fill(self.default)
        self.background_average = np.zeros(3)
        self.foreground_average = np.zeros(3)

        for coordinate in self.background_seeds:
            self.graph[coordinate[1] - 1, coordinate[0] - 1] = 0
            #self.background_average += self.image[coordinate[1], coordinate[0]]

        #self.background_average /= len(self.background_seeds)

        for coordinate in self.foreground_seeds:
            self.graph[coordinate[1] - 1, coordinate[0] - 1] = 1
            #self.foreground_average += self.image[coordinate[1], coordinate[0]]

        #self.foreground_average /= len(self.foreground_seeds)

    def populate_graph(self):
        self.nodes = []
        self.edges = []

        # make all s and t connections for the graph
        for (y, x), value in np.ndenumerate(self.graph):
            # this is a background pixel
            if value == 0.0:
                self.nodes.append((self.get_node_num(x, y, self.image.shape), self.MAXIMUM, 0))

            # this is a foreground node
            elif value == 1.0:
                self.nodes.append((self.get_node_num(x, y, self.image.shape), 0, self.MAXIMUM))

            else:
                '''d_f = np.power(self.image[y, x] - self.foreground_average, 2)
                d_b = np.power(self.image[y, x] - self.background_average, 2)
                d_f = np.sum(d_f)
                d_b = np.sum(d_b)
                e_f = d_f / (d_f + d_b)
                e_b = d_b / (d_f + d_b)'''
                self.nodes.append((self.get_node_num(x, y, self.image.shape), 0, 0))

                '''if e_f > e_b:
                    self.graph[y, x] = 1.0
                else:
                    self.graph[y, x] = 0.0'''

        for (y, x), value in np.ndenumerate(self.graph):
            if y == self.graph.shape[0] - 1 or x == self.graph.shape[1] - 1:
                continue
            my_index = self.get_node_num(x, y, self.image.shape)

            neighbor_index = self.get_node_num(x+1, y, self.image.shape)
            g = 1 / (1 + np.sum(np.power(self.image[y, x] - self.image[y, x+1], 2)))
            self.edges.append((my_index, neighbor_index, g))

            neighbor_index = self.get_node_num(x, y+1, self.image.shape)
            g = 1 / (1 + np.sum(np.power(self.image[y, x] - self.image[y+1, x], 2)))
            self.edges.append((my_index, neighbor_index, g))

    def cut_graph(self):
        self.segment_overlay = np.zeros_like(self.segment_overlay)
        self.mask = np.zeros_like(self.image, dtype=bool)
        g = maxflow.Graph[float](len(self.nodes), len(self.edges))
        nodelist = g.add_nodes(len(self.nodes))

        for node in self.nodes:
            g.add_tedge(nodelist[node[0]], node[1], node[2])

        for edge in self.edges:
            g.add_edge(edge[0], edge[1], edge[2], edge[2])

        flow = g.maxflow()

        for index in range(len(self.nodes)):
            if g.get_segment(index) == 1:
                xy = self.get_xy(index, self.image.shape)
                self.segment_overlay[xy[1], xy[0]] = (255, 0, 255)
                #self.mask[xy[1], xy[0]] = (True, True, True)
                self.mask[xy[0], xy[1]] = (True, True, True)

    def swap_overlay(self, overlay_num):
        self.current_overlay = overlay_num

    def save_image(self, outfilename):
        if self.mask is None:
            print('Please segment the image before saving.')
            return
        print(outfilename)
        # print(self.image.name())
        to_save = np.zeros_like(self.image)

        np.copyto(to_save, self.image, where=self.mask)
        cv2.imwrite(outfilename, to_save)

    @staticmethod
    def get_node_num(x, y, array_shape):
        return y * array_shape[1] + x

    @staticmethod
    def get_xy(nodenum, array_shape):
        return (nodenum % array_shape[1]), (int(nodenum / array_shape[1]))