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/1tyc66hb907ez_ODTD67u4Q?pwd=bvry 提取码: bvry
用法：
- 下载权重：https://github.com/facebookresearch/segment-anything
config_metrics.py配置比例尺sam_demo.py测量
发表于 2023-9-3 10:57:01