Implemented Spokes Ellipse Fitting:

- added new metrics to inference script to evaluate detection as well as segmentation - added commands readme with instructions on running the code - added requirement files for conda and pip
biorobotics · Dec 14, 2022 · 1838747 · 1838747
1 parent 70035a5
commit 1838747
Show file tree

Hide file tree

Showing 13 changed files with 853 additions and 155 deletions.
diff --git a/.gitignore b/.gitignore
@@ -25,8 +25,6 @@
 
 storage.googleapis.com
 runs/*
-runs_old/*
-data/*
 !data/hyps/*
 !data/images/zidane.jpg
 !data/images/bus.jpg

diff --git a/commands.md b/commands.md
@@ -0,0 +1,151 @@
+## Commands to run and evaluate YOLO v3-tiny models for vessel segmentation
+
+***
+
+### Table of Contents:
+1. [Requirements](#requirements)
+2. [Pre-Training](#repository-overview)
+3. [Training the YOLO v3-tiny](#requirements)
+4. [Evaluation](#building-the-repository)
+5. [Ellipse Fitting](#running-the-robot)
+
+***
+
+### Requirements
+
+Create a conda environment using the `conda_req.txt` file.
+
+```bash
+>>> conda create --name <env_name> --file conda_req.txt
+>>> conda activate <env_name>
+```
+
+If using pip instead of conda, install the requirements using the `pip_req.txt` file.
+
+```bash
+>>> pip install -r pip_req.txt
+```
+
+***
+
+### Pre-Training
+
+Create the `dataset` directory and add the `name_of_dataset` dataset to this directory. The dataset should follow the following folder structure and name convention.
+
+~~~{.bash}
+dataset
+└── name_of_dataset
+    ├── train
+    │   ├── images      # Images of vessels (256*256*3, .png images)
+    │   └── labels_bb   # Bounding box labels (txt)
+    ├── valid
+    │   ├── images      # Images of vessels (256*256*3, .png images)
+    │   └── labels_bb   # Bounding box labels (txt)
+    └── test
+        ├── images      # Images of vessels (256*256*3, .png images)
+        └── labels_bb   # Bounding box labels (txt)
+~~~
+
+In the `data` directory, create a YAML file coresponing your dataset, with the following contents.
+
+```yaml
+path: dataset/mixed_pig_fukuda_dataset1     # dataset root dir
+
+train: train/images                         # train images (relative to 'path')
+val: valid/images                           # val images (relative to 'path') 128 images
+test: test/images                           # test images (optional)
+
+# Classes
+nc: 2                                       # number of classes
+names: ['artery', 'vein']                   # name of classes
+```
+For vessel segmentation in ultrasound images, we have already provided the `data/vessel_multiclass.yaml` file.
+
+Create a `runs` directory to store the training results. The `runs` directory should have `train` and `val` subdirectories to store the training and validation results respectively.
+
+***
+
+### Training the YOLO v3-tiny
+
+To train the YOLO v3-tiny model, run the `train.py` script. The `train.py` script takes multiple arguments, but here are a few important ones.
+
+~~~{.bash}
+>>> python train.py -h
+usage: train.py [-h] [--weights WEIGHTS] [--data DATA][--epochs EPOCHS]
+                [--batch-size BATCH_SIZE] [--imgsz IMGSZ] [--name NAME]
+                [--single-cls] [--augment]
+optional arguments:
+  -h, --help            show this help message and exit
+  --weights WEIGHTS     initial weights path
+  --data DATA           dataset.yaml path
+  --epochs EPOCHS       number of epochs
+  --batch-size BATCH_SIZE
+                        total batch size for all GPUs, -1 for autobatch
+  --imgsz IMGSZ, --img IMGSZ, --img-size IMGSZ
+                        train, val image size (pixels)
+  --single-cls          train multi-class data as single-class
+  --name NAME           save to project/name
+  --augment             data augmentation needed
+~~~
+
+Here are a couple of examples on how to train the YOLO v3-tiny model. The first one shows single class training, and the second one shows multi-class training. The third one demonstrates how these can be run with with data augmentation.
+
+~~~{.bash}
+>>> python train.py --data vessel_multiclass.yaml --weights yolov3-tiny.pt --img 256 --epoch 300 --batch-size 16 --name "yolo_singleclass" --single-cls
+~~~
+~~~{.bash}
+>>> python train.py --data vessel_multiclass.yaml --weights yolov3-tiny.pt --img 256 --epoch 300 --batch-size 16 --name "yolo_multiclass"
+~~~
+~~~{.bash}
+>>> python train.py --data vessel_multiclass.yaml --weights yolov3-tiny.pt --img 256 --epoch 300 --batch-size 16 --name "yolo_multiclass_aug" --augment
+~~~
+
+***
+
+### Evaluation
+
+To evaluate the YOLO v3-tiny trained model, run the `infer.py` script. This script is hardcoded to perform inference on the `Pig_A`, `Pig_B`, `Pig_C` and `Pig_D` datasets. Download these datasets (.zip files) from the TRACIR shared Google Drive under the `datasets` folder. The `infer.py` script takes multiple arguments, but here are a few important ones.
+
+~~~{.bash}
+>>> python infer.py -h
+usage: infer.py [-h] [--data DATA] [--weights WEIGHTS [WEIGHTS ...]]
+                [--batch-size BATCH_SIZE] [--imgsz IMGSZ] [--augment]
+                [--name NAME] [--single-cls]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --data DATA           dataset.yaml path
+  --weights WEIGHTS [WEIGHTS ...]
+                        model.pt path(s)
+  --batch-size BATCH_SIZE
+                        batch size
+  --imgsz IMGSZ, --img IMGSZ, --img-size IMGSZ
+                        inference size (pixels)
+  --augment             augmented inference
+  --name NAME           save to project/name
+  --single-cls          treat as single-class dataset
+~~~
+
+Here are a couple of examples on how to run the evaluation script. The first one shows single class testing, and the second one shows multi-class testing. The third one demonstrates how these can be run with with data augmentation.
+
+~~~{.bash}
+>>> python infer.py --data vessel_multiclass.yaml --weights runs/train/yolo_singleclass/weights/best.pt --img 256 --name "test" --single-cls
+~~~
+~~~{.bash}
+>>> python infer.py --data vessel_multiclass.yaml --weights runs/train/yolo_multiclass/weights/best.pt --img 256 --name "test"
+~~~
+~~~{.bash}
+>>> python infer.py --data vessel_multiclass.yaml --weights runs/train/yolo_singleclass_aug/weights/best.pt --img 256 --name "test" --single-cls --augment
+~~~
+
+***
+
+### Ellipse Fitting
+
+After YOLO predicts the bounding boxes, we fit an ellipse to each bounding box using the Spokes Ellipse algorithm, to perform "segmentation" with YOLO. The ellipse fitting is already taken care of in the `infer.py` script.
+
+Ellipse fitting happens in 2 steps, `binary_search` and `fit_ellipse`. The `binary_search` function is the "Spokes" part of the algorithm. It assumes that the center of the bounding box is the center of the ellipse, and creates 8 spokes. Then the alogrithm travels along these spokes in the image to find the point where the intensity of the image changes (from dark to light, the boundary of the vessel), thus performing a binary search. The `binary_search` function returns the 8 points (list of x and y) that are the endpoints of the spokes.
+
+The `fit_ellipse` function is the "Ellipse" part of the algorithm. It takes the 8 end points of the spokes and fits an ellipse on it using the least squares method. The optimization code implementation is taken from this [link](http://juddzone.com/ALGORITHMS/least_squares_ellipse.html). The `fit_ellipse` function returns the center of the ellipse, the major and minor axes, and the angle of the ellipse.
+
+***
diff --git a/commands.txt b/commands.txt
diff --git a/conda_req.txt b/conda_req.txt
@@ -0,0 +1,183 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=main
+_openmp_mutex=5.1=1_gnu
+absl-py=1.3.0=pypi_0
+backcall=0.2.0=pypi_0
+beautifulsoup4=4.11.1=pypi_0
+blas=1.0=mkl
+bzip2=1.0.8=h7b6447c_0
+ca-certificates=2022.10.11=h06a4308_0
+cachetools=4.2.4=pypi_0
+certifi=2021.5.30=py36h06a4308_0
+charset-normalizer=2.0.12=pypi_0
+click=8.0.4=pypi_0
+cuda=11.7.1=0
+cuda-cccl=11.7.91=0
+cuda-command-line-tools=11.7.1=0
+cuda-compiler=11.7.1=0
+cuda-cudart=11.7.99=0
+cuda-cudart-dev=11.7.99=0
+cuda-cuobjdump=11.7.91=0
+cuda-cupti=11.7.101=0
+cuda-cuxxfilt=11.7.91=0
+cuda-demo-suite=11.8.86=0
+cuda-documentation=11.8.86=0
+cuda-driver-dev=11.7.99=0
+cuda-gdb=11.8.86=0
+cuda-libraries=11.7.1=0
+cuda-libraries-dev=11.7.1=0
+cuda-memcheck=11.8.86=0
+cuda-nsight=11.8.86=0
+cuda-nsight-compute=11.8.0=0
+cuda-nvcc=11.7.99=0
+cuda-nvdisasm=11.8.86=0
+cuda-nvml-dev=11.7.91=0
+cuda-nvprof=11.8.87=0
+cuda-nvprune=11.7.91=0
+cuda-nvrtc=11.7.99=0
+cuda-nvrtc-dev=11.7.99=0
+cuda-nvtx=11.7.91=0
+cuda-nvvp=11.8.87=0
+cuda-runtime=11.7.1=0
+cuda-sanitizer-api=11.8.86=0
+cuda-toolkit=11.7.1=0
+cuda-tools=11.7.1=0
+cuda-visual-tools=11.7.1=0
+cudatoolkit=11.3.1=ha36c431_9
+cycler=0.11.0=pypi_0
+dataclasses=0.8=pyh4f3eec9_6
+decorator=5.1.1=pypi_0
+docker-pycreds=0.4.0=pypi_0
+ffmpeg=4.3=hf484d3e_0
+freetype=2.11.0=h70c0345_0
+gds-tools=1.4.0.31=0
+gitdb=4.0.9=pypi_0
+gitpython=3.1.18=pypi_0
+gmp=6.2.1=h295c915_3
+gnutls=3.6.15=he1e5248_0
+google-auth=2.13.0=pypi_0
+google-auth-oauthlib=0.4.6=pypi_0
+grpcio=1.48.2=pypi_0
+idna=3.4=pypi_0
+importlib-metadata=4.8.3=pypi_0
+importlib-resources=5.4.0=pypi_0
+intel-openmp=2022.1.0=h9e868ea_3769
+ipdb=0.13.9=pypi_0
+ipython=7.16.3=pypi_0
+ipython-genutils=0.2.0=pypi_0
+jedi=0.17.2=pypi_0
+jpeg=9e=h7f8727e_0
+kiwisolver=1.3.1=pypi_0
+lame=3.100=h7b6447c_0
+lcms2=2.12=h3be6417_0
+lerc=3.0=h295c915_0
+libcublas=11.11.3.6=0
+libcublas-dev=11.11.3.6=0
+libcufft=10.9.0.58=0
+libcufft-dev=10.9.0.58=0
+libcufile=1.4.0.31=0
+libcufile-dev=1.4.0.31=0
+libcurand=10.3.0.86=0
+libcurand-dev=10.3.0.86=0
+libcusolver=11.4.1.48=0
+libcusolver-dev=11.4.1.48=0
+libcusparse=11.7.5.86=0
+libcusparse-dev=11.7.5.86=0
+libdeflate=1.8=h7f8727e_5
+libedit=3.1.20210910=h7f8727e_0
+libffi=3.2.1=hf484d3e_1007
+libgcc-ng=11.2.0=h1234567_1
+libgomp=11.2.0=h1234567_1
+libiconv=1.16=h7f8727e_2
+libidn2=2.3.2=h7f8727e_0
+libnpp=11.8.0.86=0
+libnpp-dev=11.8.0.86=0
+libnvjpeg=11.9.0.86=0
+libnvjpeg-dev=11.9.0.86=0
+libpng=1.6.37=hbc83047_0
+libstdcxx-ng=11.2.0=h1234567_1
+libtasn1=4.16.0=h27cfd23_0
+libtiff=4.4.0=hecacb30_0
+libunistring=0.9.10=h27cfd23_0
+libuv=1.40.0=h7b6447c_0
+libwebp-base=1.2.4=h5eee18b_0
+lxml=4.9.1=pypi_0
+lz4-c=1.9.3=h295c915_1
+markdown=3.3.7=pypi_0
+matplotlib=3.3.4=pypi_0
+mkl=2020.2=256
+mkl-service=2.3.0=py36he8ac12f_0
+mkl_fft=1.3.0=py36h54f3939_0
+mkl_random=1.1.1=py36h0573a6f_0
+natsort=8.2.0=pypi_0
+ncurses=6.3=h5eee18b_3
+nettle=3.7.3=hbbd107a_1
+nsight-compute=2022.3.0.22=0
+numpy=1.19.5=pypi_0
+numpy-base=1.19.2=py36hfa32c7d_0
+oauthlib=3.2.2=pypi_0
+olefile=0.46=py36_0
+opencv-python=4.6.0.66=pypi_0
+openh264=2.1.1=h4ff587b_0
+openjpeg=2.4.0=h3ad879b_0
+openssl=1.1.1s=h7f8727e_0
+pandas=1.1.5=pypi_0
+parso=0.7.1=pypi_0
+pathtools=0.1.2=pypi_0
+pexpect=4.8.0=pypi_0
+pickleshare=0.7.5=pypi_0
+pillow=8.4.0=pypi_0
+pip=21.2.2=py36h06a4308_0
+promise=2.3=pypi_0
+prompt-toolkit=3.0.31=pypi_0
+protobuf=3.19.6=pypi_0
+psutil=5.9.2=pypi_0
+ptyprocess=0.7.0=pypi_0
+pyasn1=0.4.8=pypi_0
+pyasn1-modules=0.2.8=pypi_0
+pygments=2.13.0=pypi_0
+pyparsing=3.0.9=pypi_0
+python=3.6.9=h265db76_0
+python-dateutil=2.8.2=pypi_0
+pytorch=1.10.2=py3.6_cuda11.3_cudnn8.2.0_0
+pytorch-cuda=11.7=h67b0de4_0
+pytorch-mutex=1.0=cuda
+pytz=2022.4=pypi_0
+pyyaml=6.0=pypi_0
+readline=7.0=h7b6447c_5
+requests=2.27.1=pypi_0
+requests-oauthlib=1.3.1=pypi_0
+rsa=4.9=pypi_0
+scipy=1.5.4=pypi_0
+seaborn=0.11.2=pypi_0
+sentry-sdk=1.9.10=pypi_0
+setproctitle=1.2.3=pypi_0
+setuptools=58.0.4=py36h06a4308_0
+shortuuid=1.0.9=pypi_0
+six=1.16.0=pyhd3eb1b0_1
+smmap=5.0.0=pypi_0
+soupsieve=2.3.2.post1=pypi_0
+sqlite=3.33.0=h62c20be_0
+tensorboard=2.10.1=pypi_0
+tensorboard-data-server=0.6.1=pypi_0
+tensorboard-plugin-wit=1.8.1=pypi_0
+thop=0.1.1-2209072238=pypi_0
+tk=8.6.12=h1ccaba5_0
+toml=0.10.2=pypi_0
+torch=1.10.1=pypi_0
+torchaudio=0.10.2=py36_cu113
+torchvision=0.11.2=pypi_0
+tqdm=4.64.1=pypi_0
+traitlets=4.3.3=pypi_0
+typing_extensions=4.1.1=pyh06a4308_0
+urllib3=1.26.12=pypi_0
+wandb=0.13.4=pypi_0
+wcwidth=0.2.5=pypi_0
+werkzeug=2.0.3=pypi_0
+wheel=0.37.1=pyhd3eb1b0_0
+xz=5.2.6=h5eee18b_0
+zipp=3.6.0=pypi_0
+zlib=1.2.12=h5eee18b_3
+zstd=1.5.2=ha4553b6_0