diff --git a/trail_detection_node/README.md b/trail_detection_node/README.md new file mode 100644 index 00000000..f6e7b881 --- /dev/null +++ b/trail_detection_node/README.md @@ -0,0 +1,29 @@ +SETUP GUIDE +============= +1. Create a new folder named ```model``` in trail_detection_node/trail_detection_node +2. Download ```psp_resnet50_pascal_voc_best_model.pth``` in NCRN Google Drive (Trail detection model folder) and put the file in the ```model``` folder +3. Download ```resnet50-25c4b509.pth``` in NCRN Google Drive (Trail detection model folder) and put the file in folder ```~/.torch/models``` +4. Build the package + +RUN the node +============= +``` +ros2 run trail_detection_node trail_detection +``` + +Node Usage +============= +The node has three modes, which can be controlled by setting the ```only_camera_mode```(line 175) and ```visualize```(line 188) to different value: +1. Trail centerline publisher(default, ```only_camera_mode=False```, ```visualize=False```): subscribes to camera and lidar and publishes the centerline point(posestamp type msg) in ```velodyne``` frame in ```trail_location``` topic +2. Trail centerline visualizer(```only_camera_mode=False```, ```visualize=True```): subscribes to camera and lidar and shows the chosen centerline point in **white**, the lidar points in **blue**, center points with no lidar points around in **yellow**, and center points with lidar points around in **red** in the image in the pop up window +3. Only Camera mode(```only_camera_mode=True```): subscribes only to camera, runs the segmentation model, and shows the segmented trail in image in the pop up window + +Possible issues +============= +1. **Model doesn't load successfully:** + Change the directory in the ```load model``` function in both ```trail_detection_node/trail_detection_node/v2_trailDetectionNode.py```(line 103 and 104) and ```trail_detection_node/trail_detection_node/visualizer_v2_trailDetectionNode.py```(line 107 and 110) +2. **Subfolder package relative import error:** + Temporary fix: Move the script to the same folder and change the import code + + + diff --git a/trail_detection_node/package.xml b/trail_detection_node/package.xml new file mode 100644 index 00000000..aa18aeb5 --- /dev/null +++ b/trail_detection_node/package.xml @@ -0,0 +1,21 @@ + + + + trail_detection_node + 0.0.0 + TODO: Package description + zhaodong + TODO: License declaration + + sensor_msgs + cv_bridge + + ament_copyright + ament_flake8 + ament_pep257 + python3-pytest + + + ament_python + + diff --git a/trail_detection_node/resource/trail_detection_node b/trail_detection_node/resource/trail_detection_node new file mode 100644 index 00000000..e69de29b diff --git a/trail_detection_node/setup.cfg b/trail_detection_node/setup.cfg new file mode 100644 index 00000000..3cedd943 --- /dev/null +++ b/trail_detection_node/setup.cfg @@ -0,0 +1,4 @@ +[develop] +script_dir=$base/lib/trail_detection_node +[install] +install_scripts=$base/lib/trail_detection_node diff --git a/trail_detection_node/setup.py b/trail_detection_node/setup.py new file mode 100644 index 00000000..83dd74cc --- /dev/null +++ b/trail_detection_node/setup.py @@ -0,0 +1,26 @@ +from setuptools import setup + +package_name = 'trail_detection_node' + +setup( + name=package_name, + version='0.0.0', + packages=[package_name], + data_files=[ + ('share/ament_index/resource_index/packages', + ['resource/' + package_name]), + ('share/' + package_name, ['package.xml']), + ], + install_requires=['setuptools'], + zip_safe=True, + maintainer='zhaodong', + maintainer_email='zhaodong.jiang@mail.utoronto.ca', + description='TODO: Package description', + license='TODO: License declaration', + tests_require=['pytest'], + entry_points={ + 'console_scripts': [ + 'trail_detection = trail_detection_node.trailDetectionNode:main', + ], + }, +) diff --git a/trail_detection_node/test.py b/trail_detection_node/test.py new file mode 100644 index 00000000..c925cc4c --- /dev/null +++ b/trail_detection_node/test.py @@ -0,0 +1,3 @@ +from trail_detection_node.trailDetectionNode import * + +main() diff --git a/trail_detection_node/test/test_copyright.py b/trail_detection_node/test/test_copyright.py new file mode 100644 index 00000000..97a39196 --- /dev/null +++ b/trail_detection_node/test/test_copyright.py @@ -0,0 +1,25 @@ +# Copyright 2015 Open Source Robotics Foundation, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ament_copyright.main import main +import pytest + + +# Remove the `skip` decorator once the source file(s) have a copyright header +@pytest.mark.skip(reason='No copyright header has been placed in the generated source file.') +@pytest.mark.copyright +@pytest.mark.linter +def test_copyright(): + rc = main(argv=['.', 'test']) + assert rc == 0, 'Found errors' diff --git a/trail_detection_node/test/test_flake8.py b/trail_detection_node/test/test_flake8.py new file mode 100644 index 00000000..27ee1078 --- /dev/null +++ b/trail_detection_node/test/test_flake8.py @@ -0,0 +1,25 @@ +# Copyright 2017 Open Source Robotics Foundation, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ament_flake8.main import main_with_errors +import pytest + + +@pytest.mark.flake8 +@pytest.mark.linter +def test_flake8(): + rc, errors = main_with_errors(argv=[]) + assert rc == 0, \ + 'Found %d code style errors / warnings:\n' % len(errors) + \ + '\n'.join(errors) diff --git a/trail_detection_node/test/test_pep257.py b/trail_detection_node/test/test_pep257.py new file mode 100644 index 00000000..b234a384 --- /dev/null +++ b/trail_detection_node/test/test_pep257.py @@ -0,0 +1,23 @@ +# Copyright 2015 Open Source Robotics Foundation, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ament_pep257.main import main +import pytest + + +@pytest.mark.linter +@pytest.mark.pep257 +def test_pep257(): + rc = main(argv=['.', 'test']) + assert rc == 0, 'Found code style errors / warnings' diff --git a/trail_detection_node/trail_detection_node/__init__.py b/trail_detection_node/trail_detection_node/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trail_detection_node/trail_detection_node/model_loader.py b/trail_detection_node/trail_detection_node/model_loader.py new file mode 100644 index 00000000..d8934ae3 --- /dev/null +++ b/trail_detection_node/trail_detection_node/model_loader.py @@ -0,0 +1,413 @@ +import torch.nn as nn +import torch.nn.functional as F + +import torch +import os + +# all the code below are copied from https://github.com/Tramac/awesome-semantic-segmentation-pytorch + +class BasicBlockV1b(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BasicBlockV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(True) + self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation, + dilation=previous_dilation, bias=False) + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class BottleneckV1b(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BottleneckV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = norm_layer(planes) + self.conv2 = nn.Conv2d(planes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn2 = norm_layer(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + +class ResNetV1b(nn.Module): + + def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False, + zero_init_residual=False, norm_layer=nn.BatchNorm2d): + self.inplanes = 128 if deep_stem else 64 + super(ResNetV1b, self).__init__() + if deep_stem: + self.conv1 = nn.Sequential( + nn.Conv2d(3, 64, 3, 2, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 64, 3, 1, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 128, 3, 1, 1, bias=False) + ) + else: + self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(True) + self.maxpool = nn.MaxPool2d(3, 2, 1) + self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) + if dilated: + self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer) + else: + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, BottleneckV1b): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlockV1b): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False), + norm_layer(planes * block.expansion), + ) + + layers = [] + if dilation in (1, 2): + layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + elif dilation == 4: + layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + else: + raise RuntimeError("=> unknown dilation size: {}".format(dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append(block(self.inplanes, planes, dilation=dilation, + previous_dilation=dilation, norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + +_model_sha1 = {name: checksum for checksum, name in [ + ('25c4b50959ef024fcc050213a06b614899f94b3d', 'resnet50'), + ('2a57e44de9c853fa015b172309a1ee7e2d0e4e2a', 'resnet101'), + ('0d43d698c66aceaa2bc0309f55efdd7ff4b143af', 'resnet152'), +]} + +def short_hash(name): + if name not in _model_sha1: + raise ValueError('Pretrained model for {name} is not available.'.format(name=name)) + return _model_sha1[name][:8] + +def get_resnet_file(name, root='~/.torch/models'): + file_name = '{name}-{short_hash}'.format(name=name, short_hash=short_hash(name)) + root = os.path.expanduser(root) + + file_path = os.path.join(root, file_name + '.pth') + sha1_hash = _model_sha1[name] + if os.path.exists(file_path): + return file_path + + else: + print('Model file {} is not found.'.format(file_path)) + +def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs) + if pretrained: + model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False) + return model + +class SeparableConv2d(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, + dilation=1, bias=False, norm_layer=nn.BatchNorm2d): + super(SeparableConv2d, self).__init__() + self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias) + self.bn = norm_layer(inplanes) + self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.pointwise(x) + return x + +# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py +class JPU(nn.Module): + def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs): + super(JPU, self).__init__() + + self.conv5 = nn.Sequential( + nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv4 = nn.Sequential( + nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv3 = nn.Sequential( + nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + + self.dilation1 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation2 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation3 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation4 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False), + norm_layer(width), + nn.ReLU(True)) + + def forward(self, *inputs): + feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])] + size = feats[-1].size()[2:] + feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True) + feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True) + feat = torch.cat(feats, dim=1) + feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)], + dim=1) + + return inputs[0], inputs[1], inputs[2], feat + + + +class SegBaseModel(nn.Module): + r"""Base Model for Semantic Segmentation + + Parameters + ---------- + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + """ + + def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs): + super(SegBaseModel, self).__init__() + dilated = False if jpu else True + self.aux = aux + self.nclass = nclass + if backbone == 'resnet50': + self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + + self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None + + def base_forward(self, x): + """forwarding pre-trained network""" + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + x = self.pretrained.maxpool(x) + c1 = self.pretrained.layer1(x) + c2 = self.pretrained.layer2(c1) + c3 = self.pretrained.layer3(c2) + c4 = self.pretrained.layer4(c3) + + if self.jpu: + return self.jpu(c1, c2, c3, c4) + else: + return c1, c2, c3, c4 + + def evaluate(self, x): + """evaluating network with inputs and targets""" + return self.forward(x)[0] + + def demo(self, x): + pred = self.forward(x) + if self.aux: + pred = pred[0] + return pred + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels), + nn.ReLU(inplace=True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + +''' +PSP NET +''' + +def _PSP1x1Conv(in_channels, out_channels, norm_layer, norm_kwargs): + return nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + +class _PyramidPooling(nn.Module): + def __init__(self, in_channels, **kwargs): + super(_PyramidPooling, self).__init__() + out_channels = int(in_channels / 4) + self.avgpool1 = nn.AdaptiveAvgPool2d(1) + self.avgpool2 = nn.AdaptiveAvgPool2d(2) + self.avgpool3 = nn.AdaptiveAvgPool2d(3) + self.avgpool4 = nn.AdaptiveAvgPool2d(6) + self.conv1 = _PSP1x1Conv(in_channels, out_channels, **kwargs) + self.conv2 = _PSP1x1Conv(in_channels, out_channels, **kwargs) + self.conv3 = _PSP1x1Conv(in_channels, out_channels, **kwargs) + self.conv4 = _PSP1x1Conv(in_channels, out_channels, **kwargs) + + def forward(self, x): + size = x.size()[2:] + feat1 = F.interpolate(self.conv1(self.avgpool1(x)), size, mode='bilinear', align_corners=True) + feat2 = F.interpolate(self.conv2(self.avgpool2(x)), size, mode='bilinear', align_corners=True) + feat3 = F.interpolate(self.conv3(self.avgpool3(x)), size, mode='bilinear', align_corners=True) + feat4 = F.interpolate(self.conv4(self.avgpool4(x)), size, mode='bilinear', align_corners=True) + return torch.cat([x, feat1, feat2, feat3, feat4], dim=1) + +class _PSPHead(nn.Module): + def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_PSPHead, self).__init__() + self.psp = _PyramidPooling(2048, norm_layer=norm_layer, norm_kwargs=norm_kwargs) + self.block = nn.Sequential( + nn.Conv2d(4096, 512, 3, padding=1, bias=False), + norm_layer(512, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(512, nclass, 1) + ) + + def forward(self, x): + x = self.psp(x) + return self.block(x) + +class PSPNet(SegBaseModel): + r"""Pyramid Scene Parsing Network + + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Zhao, Hengshuang, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia. + "Pyramid scene parsing network." *CVPR*, 2017 + """ + + def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs): + super(PSPNet, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs) + self.head = _PSPHead(nclass, **kwargs) + if self.aux: + self.auxlayer = _FCNHead(1024, nclass, **kwargs) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + _, _, c3, c4 = self.base_forward(x) + outputs = [] + x = self.head(c4) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + return tuple(outputs) \ No newline at end of file diff --git a/trail_detection_node/trail_detection_node/setup.py b/trail_detection_node/trail_detection_node/setup.py new file mode 100644 index 00000000..ff6a39fd --- /dev/null +++ b/trail_detection_node/trail_detection_node/setup.py @@ -0,0 +1,125 @@ +from setuptools import find_packages, setup + + +def readme(): + with open('README.md', encoding='utf-8') as f: + content = f.read() + return content + + +version_file = 'mmseg/version.py' + + +def get_version(): + with open(version_file, 'r') as f: + exec(compile(f.read(), version_file, 'exec')) + return locals()['__version__'] + + +def parse_requirements(fname='requirements.txt', with_version=True): + """Parse the package dependencies listed in a requirements file but strips + specific versioning information. + + Args: + fname (str): path to requirements file + with_version (bool, default=False): if True include version specs + + Returns: + List[str]: list of requirements items + + CommandLine: + python -c "import setup; print(setup.parse_requirements())" + """ + import sys + from os.path import exists + import re + require_fpath = fname + + def parse_line(line): + """Parse information from a line in a requirements text file.""" + if line.startswith('-r '): + # Allow specifying requirements in other files + target = line.split(' ')[1] + for info in parse_require_file(target): + yield info + else: + info = {'line': line} + if line.startswith('-e '): + info['package'] = line.split('#egg=')[1] + else: + # Remove versioning from the package + pat = '(' + '|'.join(['>=', '==', '>']) + ')' + parts = re.split(pat, line, maxsplit=1) + parts = [p.strip() for p in parts] + + info['package'] = parts[0] + if len(parts) > 1: + op, rest = parts[1:] + if ';' in rest: + # Handle platform specific dependencies + # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies + version, platform_deps = map(str.strip, + rest.split(';')) + info['platform_deps'] = platform_deps + else: + version = rest # NOQA + info['version'] = (op, version) + yield info + + def parse_require_file(fpath): + with open(fpath, 'r') as f: + for line in f.readlines(): + line = line.strip() + if line and not line.startswith('#'): + for info in parse_line(line): + yield info + + def gen_packages_items(): + if exists(require_fpath): + for info in parse_require_file(require_fpath): + parts = [info['package']] + if with_version and 'version' in info: + parts.extend(info['version']) + if not sys.version.startswith('3.4'): + # apparently package_deps are broken in 3.4 + platform_deps = info.get('platform_deps') + if platform_deps is not None: + parts.append(';' + platform_deps) + item = ''.join(parts) + yield item + + packages = list(gen_packages_items()) + return packages + + +if __name__ == '__main__': + setup( + name='mmsegmentation', + version=get_version(), + description='Open MMLab Semantic Segmentation Toolbox and Benchmark', + long_description_content_type='text/markdown', + author='MMSegmentation Authors', + author_email='openmmlab@gmail.com', + keywords='computer vision, semantic segmentation', + url='http://github.com/open-mmlab/mmsegmentation', + packages=find_packages(exclude=('configs', 'tools', 'demo')), + classifiers=[ + 'Development Status :: 4 - Beta', + 'License :: OSI Approved :: Apache Software License', + 'Operating System :: OS Independent', + 'Programming Language :: Python :: 3.6', + 'Programming Language :: Python :: 3.7', + 'Programming Language :: Python :: 3.8', + ], + license='Apache License 2.0', + setup_requires=parse_requirements('requirements/build.txt'), + tests_require=parse_requirements('requirements/tests.txt'), + install_requires=parse_requirements('requirements/runtime.txt'), + extras_require={ + 'all': parse_requirements('requirements.txt'), + 'tests': parse_requirements('requirements/tests.txt'), + 'build': parse_requirements('requirements/build.txt'), + 'optional': parse_requirements('requirements/optional.txt'), + }, + ext_modules=[], + zip_safe=False) diff --git a/trail_detection_node/trail_detection_node/trailDetectionNode.py b/trail_detection_node/trail_detection_node/trailDetectionNode.py new file mode 100644 index 00000000..e4eef73e --- /dev/null +++ b/trail_detection_node/trail_detection_node/trailDetectionNode.py @@ -0,0 +1,420 @@ +import torch +from .model_loader import PSPNet +from PIL import Image as ImagePIL +from torchvision import transforms +import cv2 +import os + +import rclpy +from rclpy.node import Node +from geometry_msgs.msg import PoseStamped +from sensor_msgs.msg import Image, PointCloud2 +import sensor_msgs_py.point_cloud2 as pc2 +import numpy as np +import math +from cv_bridge import CvBridge + +import message_filters +''' + The transformation matrix as well as the coordinate conversion and depth estimation functions (from line 21 to line 98) are copied from human_detection_node +''' +camera_transformation_k = np.array([ + [628.5359544,0,676.9575694], + [0,627.7249542,532.7206716], + [0,0,1]]) + +rotation_matrix = np.array([ + [-0.007495781893,-0.0006277316155,0.9999717092], + [-0.9999516401,-0.006361853422,-0.007499625104], + [0.006366381192,-0.9999795662,-0.0005800141927]]) + +rotation_matrix = rotation_matrix.T + +translation_vector = np.array([-0.06024059837, -0.08180891509, -0.3117851288]) +image_width=1280 +image_height=1024 + +def convert_to_lidar_frame(uv_coordinate): + """ + convert 2d camera coordinate + depth into 3d lidar frame + """ + point_cloud = np.empty( (3,) , dtype=float) + point_cloud[2] = uv_coordinate[2] + point_cloud[1] = ( image_height - uv_coordinate[1] )*point_cloud[2] + point_cloud[0] = uv_coordinate[0]*point_cloud[2] + + inverse_camera_transformation_k = np.linalg.inv(camera_transformation_k) + inverse_rotation_matrix = np.linalg.inv(rotation_matrix) + point_cloud = inverse_camera_transformation_k @ point_cloud + point_cloud = inverse_rotation_matrix @ (point_cloud-translation_vector) + return point_cloud + +def convert_to_camera_frame(point_cloud): + """ + convert 3d lidar data into 2d coordinate of the camera frame + depth + """ + length = point_cloud.shape[0] + translation = np.tile(translation_vector, (length, 1)).T + + point_cloud = point_cloud.T + point_cloud = rotation_matrix@point_cloud + translation + point_cloud = camera_transformation_k @ point_cloud + + uv_coordinate = np.empty_like(point_cloud) + + """ + uv = [x/z, y/z, z], and y is opposite so the minus imageheight + """ + uv_coordinate[0] = point_cloud[0] / point_cloud[2] + uv_coordinate[1] = image_height - point_cloud[1] / point_cloud[2] + uv_coordinate[2] = point_cloud[2] + + uv_depth = uv_coordinate[2, :] + filtered_uv_coordinate = uv_coordinate[:, uv_depth >= 0] + return filtered_uv_coordinate + +def estimate_depth(x, y, np_2d_array): + """ + estimate the depth by finding points closest to x,y from thhe 2d array + """ + # Calculate the distance between each point and the target coordinates (x, y) + distances_sq = (np_2d_array[0,:] - x) ** 2 + (np_2d_array[1,:] - y) ** 2 + + # Find the indices of the k nearest points + k = 5 # Number of nearest neighbors we want + closest_indices = np.argpartition(distances_sq, k)[:k] + pixel_distance_threshold = 2000 + + valid_indices = [idx for idx in closest_indices if distances_sq[idx]<=pixel_distance_threshold] + if len(valid_indices) == 0: + # lidar points disappears usually around 0.4m + distance_where_lidar_stops_working = -1 + return distance_where_lidar_stops_working + + filtered_indices = np.array(valid_indices) + # Get the depth value of the closest point + closest_depths = np_2d_array[2,filtered_indices] + + return np.mean(closest_depths) + +def load_model(device): + # model = FCN32s(nclass=2, backbone='vgg16', pretrained_base=True, pretrained=True) + model = PSPNet(nclass=2, backbone='resnet50', pretrained_base=True) + + # model_location = '32s_batch8/fcn32s_vgg16_pascal_voc_best_model.pth' + # model_location = '32s_self_labelled/fcn32s_vgg16_pascal_voc_best_model.pth' + model_location = 'psp2/psp_resnet50_pascal_voc_best_model.pth' + + if os.path.isfile(f'src/trail_detection_node/trail_detection_node/model/{model_location}'): + model.load_state_dict(torch.load(f'src/trail_detection_node/trail_detection_node/model/{model_location}',map_location=torch.device('cuda:0'))) + else: + model.load_state_dict(torch.load(f'trail_detection_node/model/{model_location}',map_location=torch.device('cuda:0'))) + model = model.to(device) + + model.eval() + print('Finished loading model!') + + return model + +def find_route(model, device, cv_image): + PIL_image = ImagePIL.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)) + transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ]) + image = transform(PIL_image).unsqueeze(0).to(device) + with torch.no_grad(): + output = model(image) + # pred is prediction generated by the model, route is the variable for the center line + pred = torch.argmax(output[0], 1).squeeze(0).cpu().data.numpy() + pred[pred == 0] = 255 # only see the trail + pred[pred == 1] = 0 + pred = np.array(pred, dtype=np.uint8) + + # # prediction visualize + # cv2.imshow('prediction',pred) + # cv2.waitKey(25) + + route = np.zeros_like(pred) + # calculate the center line by taking the average + row_num = 0 + for row in pred: + white_pixels = list(np.nonzero(row)[0]) + if white_pixels: + average = (white_pixels[0] + white_pixels[-1]) / 2 + route[row_num][round(average)] = 255 + row_num = row_num + 1 + return route, pred + +def equalize_hist_rgb(img): + # Split the image into its color channels + r, g, b = cv2.split(img) + + # Equalize the histograms for each channel + r_eq = cv2.equalizeHist(r) + g_eq = cv2.equalizeHist(g) + b_eq = cv2.equalizeHist(b) + + # Merge the channels + img_eq = cv2.merge((r_eq, g_eq, b_eq)) + # img_eq = cv2.fastNlMeansDenoisingColored(img_eq,None,10,20,7,42) + # img_eq = cv2.GaussianBlur(img_eq, (25,25), 0) + return img_eq + +''' + The traildetector node has two subscriptions(lidar and camera) and one publisher(trail position). After it receives msgs from both lidar and camera, + it detects the trail in the image and sends the corresponding lidar position as the trail location. + The msgs are synchornized before processing, using buffer and sync function. + To find the path, the node will process the image, find a line to follow (by taking the average of left and right of the path), estimate the lidar points depth, + and choose to go to the closest point. + This ediction of node assumes that the path is pointing frontward and has only one path in front. +''' +class trailDetector(Node): + def __init__(self, model, device): + super().__init__('trail_detector') + self.only_camera_mode = False + self.bridge = CvBridge() + + # load model and device + self.model = model + self.device = device + + # define trail publication + self.trail_publisher = self.create_publisher( + PoseStamped, + 'trail_location', + 10) + + self.visualize = False + + if self.only_camera_mode: + # define camera subscription + self.get_logger().info("Mode: Only Camera") + self.camera_subscription = self.create_subscription( + Image, + 'camera', + self.only_camera_callback, + 10) + self.camera_subscription + else: + # create subscribers and synchronizer + self.get_logger().info("Mode: Lidar and Camera") + self.image_sub = message_filters.Subscriber(self, Image, 'camera') + self.lidar_sub = message_filters.Subscriber(self, PointCloud2, 'velodyne_points') + queue_size = 30 + ts = message_filters.ApproximateTimeSynchronizer([self.image_sub, self.lidar_sub], queue_size, 0.5) + ts.registerCallback(self.trail_callback) + + + + + def trail_callback(self, camera_msg, lidar_msg): + # print("Message received!") + # process lidar msg + point_gen = pc2.read_points( + lidar_msg, field_names=( + "x", "y", "z"), skip_nans=True) + points = [[x, y, z] for x, y, z in point_gen] + points = np.array(points) + points2d = convert_to_camera_frame(points) + + # process camera msg + cv_image = self.bridge.imgmsg_to_cv2(camera_msg, desired_encoding='passthrough') + + # # increase brightness + # M = np.ones(cv_image.shape, dtype = 'uint8') * 50 # increase brightness by 50 + # cv_image = cv2.add(cv_image, M) + + # # image visualizer + # cv2.imshow('raw image',cv_image) + # cv2.waitKey(25) + + route, _ = find_route(self.model, self.device, cv_image) + + # # visualize route purely + # cv2.imshow("white_route",route) + # cv2.waitKey(25) + + route_indices = list(zip(*np.nonzero(route))) + if not route_indices: + self.get_logger().info("No centerline found!") + return + + #filter points that have no lidar points near it + filtered_route_indices = [] + for index in route_indices: + point = [] + u = index[1] + v = image_height - index[0] + point.append(u) + point.append(v) + point.append(estimate_depth(u, v, points2d)) + if point[2] == -1: + continue + else: + filtered_route_indices.append(point) + + if not filtered_route_indices: + self.get_logger().info("No usable centerline found!") + return + + # find the corresponding lidar points using the center line pixels + filtered_3dPoints = [] + for index in filtered_route_indices: + # for index in route_indices: + point = [] + point.append(index[0]) + point.append(index[1]) + point.append(index[2]) + + # point.append(index[0]) + # point.append(index[1]) + # point.append(estimate_depth(index[0], index[1], points2d)) + + point_3d = convert_to_lidar_frame(point) + filtered_3dPoints.append(point_3d) + + filtered_3dPoints = np.array(filtered_3dPoints) + # find the nearest 3d point and set that as goal + distances_sq = filtered_3dPoints[:,0]**2 + filtered_3dPoints[:,1]**2 + filtered_3dPoints[:,2]**2 + self.smallest_index = np.argmin(distances_sq) + # print(math.sqrt(distances_sq[smallest_index])) + # smallest_index = np.argmin(filtered_3dPoints[:,2]) + + if self.visualize == True: + self.visualization(route_indices, points2d, filtered_route_indices, cv_image) + else: + self.publish_msg(lidar_msg, filtered_3dPoints) + + + def visualization(self, route_indices, points2d, filtered_route_indices, cv_image): + # visualize after-pocessed image + visualize_cv_image = cv_image + # print(f"{visualize_cv_image.shape[0]}") + circle_y = route_indices[self.smallest_index][0] + circle_x = route_indices[self.smallest_index][1] + + # visualize the lidar points in image + uv_x, uv_y, uv_z = points2d + for index_lidar in range(points2d.shape[1]): + # print(f"{int(uv_x[index_lidar])},{int(uv_y[index_lidar])}") + cv2.circle(visualize_cv_image, (int(uv_x[index_lidar]), int(image_height - uv_y[index_lidar])), radius=5, color=(255, 0, 0), thickness=-1) + + # visualize center line in image + for index_circle in range(len(route_indices)): + if index_circle == self.smallest_index: + continue + else: + green_circle_row = route_indices[index_circle][0] + green_circle_coloum = route_indices[index_circle][1] + cv2.circle(visualize_cv_image, (green_circle_coloum, green_circle_row), radius=5, color=(0, 125, 125), thickness=-1) + + # visualize filtered center line in image + for index_circle in range(len(filtered_route_indices)): + if index_circle == self.smallest_index: + continue + else: + red_circle_row = filtered_route_indices[index_circle][1] + red_circle_coloum = filtered_route_indices[index_circle][0] + cv2.circle(visualize_cv_image, (red_circle_coloum, image_height - red_circle_row), radius=5, color=(0, 0, 255), thickness=-1) + + # visualize the chosen point in image + cv2.circle(visualize_cv_image, (circle_x, circle_y), radius=7, color=(255, 255, 255), thickness=-1) + + cv2.imshow('circled image',visualize_cv_image) + cv2.waitKey(25) + + def publish_msg(self, lidar_msg, filtered_3dPoints): + # publsih message + trail_location_msg = PoseStamped() + trail_location_msg.header.stamp = lidar_msg.header.stamp + trail_location_msg.header.frame_id = "velodyne" + #position + trail_location_msg.pose.position.x = filtered_3dPoints[self.smallest_index][0] + trail_location_msg.pose.position.y = filtered_3dPoints[self.smallest_index][1] + trail_location_msg.pose.position.z = filtered_3dPoints[self.smallest_index][2] + #orientation + yaw = math.atan2(filtered_3dPoints[self.smallest_index][1], filtered_3dPoints[self.smallest_index][0]) + trail_location_msg.pose.orientation.x = 0.0 + trail_location_msg.pose.orientation.y = 0.0 + trail_location_msg.pose.orientation.z = math.sin(yaw/2) + trail_location_msg.pose.orientation.w = math.cos(yaw / 2) + self.get_logger().info("location published!") + self.get_logger().info(f"Point: {filtered_3dPoints[self.smallest_index][0]}, {filtered_3dPoints[self.smallest_index][1]}, {filtered_3dPoints[self.smallest_index][2]}") + self.trail_publisher.publish(trail_location_msg) + + def only_camera_callback(self, camera_msg): + # process camera msg + cv_image = self.bridge.imgmsg_to_cv2(camera_msg, desired_encoding='passthrough') + # cv_image = cv2.resize(cv_image,(480,480)) + + # cv_image = cv2.fastNlMeansDenoisingColored(cv_image,None,5,5,5,11) + + # # histogram equalization method 1 + # cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2YCrCb) + # cv_image[:,:,0] = cv2.equalizeHist(cv_image[:,:,0]) + # cv_image = cv2.cvtColor(cv_image, cv2.COLOR_YCrCb2BGR) + + # histogram equalization method 2 + # cv_image = equalize_hist_rgb(cv_image) + + # # histogram equalization method 3 CLAHE + # image_bw = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY) + # clahe = cv2.createCLAHE(clipLimit=5) + # cv_image = clahe.apply(image_bw) + 30 + + # # histogram equalization method 4 CLAHE + # cla = cv2.createCLAHE(clipLimit=4.0) + # H, S, V = cv2.split(cv2.cvtColor(cv_image, cv2.COLOR_BGR2HSV)) + # eq_V = cla.apply(V) + # cv_image = cv2.cvtColor(cv2.merge([H, S, eq_V]), cv2.COLOR_HSV2BGR) + + # # increase brightness + # M = np.ones(cv_image.shape, dtype = 'uint8') * 50 # increase brightness by 50 + # cv_image = cv2.add(cv_image, M) + + # # image visualizer + # cv2.imshow('raw image',cv_image) + # cv2.waitKey(25) + + route, pred = find_route(self.model, self.device, cv_image) + + # # visualize route purely + # cv2.imshow("white_route",route) + # cv2.waitKey(25) + + # increase the brightness of image where is predicted as road + sign = cv2.cvtColor(pred, cv2.COLOR_GRAY2BGR) + cv_image = cv2.add(cv_image, sign) + + route_indices = list(zip(*np.nonzero(route))) + if not route_indices: + self.get_logger().info("No centerline found!") + return + + # for index_circle in range(len(route_indices)): + # red_circle_x = route_indices[index_circle][0] + # red_circle_y = route_indices[index_circle][1] + # cv2.circle(cv_image, (red_circle_y, red_circle_x), radius=5, color=(0, 0, 255), thickness=-1) + + cv2.imshow('circled image',cv_image) + cv2.waitKey(25) + + + + +def main(args=None): + print(torch.cuda.is_available()) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = load_model(device) + + rclpy.init(args=args) + trailDetectorNode = trailDetector(model, device) + rclpy.spin(trailDetectorNode) + + trailDetectorNode.destroy_node() + rclpy.shutdown() + cv2.destroyAllWindows() + +if __name__ == '__main__': + main()