from PIL import Image
from utils import detect_image
import torch
from model import Unet_vgg

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
name_classes = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair",
                "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

model = Unet_vgg(num_classes=21, pretrained=True).to(device).eval()
checkpoint = torch.load('vgg_pretrain.pth')
model.load_state_dict(checkpoint)

img_path = './VOCdevkit/VOC2007/JPEGImages/2007_009052.jpg'
image = Image.open(img_path)
r_image = detect_image(model, image, device=device)

# r_image.show()
r_image.save('./predict.png')