22 lines
715 B
Python
22 lines
715 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 = "/home/xiazj/ai-station-code/tmp/dimaoshibie/d94afe94-2fae-4ce5-9ee4-94ac1d699337/dimao2.jpg"
|
|
predict(path) |