# UNet-MobileNet-Pytorch

**Repository Path**: buptybx/UNet-MobileNet-Pytorch

## Basic Information

- **Project Name**: UNet-MobileNet-Pytorch
- **Description**: Using MobileNet as the backbone of UNet
- **Primary Language**: Unknown
- **License**: MulanPSL-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 2
- **Created**: 2022-03-15
- **Last Updated**: 2022-03-15

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# UNet-MobileNet-Pytorch

### How to run
Glone this repo//
//Download the pretrained model and dataset from my baidunetdisk below
//Put them into the project content
//Create an environment and run:

`pip install -r requirements.txt`

### Prediction

To predict a single image and save it:

`python predict.py -i image.jpg -o output.jpg`

To predict a multiple images and show them without saving them:

`python predict.py -i image1.jpg image2.jpg --viz --no-save`

You can specify which model file to use with `--model MODEL.pth/pt`.

### Training

```shell script
> python train.py -h
usage: train.py [-h] [-e E] [-b [B]] [-l [LR]] [-f LOAD] [-s SCALE] [-v VAL]

Train the UNet on images and target masks

optional arguments:
  -h, --help            show this help message and exit
  -e E, --epochs E      Number of epochs (default: 5)
  -b [B], --batch-size [B]
                        Batch size (default: 1)
  -l [LR], --learning-rate [LR]
                        Learning rate (default: 0.1)
  -f LOAD, --load LOAD  Load model from a .pth file (default: False)
  -s SCALE, --scale SCALE
                        Downscaling factor of the images (default: 0.5)
  -v VAL, --validation VAL
                        Percent of the data that is used as validation (0-100)
                        (default: 15.0)

```
By default, the `scale` is 0.5, so if you wish to obtain better results (but use more memory), set it to 1.

The input images and target masks should be in the `data/imgs` and `data/masks` folders respectively.

## Tensorboard
You can visualize in real time the train and test losses, the weights and gradients, along with the model predictions with tensorboard:

`tensorboard --logdir=runs`

You can find a reference training run with the Caravana dataset on [TensorBoard.dev](https://tensorboard.dev/experiment/1m1Ql50MSJixCbG1m9EcDQ/#scalars&_smoothingWeight=0.6) (only scalars are shown currently).

## Reference
https://github.com/milesial/Pytorch-UNet