Skip to content
ModelListedExternal
Download

Aisegmentation

This is a semantic segmentation model trained to identify offroad terrains from images.

KindOtherBase modeldeeplabv3_resnet50VersionRevision 69bc659

Model source

Kind
Other
File size
161 MB
Base model
deeplabv3_resnet50
Version
vRevision 69bc659
Updated
Sep 15, 2026
Downloads
0
Tasks
Image segmentation
Source
Hugging Face

Source description

Offroad Terrain Segmentation Model

This is a semantic segmentation model trained to identify offroad terrains from images.

Model Details

  • Model Architecture: DeepLabV3 with ResNet50 backbone
  • Pre-training: Initialized with weights pre-trained on COCO dataset
  • Dataset: 'Offroad_Segmentation_Training_Dataset'
  • Input: RGB images (540x960 pixels)
  • Output: Segmentation mask with 2 classes (e.g., background, offroad terrain)
  • Training Epochs: 10
  • Batch Size: 16
  • Learning Rate: 0.001

Training Metrics

  • Final Training Loss: 0.0682
  • Final Validation Loss: 0.0785
  • Final Training Mean IoU: 0.2789
  • Final Validation Mean IoU: 0.2690

How to use

import torch
import torchvision.models.segmentation as models
from torchvision import transforms
from PIL import Image
import cv2
import numpy as np

# Load the model architecture
model = models.deeplabv3_resnet50(pretrained=False) # Set pretrained=False as we load custom weights
model.classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=1) # Adjust output channels for 2 classes

# Load the state dictionary
model_path = "deeplabv3_resnet50_offroad.pth" # Path to your saved model
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()

# Preprocessing transformations
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def predict_mask(image_path, model, transform):
    image = Image.open(image_path).convert("RGB")
    original_size = image.size

    # Resize to model's expected input size (960x540 for this model, or handle dynamic resizing)
    # For simplicity, assuming model input was trained on fixed size, let's resize
    image_tensor = transform(image).unsqueeze(0) # Add batch dimension

    with torch.no_grad():
        output = model(image_tensor)['out']

    # Get the predicted class for each pixel
    predicted_mask = torch.argmax(output.squeeze(), dim=0).cpu().numpy()

    # Resize mask back to original image size if necessary
    predicted_mask_resized = cv2.resize(predicted_mask.astype(np.uint8), original_size, interpolation=cv2.INTER_NEAREST)

    return predicted_mask_resized

# Example usage:
# Assuming you have an image 'test_image.jpg'
# mask = predict_mask('test_image.jpg', model, transform)
# plt.imshow(mask)
# plt.show()

Sources

1 source
Hugging FacePrimaryVersion 69bc659279078c8cd2a540b3158e366e4ad88489
huggingface.co/yogikushle/aisegmentation/tree/69bc659279078c8cd2a540b3158e366e4ad88489

Verified Sep 15

Model artifacts

1 artifact

deeplabv3_resnet50_offroad.pth

pth · 161 MB · SHA-256 331c810f137b…edf5 · Hugging Face

Download

Source excerpts

2 excerpts

--- base_model: deeplabv3_resnet50 model_name: offroad_segmentation tags: - image-segmentation - pytorch - computer-vision - deeplabv3 widget: - src: https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/widget-images/image-segmentation.png example_title: Image segmentation example --- # Offroad Terrain Segmentation Model This is a semantic segmentation model trained to identify offroad terrains from images. ## Model Details - **Model Architecture**: DeepLabV3 with ResNet50 backbone - **Pre-training**: Initialized with weights pre-trained on COCO dataset - **Dataset**: 'Offroad_Segmentation_Training_Dataset' - **Input**: RGB images (540x960 pixels) - **Output**: Segmentation mask with 2 classes (e.g., background, offroad terrain) - **Training Epochs**: 10 - **Batch Size**: 16 - **Learning Rate**: 0.001 ## Training Metrics - **Final Training Loss**: 0.0682 - **Final Validation Loss**: 0.0785 - **Final Training Mean IoU**: 0.2789 - **Final Validation Mean IoU**: 0.2690 ## How to use ```python import torch import torchvision.models.segmentation as models from torchvision import transforms from PIL import Image import cv2 import numpy as np # Load the model architecture model = models.deeplabv3_resnet50(pretrained=False) # Set pretrained=False as we load custom weights model.classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=1) # Adjust output channels for 2 classes # Load the state dictionary model_path = "deeplabv3_resnet50_offroad.pth" # Path to your saved model model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) model.eval() # Preprocessing transformations transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) def predict_mask(image_path, model, transform): image = Image.open(image_path).convert("RGB") original_size = image.size # Resize to model's expected input size (960x540 for this model, or handle dynamic resizing) # For simplicity, assuming model input was trained on fixed size, let's resize image_tensor = transform(image).unsqueeze(0) # Add batch dimension with torch.no_grad(): output = model(image_tensor)['out'] # Get the predicted class for each pixel predicted_mask = torch.argmax(output.squeeze(), dim=0).cpu().numpy() # Resize mask back to original image size if necessary predicted_mask...

Source context: 0 downloads · 0 likes · Pipeline image-segmentation · Repo yogikushle/aisegmentation