恩培-计算机视觉

 找回密码
 立即注册
搜索
查看: 1704|回复: 0

SAM测量不规则图形面积

[复制链接]

157

主题

103

回帖

3872

积分

管理员

Rank: 9Rank: 9Rank: 9

积分
3872
发表于 2023-9-3 10:57:01 | 显示全部楼层 |阅读模式

SAM测量不规则图形面积

项目 内容
更新人 @恩培
更新时间 2023-09-01
标题 SAM测量不规则图形面积
方法 Python、Opencv、SAM、FastSAM
相关付费内容 三、项目三:Segment Anything Model (SAM) 分割测量面积

[toc]

1. 视频/截图

https://www.bilibili.com/video/BV1Xk4y1w7xR/?spm_id_from=333.999.0.0

2. 核心代码

# -*- coding : utf-8 -*-
import cv2
import numpy as np
from sam_model import SegmentModel
import argparse
# 导入PIL
from PIL import Image, ImageDraw, ImageFont


class AreaCalculator:
    def __init__(self, img_file):
        metric_cofig_file = "./config/metric_cofig.txt"
        metric_points = self.get_image_metric(metric_cofig_file)
        metric_length = 427  # 根据实际情况设置,单位是m

        # 计算单位面积在图片中的像素个数
        self.area_per_pixel = self.calculate_area_per_pixel(metric_points, metric_length)

        # 加载模型
        model_path = "weights/sam_vit_h_4b8939.pth"
        self.model = SegmentModel(model_path=model_path, model_type="vit_h")
        # 读取图片
        self.image_file = img_file
        self.origin_img_data = cv2.imread(img_file)
        self.image_data = self.origin_img_data.copy()

        # 初始化image embedding
        self.model.init_embedding(img_file)

        self.segment_pts = []
        self.ignore_pts = []

    # 读取配置文件中的点
    def get_image_metric(self, metric_cofig_file): 
        # 文件每行一个点:x,y
        with open(metric_cofig_file, "r", encoding="utf8") as fr:
            lines = fr.readlines()
            points = []
            for line in lines:
                line = line.strip()
                if line == "":
                    continue
                x, y = line.split(",")
                points.append((int(x), int(y)))
            return points

    # 计算单位面积在图片中的像素个数
    def calculate_area_per_pixel(self, metric_points, metric_length):
        # 计算点之间的像素距离
        distance = np.linalg.norm(np.array(metric_points[0]) - np.array(metric_points[1]))
        print("distance is ", distance)
        # 计算单位面积在图片中的像素个数
        area_per_pixel = (distance * distance )/ (metric_length * metric_length)
        print("area_per_pixel is ", area_per_pixel)
        return area_per_pixel

    # 计算面积
    def calculate_area(self, last_stage = False):

        input_point = np.array(self.segment_pts + self.ignore_pts) # 输入点
        input_label = np.array([1] * len(self.segment_pts) + [0] * len(self.ignore_pts)) # 输入点的标签

        maskes, scores = self.model.predict(input_point, input_label) # 预测
        return_by_max_score = True  # False # 设置是否选取最大score作为输出,否则输出最大面积的mask
        areas = [mask.sum() for mask in maskes]
        # np.sum(mask == True)

        if return_by_max_score:
            idx = np.argmax(scores)  # 获取scores最大的索引
        else:
            idx = np.argmax(np.asarray(areas)) # 获取面积最大的索引

        mask = maskes[idx] # 获取mask
        # 获取mask中True的个数
        num_mask = mask.sum()


        area = num_mask / self.area_per_pixel  # 计算面积

        # 使用图片数组的方式贴图

        if last_stage:
            # 若mask为true,则将像素置为绿色,背景为黑色
            color = (0,255,0)
            self.image_data[mask > 0,0] = color[0]
            self.image_data[mask > 0,1] = color[1]
            self.image_data[mask > 0,2] = color[2]
            self.image_data[mask == 0] = 0
        else:
            # 若mask为true,则将像素置为半透明
            alpha = 0.5
            color = (255,0,255)
            self.image_data[mask > 0,0] = self.image_data[mask > 0,0] * alpha + color[0] * (1-alpha)
            self.image_data[mask > 0,1] = self.image_data[mask > 0,1] * alpha + color[1] * (1-alpha)
            self.image_data[mask > 0,2] = self.image_data[mask > 0,2] * alpha + color[2] * (1-alpha)

            # 绘制点
            for pt in self.segment_pts:
                cv2.circle(self.image_data, pt, 20, (0, 255, 0), -1)
            for pt in self.ignore_pts:
                cv2.circle(self.image_data, pt, 20, (0, 0, 255), -1)


        color = (0,0,0)
        alpha = 0.2
        l,t = 100, 100
        r,b = l+450,t+140
        self.image_data[t:b,l:r,0] = self.image_data[t:b,l:r,0] * alpha + color[0] * (1-alpha)
        self.image_data[t:b,l:r,1] = self.image_data[t:b,l:r,1] * alpha + color[1] * (1-alpha)
        self.image_data[t:b,l:r,2] = self.image_data[t:b,l:r,2] * alpha + color[2] * (1-alpha)

        self.image_data = self.cv2AddChineseText(self.image_data, "像素个数:" + str(int(num_mask)), (l+50,t+20), textColor=(255, 255, 255), textSize=40)
        if last_stage:
            # 显示面积
            area = round(area/1000000, 3)
            self.image_data = self.cv2AddChineseText(self.image_data, "面积:{} km^2".format(area), (l+50,t+80), textColor=(0, 255, 0), textSize=40)


    # 鼠标事件
    def set_seg_point_event(self, event, x, y, flags, params):
        if event == cv2.EVENT_LBUTTONUP: # 左键点击
            self.image_data = self.origin_img_data.copy()
            print("click L, " + str([x, y]))
            self.segment_pts.append((x, y))
            # 计算面积
            self.calculate_area()
        elif event == cv2.EVENT_RBUTTONUP: # 右键点击
            self.image_data = self.origin_img_data.copy()
            print("click R, " + str([x, y]))
            self.ignore_pts.append((x, y))
            # 计算面积
            self.calculate_area()


    # 绘制中文
    def cv2AddChineseText(self, img, text, position, textColor=(0, 255, 0), textSize=30):

        if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
            img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        # 创建一个可以在给定图像上绘图的对象
        draw = ImageDraw.Draw(img)
        # 字体的格式
        fontStyle = ImageFont.truetype(
            "./fonts/MSYH.TTC", textSize, encoding="utf-8")
        # 绘制文本
        draw.text(position, text, textColor, font=fontStyle)
        # 转换回OpenCV格式
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

    # 运行
    def run(self):
        cv2.namedWindow('image')
        cv2.setMouseCallback('image',self.set_seg_point_event) # 设置鼠标事件
        while(True):
            cv2.imshow('image',self.image_data)
            # ESC退出
            if cv2.waitKey(20) & 0xFF == 27:
                break
            # 重置
            elif cv2.waitKey(20) & 0xFF == ord('r'):
                self.image_data = self.origin_img_data.copy()
                self.ignore_pts = []
                self.segment_pts = []
            # enter计算
            elif cv2.waitKey(20) & 0xFF == 13:
                self.calculate_area(last_stage=True)

        cv2.destroyAllWindows()


# parse args
parser = argparse.ArgumentParser()
parser.add_argument("--img_file", type=str, default="./imgs/test.jpg")
args = parser.parse_args()
# init
area_app = AreaCalculator(args.img_file)
area_app.run()

3.代码及附件

链接:https://pan.baidu.com/s/1mYao_zjB_74Kaong_BKyrg?pwd=9v0r

用法:

  1. 下载权重:https://github.com/facebookresearch/segment-anything
  2. config_metrics.py配置比例尺
  3. sam_demo.py测量
回复

使用道具 举报

您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

恩培-计算机视觉

GMT+8, 2024-7-20 20:44 , Processed in 0.075457 second(s), 23 queries .

Powered by Discuz! X3.4

© 2001-2023 Discuz! Team.

快速回复 返回顶部 返回列表