下面分别是原图,缺陷图,返回的二值化图像,实现代码在文后。
设置模型参数,加载权重文件的方式与web_demo.py中的相同
from model.openllama import OpenLLAMAPEFTModel
import torch
import numpy as np
import argparse
import matplotlib.pyplot as plt
import cv2
from PIL import Image as PILImage
import os
torch.cuda.set_device(5)
parser = argparse.ArgumentParser("AnomalyGPT", add_help=True)
command_args = parser.parse_args()
# init the model
args = {
'model': 'openllama_peft',
'imagebind_ckpt_path': 'pretrained_ckpt/imagebind_ckpt/imagebind_huge.pth',
'vicuna_ckpt_path': 'pretrained_ckpt/vicuna_ckpt/7b_v0',
'anomalygpt_ckpt_path': 'code/ckpt/train_supervised/pytorch_model.pt',
'delta_ckpt_path': 'pretrained_ckpt/pandagpt_ckpt/7b/pytorch_model.pt',
'stage': 2,
'max_tgt_len': 128,
'lora_r': 32,
'lora_alpha': 32,
'lora_dropout': 0.1
}
# 初始化模型
model = OpenLLAMAPEFTModel(**args)
# 加载模型权重
delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cuda:5'))
model.load_state_dict(delta_ckpt, strict=False)
delta_ckpt = torch.load(args['anomalygpt_ckpt_path'], map_location=torch.device('cuda:5'))
model.load_state_dict(delta_ckpt, strict=False)
model = model.eval().half().cuda()
print(f'[!] init the 7b model over ...')
下面是封装好的缺陷检测函数
def detect_anomalies(input_image_path, normal_image_path, max_length, top_p, temperature):
# 准备 prompt
prompt_text = "Is there any anomaly in the image?" # 这里简化了 prompt 的构建
if not os.path.exists(input_image_path):
print("Input image file does not exist.")
return None
if normal_image_path and not os.path.exists(normal_image_path):
print("Normal image file does not exist.")
return None
# 模型预测
try:
response, pixel_output = model.generate({
'prompt': prompt_text,
'image_paths': [input_image_path] if input_image_path else [],
'normal_img_paths': [normal_image_path] if normal_image_path else [],
'audio_paths': [],
'video_paths': [],
'thermal_paths': [],
'top_p': top_p,
'temperature': temperature,
'max_tgt_len': max_length,
'modality_embeds': [] # 假设这里不需要模态嵌入
}, web_demo=True)
except Exception as e:
print(f"Error in model generation: {e}")
return None
if pixel_output is None:
print("No output generated by the model.")
return None
return response,pixel_output,
打开本地缺陷图片和原始图片,设置好其他参数,调用检测函数,即可得到检测返回的文本信息和返回的缺陷图片
nput_image_path = 'defectTest/inpaintA.png'
normal_image_path = 'normalImage/U101_90500000483A_,rgb.179.756.png' # 如果没有正常图像,可以设置为 None
max_length = 512
top_p = 0.01
temperature = 1.0
image = cv2.imread(input_image_path)
image_width = image.shape[1]
image_height = image.shape[0]
response,pixel_output = detect_anomalies(input_image_path, normal_image_path, max_length, top_p, temperature)
对输出图像进行处理和保存
output = pixel_output.cpu().numpy().squeeze()
output = np.clip(output, 0, 1)
output = (output * 255).astype(np.uint8)
output = PILImage.fromarray(output)
output = output.resize((image_width, image_height), PILImage.LANCZOS)
output.save('output.png')
将输出图像转化为二值化图
gray_image = output.convert("L")
threshold =50
binary_image = gray_image.point(lambda x: 0 if x < threshold else 255, '1')
binary_image.save('binary_image.png')
获取缺陷区域的中心点坐标,可达到定位缺陷的目的
gray_image = output.convert("L")
threshold =50
binary_image = gray_image.point(lambda x: 0 if x < threshold else 255, '1')
binary_image.save('binary_image.png')
contours, hierarchy = cv2.findContours(cv2.imread('binary_image.png', cv2.IMREAD_GRAYSCALE), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 存储所有中心的坐标
centers = []
# 遍历所有轮廓
for contour in contours:
# 计算轮廓的矩
M = cv2.moments(contour)
# 检查矩是否有效(即,轮廓是否包含像素)
if M["m00"] != 0:
# 计算质心(中心)
cX = int(M["m10"] / M["m00"])
cY = int(M["m01"] / M["m00"])
# 将中心坐标添加到列表中
centers.append((cX, cY))
# 如果没有找到任何轮廓,则centers列表将为空
if not centers:
print("No white regions found.")
else:
# 打印所有中心的坐标
for center in centers:
print(f"Center: {center}")
想要输出图像保存原始大小,需要屏蔽图像剪裁操作,然后再将输出图像恢复原始大小 ,这会影响检测精确度,如果只追求精确度不需要恢复原始大小的可以忽略这一步
返回结果