Skip to content
master
Go to file
Code
This branch is 8 commits ahead of clcarwin:master.

Latest commit

 

Git stats

Files

Permalink
Failed to load latest commit information.
Type
Name
Latest commit message
Commit time
 
 
 
 
 
 
 
 

README.md

ENet-PyTorch

Fork based on convert_torch_to_pytorch and adapted for image segmentation with ENet in PyTorch.

Conversion from torch ENet

  • Convert the source torch7 trained model from cutorch/cudnn weights to CPU weights for conversion:
$ th convert_torch_gpu_to_cpu.lua -model model-best.net
saved cpu-model-best.net
  • Convert the lua model to a pytorch script and state dict:
$ python convert_torch.py --model cpu-model-best.net --output enet_pytorch

This produces a standalone script enet_pytorch.py defining a torch.nn module, and an associated state dict.

Example usage:

import torch
from torch.autograd import Variable
import torch.nn.functional as F
import cv2
from enet_pytorch import enet_pytorch

enet_pytorch.load_state_dict(torch.load('enet_pytorch.pth'))
enet_pytorch.cuda().eval()

img = 'frankfurt_000001_030067_leftImg8bit.png'
im = cv2.imread(img).astype(np.float32)[:, :, ::-1]/255
inp = torch.from_numpy(im.transpose(2, 0, 1)).unsqueeze(0).cuda()
inp = F.upsample(Variable(inp), (512, 1024), mode='bilinear').data

out = enet_pytorch(Variable(inp))

Tested to reproduce ENet original results.

Model and weights

For reference the output on the author's cityscapes model and associated converted weights are provided on https://goo.gl/Rvtqxr. The enet_cityscapes_label.py specifies the conversion of training IDs to original cityscapes dataset IDs for this model.

Scores on the validation set

The author's converted model's scores on cityscapes validation set are as follows

Class IoU Class iIoU Category IoU Category iIoU
60.4 36.2 80.0 63.9

Implementation notes

The max-unpooling modules unpool the last pooling module (FIFO), which is appropriate for encoder-decoder segmentation networks. This converter is not tested for other networks than ENet (but should remain compatible with the models tested with convert_torch_to_pytorch).

References

License

The converter code developed in this repository follows the MIT license. ENet's model and weights (provided externally to this repository) are part of ENet and follow the CC BY-NC 4.0 license.

About

An ENet implementation in PyTorch, based on convert_torch_to_pytorch

Resources

License

Releases

No releases published

Packages

No packages published
You can’t perform that action at this time.