Views
No views yet
1import torch
2import torchvision.models.segmentation as models
3from torchvision import transforms
4from PIL import Image
5import cv2
6import numpy as np
7
8# Load the model architecture
9model = models.deeplabv3_resnet50(pretrained=False) # Set pretrained=False as we load custom weights
10model.classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=1) # Adjust output channels for 2 classes
11
12# Load the state dictionary
13model_path = "deeplabv3_resnet50_offroad.pth" # Path to your saved model
14model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
15model.eval()
16
17# Preprocessing transformations
18transform = transforms.Compose([
19 transforms.ToTensor(),
20 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
21])
22
23def predict_mask(image_path, model, transform):
24 image = Image.open(image_path).convert("RGB")
25 original_size = image.size
26
27 # Resize to model's expected input size (960x540 for this model, or handle dynamic resizing)
28 # For simplicity, assuming model input was trained on fixed size, let's resize
29 image_tensor = transform(image).unsqueeze(0) # Add batch dimension
30
31 with torch.no_grad():
32 output = model(image_tensor)['out']
33
34 # Get the predicted class for each pixel
35 predicted_mask = torch.argmax(output.squeeze(), dim=0).cpu().numpy()
36
37 # Resize mask back to original image size if necessary
38 predicted_mask_resized = cv2.resize(predicted_mask.astype(np.uint8), original_size, interpolation=cv2.INTER_NEAREST)
39
40 return predicted_mask_resized
41
42# Example usage:
43# Assuming you have an image 'test_image.jpg'
44# mask = predict_mask('test_image.jpg', model, transform)
45# plt.imshow(mask)
46# plt.show()