22 lines
680 B
Python
22 lines
680 B
Python
|
import time
|
||
|
|
import os
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from PIL import Image
|
||
|
|
from tqdm import tqdm
|
||
|
|
from segformer import SegFormer_Segmentation
|
||
|
|
|
||
|
|
def model():
|
||
|
|
return SegFormer_Segmentation()
|
||
|
|
|
||
|
|
def predict(img):
|
||
|
|
SegFormer = model()
|
||
|
|
count = True
|
||
|
|
name_classes = ["_background_", "Cropland", 'Forest', 'Grass', 'Shrub', 'Wetland', 'Water', 'Tundra', 'Impervious surface', 'Bareland', 'Ice/snow']
|
||
|
|
image = Image.open(img)
|
||
|
|
r_image, count_dict, classes_nums = SegFormer.detect_image(image, count=count, name_classes=name_classes)
|
||
|
|
r_image.show()
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
path = "D:\\project\\ai-station\\tmp\\dimaoshibie\\crop_9_14.tif"
|
||
|
|
predict(path)
|