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)