import os
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import torchvision
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from torch import threshold
img_path = r'D:/卷积神经网络的遥感图像的分割/代码/test.png'
model = load_model('model_UCMerced_LandUse.h5')
model.summary()
#定义给出的类
COCO_INSTANCE_CATEGORY_NAMES = [
'建筑物','住宅区','工厂']
#定义一个新函数get_prediction(),用于输入图像路径,对图像中的物体进行识别。
def get_prediction(img_path, threshold):
img = Image.open(img_path)
transform = T.Compose([T.ToTensor()]) #将图片二值化处理
img = transform(img)
pred = model([img])
pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
pred_score = list(pred[0]['scores'].detach().numpy())
pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
pred_boxes = pred_boxes[:pred_t + 1]
pred_class = pred_class[:pred_t + 1]
print("pred_class:", pred_class)
print("pred_boxes:", pred_boxes)
return pred_boxes, pred_class
#定义一个新函数object_detection_api(),用于调用上面定义的函数get_prediction(),得到图像中的物体识别结果,并将识别结果绘制出来
def object_detection_api(img_path, threshold=0.5, rect_th=3, text_size=3, text_th=3):
boxes, pred_cls = get_prediction(img_path, threshold)
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for i in range(len(boxes)):
cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)
cv2.putText(img,pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
plt.imshow(img)
plt.show()
if name == 'main':
object_detection_api(img_path='D:/卷积神经网络的遥感图像的分割/代码/test.png')
报错:AttributeError: 'torch.Size' object has no attribute 'rank'