trafficDetectionHandover/model_predict.py

from models_711.segWaterBuilding import SegModel
from PIL import Image
from torchvision.transforms import transforms
import numpy as np
import cv2
import os
from cv2 import getTickCount, getTickFrequency
import matplotlib.pyplot as plt

def predict_lunkuo(impth=None):
    # segmodel = SegModel()
    loop_start = getTickCount()

    pred = segmodel.eval(image=img)

    loop_time = cv2.getTickCount() - loop_start
    tool_time = loop_time / (cv2.getTickFrequency())
    running_fps = int(1 / tool_time)
    print('running_fps:', running_fps)

    preds_squeeze = pred.squeeze(0)
    preds_squeeze[preds_squeeze != 0] = 255
    preds_squeeze = np.array(preds_squeeze.cpu())
    preds_squeeze = np.uint8(preds_squeeze)
    #print('preds_squeeze:', preds_squeeze.shape)
    _, binary = cv2.threshold(preds_squeeze,220,255,cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    img_n = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
    img2 = cv2.drawContours(img_n,contours,-1,(0,0,255),8)
    # save_path = './' + '00000000000000000000000000001' + '.png'
    # cv2.imshow('image',img2)
    # cv2.waitKey(0)
    plt.figure()
    plt.imshow(img2[:,:,[2,1,0]])
    # plt.show()

# if __name__ == '__main__':
#     impth = "/home/data/lijiwen/wurenjiqifei/images/20211225巡河_10.jpg"
#     # to_tensor = transforms.Compose([
#     #     transforms.ToTensor(),
#     #     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
#     # ])
#     img = Image.open(impth).convert('RGB')
#     predict_lunkuo(impth=impth)
if __name__ == '__main__':
    impth = '/home/data/lijiwen/wurenjiqifei/bu711/'
    segmodel = SegModel()
    folders = os.listdir(impth)
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        img = Image.open(imgpath).convert('RGB')
        predict_lunkuo(impth=impth)