changed to absolute imports and added docs (#881)

This commit is contained in:
Akshay Kulkarni
2020-02-17 11:05:59 -05:00
committed by GitHub
parent f44dfb3e7a
commit 0ad3e8b8e9
2 changed files with 11 additions and 12 deletions
@@ -1,4 +0,0 @@
# For relative imports to work in Python 3.6
import os
import sys
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
@@ -2,30 +2,33 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from parts import DoubleConv, Down, Up
from models.unet.parts import DoubleConv, Down, Up
class UNet(nn.Module):
'''
Architecture based on U-Net: Convolutional Networks for Biomedical Image Segmentation
Link - https://arxiv.org/abs/1505.04597
Parameters:
num_classes (int) - Number of output classes required (default 19 for KITTI dataset)
bilinear (bool) - Whether to use bilinear interpolation or transposed
convolutions for upsampling.
'''
def __init__(self, num_classes=19, bilinear=False):
super().__init__()
self.bilinear = bilinear
self.num_classes = num_classes
self.layer1 = DoubleConv(3, 64)
self.layer2 = Down(64, 128)
self.layer3 = Down(128, 256)
self.layer4 = Down(256, 512)
self.layer5 = Down(512, 1024)
self.layer6 = Up(1024, 512, bilinear=self.bilinear)
self.layer7 = Up(512, 256, bilinear=self.bilinear)
self.layer8 = Up(256, 128, bilinear=self.bilinear)
self.layer9 = Up(128, 64, bilinear=self.bilinear)
self.layer6 = Up(1024, 512, bilinear=bilinear)
self.layer7 = Up(512, 256, bilinear=bilinear)
self.layer8 = Up(256, 128, bilinear=bilinear)
self.layer9 = Up(128, 64, bilinear=bilinear)
self.layer10 = nn.Conv2d(64, self.num_classes, kernel_size=1)
self.layer10 = nn.Conv2d(64, num_classes, kernel_size=1)
def forward(self, x):
x1 = self.layer1(x)