STDC/predict.py

from models_725.segWaterBuilding import SegModel
from PIL import Image
import numpy as np
import cv2
import os
from cv2 import getTickCount, getTickFrequency
import time

def predict_lunkuo(impth=None):
    pred, probs = segmodel.eval(image=impth)#####
    preds_squeeze = pred.squeeze(0)
    preds_squeeze[preds_squeeze != 0] = 255
    preds_squeeze = np.array(preds_squeeze.cpu())
    preds_squeeze = np.uint8(preds_squeeze)
    _, 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(impth,cv2.COLOR_RGB2BGR)
    img2 = cv2.drawContours(img_n,contours,-1,(0,0,255),8)
    
    return img2

if __name__ == '__main__':
    impth = 'images/examples'
    outpth= 'images/results'
    folders = os.listdir(impth)
    #segmodel = SegModel(device='cuda:0')
    segmodel = SegModel(device='cpu')
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        time00 = time.time()
        img = Image.open(imgpath).convert('RGB')
        img = np.array(img)
        time11 = time.time()
        img=predict_lunkuo(impth=img)
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,img )
        print('----all_process', (time.time() - time11) * 1000)