├── LICENSE ├── README.md ├── _gitattributes ├── _gitignore ├── data ├── __init__.py ├── coco.py ├── coco_labels.txt ├── config.py ├── example.jpg ├── scripts │ ├── COCO2014.sh │ ├── VOC2007.sh │ └── VOC2012.sh └── voc0712.py ├── demo ├── __init__.py ├── demo.ipynb └── live.py ├── doc ├── SSD.jpg ├── detection_example.png ├── detection_example2.png ├── detection_examples.png └── ssd.png ├── eval_refinedet.py ├── eval_refinedet.sh ├── layers ├── __init__.py ├── box_utils.py ├── functions │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── detection.cpython-36.pyc │ │ ├── detection_refinedet.cpython-36.pyc │ │ └── prior_box.cpython-36.pyc │ ├── detection.py │ ├── detection_refinedet.py │ └── prior_box.py └── modules │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── l2norm.cpython-36.pyc │ ├── multibox_loss.cpython-36.pyc │ └── refinedet_multibox_loss.cpython-36.pyc │ ├── l2norm.py │ ├── multibox_loss.py │ └── refinedet_multibox_loss.py ├── models └── refinedet.py ├── train_refinedet.py ├── train_refinedet320.sh ├── train_refinedet512.sh └── utils ├── __init__.py ├── augmentations.py ├── logging.py └── osutils.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Max deGroot, Ellis Brown 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | A higher performance [PyTorch](http://pytorch.org/) implementation of [Single-Shot Refinement Neural Network for Object Detection](https://arxiv.org/abs/1711.06897 ). The official and original Caffe code can be found [here](https://github.com/sfzhang15/RefineDet). 2 | 3 | ### Table of Contents 4 | - Performance 5 | - Installation 6 | - Datasets 7 | - Train 8 | - Evaluate 9 | - Future Work 10 | - Reference 11 | 12 | 13 | 14 | 15 | 16 | 17 | ## Performance 18 | 19 | #### VOC2007 Test 20 | 21 | ##### mAP (*Single Scale Test*) 22 | 23 | | Arch | Paper | Caffe Version | Our PyTorch Version | 24 | |:-:|:-:|:-:|:-:| 25 | | RefineDet320 | 80.0% | 79.52% | 79.81% | 26 | | RefineDet512 | 81.8% | 81.85% | 80.50% | 27 | 28 | ## Installation 29 | - Install [PyTorch](http://pytorch.org/) by selecting your environment on the website and running the appropriate command. 30 | * Note: You should use at least PyTorch0.4.0 31 | - Clone this repository. 32 | * Note: We currently only support Python 3+. 33 | - Then download the dataset by following the [instructions](#datasets) below. 34 | - We now support [Visdom](https://github.com/facebookresearch/visdom) for real-time loss visualization during training! 35 | * To use Visdom in the browser: 36 | ```Shell 37 | # First install Python server and client 38 | pip install visdom 39 | # Start the server (probably in a screen or tmux) 40 | python -m visdom.server 41 | ``` 42 | * Then (during training) navigate to http://localhost:8097/ (see the Train section below for training details). 43 | - Note: For training, we currently support [VOC](http://host.robots.ox.ac.uk/pascal/VOC/) and [COCO](http://mscoco.org/), and aim to add [ImageNet](http://www.image-net.org/) support soon. 44 | 45 | ## Datasets 46 | To make things easy, we provide bash scripts to handle the dataset downloads and setup for you. We also provide simple dataset loaders that inherit `torch.utils.data.Dataset`, making them fully compatible with the `torchvision.datasets` [API](http://pytorch.org/docs/torchvision/datasets.html). 47 | 48 | 49 | ### COCO 50 | Microsoft COCO: Common Objects in Context 51 | 52 | ##### Download COCO 2014 53 | ```Shell 54 | # specify a directory for dataset to be downloaded into, else default is ~/data/ 55 | sh data/scripts/COCO2014.sh 56 | ``` 57 | 58 | ### VOC Dataset 59 | PASCAL VOC: Visual Object Classes 60 | 61 | ##### Download VOC2007 trainval & test 62 | ```Shell 63 | # specify a directory for dataset to be downloaded into, else default is ~/data/ 64 | sh data/scripts/VOC2007.sh # 65 | ``` 66 | 67 | ##### Download VOC2012 trainval 68 | ```Shell 69 | # specify a directory for dataset to be downloaded into, else default is ~/data/ 70 | sh data/scripts/VOC2012.sh # 71 | ``` 72 | 73 | ## Training RefineDet 74 | - First download the fc-reduced [VGG-16](https://arxiv.org/abs/1409.1556) PyTorch base network weights at: https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth 75 | - By default, we assume you have downloaded the file in the `RefineDet.PyTorch/weights` dir: 76 | 77 | ```Shell 78 | mkdir weights 79 | cd weights 80 | wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth 81 | ``` 82 | 83 | - To train RefineDet320 or RefineDet512 using the train scripts `train_refinedet320.sh` and `train_refinedet512.sh`. You can manually change them as you want. 84 | 85 | ```Shell 86 | ./train_refinedet320.sh #./train_refinedet512.sh 87 | ``` 88 | 89 | - Note: 90 | * For training, an NVIDIA GPU is strongly recommended for speed. 91 | * For instructions on Visdom usage/installation, see the Installation section. 92 | * You can pick-up training from a checkpoint by specifying the path as one of the training parameters (again, see `train_refinedet.py` for options) 93 | 94 | ## Evaluation 95 | To evaluate a trained network: 96 | 97 | ```Shell 98 | ./eval_refinedet.sh 99 | ``` 100 | 101 | You can specify the parameters listed in the `eval_refinedet.py` file by flagging them or manually changing them. 102 | 103 | ## TODO 104 | We have accumulated the following to-do list, which we hope to complete in the near future 105 | - Still to come: 106 | * [ ] Support for multi-scale testing 107 | 108 | ## References 109 | - [Original Implementation (CAFFE)](https://github.com/sfzhang15/RefineDet) 110 | - A list of other great SSD ports that were sources of inspiration: 111 | * [amdegroot/ssd.pytorch](https://github.com/amdegroot/ssd.pytorch) 112 | * [lzx1413/PytorchSSD](https://github.com/lzx1413/PytorchSSD) 113 | -------------------------------------------------------------------------------- /_gitattributes: -------------------------------------------------------------------------------- 1 | *.ipynb linguist-language=Python 2 | .ipynb_checkpoints/* linguist-documentation 3 | dev.ipynb linguist-documentation 4 | -------------------------------------------------------------------------------- /_gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | 27 | # PyInstaller 28 | # Usually these files are written by a python script from a template 29 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 30 | *.manifest 31 | *.spec 32 | 33 | # Installer logs 34 | pip-log.txt 35 | pip-delete-this-directory.txt 36 | 37 | # Unit test / coverage reports 38 | htmlcov/ 39 | .tox/ 40 | .coverage 41 | .coverage.* 42 | .cache 43 | nosetests.xml 44 | coverage.xml 45 | *,cover 46 | .hypothesis/ 47 | 48 | # Translations 49 | *.mo 50 | *.pot 51 | 52 | # Django stuff: 53 | *.log 54 | local_settings.py 55 | 56 | # Flask stuff: 57 | instance/ 58 | .webassets-cache 59 | 60 | # Scrapy stuff: 61 | .scrapy 62 | 63 | # Sphinx documentation 64 | docs/_build/ 65 | 66 | # PyBuilder 67 | target/ 68 | 69 | # IPython Notebook 70 | .ipynb_checkpoints 71 | 72 | # pyenv 73 | .python-version 74 | 75 | # celery beat schedule file 76 | celerybeat-schedule 77 | 78 | # dotenv 79 | .env 80 | 81 | # virtualenv 82 | venv/ 83 | ENV/ 84 | 85 | # Spyder project settings 86 | .spyderproject 87 | 88 | # Rope project settings 89 | .ropeproject 90 | 91 | # atom remote-sync package 92 | .remote-sync.json 93 | 94 | # weights 95 | weights/ 96 | 97 | #DS_Store 98 | .DS_Store 99 | 100 | # dev stuff 101 | eval/ 102 | eval.ipynb 103 | dev.ipynb 104 | .vscode/ 105 | 106 | # not ready 107 | videos/ 108 | templates/ 109 | data/ssd_dataloader.py 110 | data/datasets/ 111 | doc/visualize.py 112 | read_results.py 113 | ssd300_120000/ 114 | demos/live 115 | webdemo.py 116 | test_data_aug.py 117 | 118 | # attributes 119 | 120 | # pycharm 121 | .idea/ 122 | 123 | # temp checkout soln 124 | data/datasets/ 125 | data/ssd_dataloader.py 126 | 127 | # pylint 128 | .pylintrc -------------------------------------------------------------------------------- /data/__init__.py: -------------------------------------------------------------------------------- 1 | from .voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES, VOC_ROOT 2 | 3 | #from .coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, COCO_ROOT, get_label_map 4 | from .config import * 5 | import torch 6 | import cv2 7 | import numpy as np 8 | 9 | def detection_collate(batch): 10 | """Custom collate fn for dealing with batches of images that have a different 11 | number of associated object annotations (bounding boxes). 12 | 13 | Arguments: 14 | batch: (tuple) A tuple of tensor images and lists of annotations 15 | 16 | Return: 17 | A tuple containing: 18 | 1) (tensor) batch of images stacked on their 0 dim 19 | 2) (list of tensors) annotations for a given image are stacked on 20 | 0 dim 21 | """ 22 | targets = [] 23 | imgs = [] 24 | for sample in batch: 25 | imgs.append(sample[0]) 26 | targets.append(torch.FloatTensor(sample[1])) 27 | return torch.stack(imgs, 0), targets 28 | 29 | 30 | def base_transform(image, size, mean): 31 | x = cv2.resize(image, (size, size)).astype(np.float32) 32 | x -= mean 33 | x = x.astype(np.float32) 34 | return x 35 | 36 | 37 | class BaseTransform: 38 | def __init__(self, size, mean): 39 | self.size = size 40 | self.mean = np.array(mean, dtype=np.float32) 41 | 42 | def __call__(self, image, boxes=None, labels=None): 43 | return base_transform(image, self.size, self.mean), boxes, labels 44 | -------------------------------------------------------------------------------- /data/coco.py: -------------------------------------------------------------------------------- 1 | from .config import HOME 2 | import os 3 | import os.path as osp 4 | import sys 5 | import torch 6 | import torch.utils.data as data 7 | import torchvision.transforms as transforms 8 | import cv2 9 | import numpy as np 10 | 11 | COCO_ROOT = osp.join(HOME, 'data/coco/') 12 | IMAGES = 'images' 13 | ANNOTATIONS = 'annotations' 14 | COCO_API = 'PythonAPI' 15 | INSTANCES_SET = 'instances_{}.json' 16 | COCO_CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 17 | 'train', 'truck', 'boat', 'traffic light', 'fire', 'hydrant', 18 | 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 19 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 20 | 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 21 | 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 22 | 'kite', 'baseball bat', 'baseball glove', 'skateboard', 23 | 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 24 | 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 25 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 26 | 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 27 | 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 28 | 'keyboard', 'cell phone', 'microwave oven', 'toaster', 'sink', 29 | 'refrigerator', 'book', 'clock', 'vase', 'scissors', 30 | 'teddy bear', 'hair drier', 'toothbrush') 31 | 32 | 33 | def get_label_map(label_file): 34 | label_map = {} 35 | labels = open(label_file, 'r') 36 | for line in labels: 37 | ids = line.split(',') 38 | label_map[int(ids[0])] = int(ids[1]) 39 | return label_map 40 | 41 | 42 | class COCOAnnotationTransform(object): 43 | """Transforms a COCO annotation into a Tensor of bbox coords and label index 44 | Initilized with a dictionary lookup of classnames to indexes 45 | """ 46 | def __init__(self): 47 | self.label_map = get_label_map(osp.join(COCO_ROOT, 'coco_labels.txt')) 48 | 49 | def __call__(self, target, width, height): 50 | """ 51 | Args: 52 | target (dict): COCO target json annotation as a python dict 53 | height (int): height 54 | width (int): width 55 | Returns: 56 | a list containing lists of bounding boxes [bbox coords, class idx] 57 | """ 58 | scale = np.array([width, height, width, height]) 59 | res = [] 60 | for obj in target: 61 | if 'bbox' in obj: 62 | bbox = obj['bbox'] 63 | bbox[2] += bbox[0] 64 | bbox[3] += bbox[1] 65 | label_idx = self.label_map[obj['category_id']] - 1 66 | final_box = list(np.array(bbox)/scale) 67 | final_box.append(label_idx) 68 | res += [final_box] # [xmin, ymin, xmax, ymax, label_idx] 69 | else: 70 | print("no bbox problem!") 71 | 72 | return res # [[xmin, ymin, xmax, ymax, label_idx], ... ] 73 | 74 | 75 | class COCODetection(data.Dataset): 76 | """`MS Coco Detection `_ Dataset. 77 | Args: 78 | root (string): Root directory where images are downloaded to. 79 | set_name (string): Name of the specific set of COCO images. 80 | transform (callable, optional): A function/transform that augments the 81 | raw images` 82 | target_transform (callable, optional): A function/transform that takes 83 | in the target (bbox) and transforms it. 84 | """ 85 | 86 | def __init__(self, root, image_set='trainval35k', transform=None, 87 | target_transform=COCOAnnotationTransform(), dataset_name='MS COCO'): 88 | sys.path.append(osp.join(root, COCO_API)) 89 | from pycocotools.coco import COCO 90 | self.root = osp.join(root, IMAGES, image_set) 91 | self.coco = COCO(osp.join(root, ANNOTATIONS, 92 | INSTANCES_SET.format(image_set))) 93 | self.ids = list(self.coco.imgToAnns.keys()) 94 | self.transform = transform 95 | self.target_transform = target_transform 96 | self.name = dataset_name 97 | 98 | def __getitem__(self, index): 99 | """ 100 | Args: 101 | index (int): Index 102 | Returns: 103 | tuple: Tuple (image, target). 104 | target is the object returned by ``coco.loadAnns``. 105 | """ 106 | im, gt, h, w = self.pull_item(index) 107 | return im, gt 108 | 109 | def __len__(self): 110 | return len(self.ids) 111 | 112 | def pull_item(self, index): 113 | """ 114 | Args: 115 | index (int): Index 116 | Returns: 117 | tuple: Tuple (image, target, height, width). 118 | target is the object returned by ``coco.loadAnns``. 119 | """ 120 | img_id = self.ids[index] 121 | target = self.coco.imgToAnns[img_id] 122 | ann_ids = self.coco.getAnnIds(imgIds=img_id) 123 | 124 | target = self.coco.loadAnns(ann_ids) 125 | path = osp.join(self.root, self.coco.loadImgs(img_id)[0]['file_name']) 126 | assert osp.exists(path), 'Image path does not exist: {}'.format(path) 127 | img = cv2.imread(osp.join(self.root, path)) 128 | height, width, _ = img.shape 129 | if self.target_transform is not None: 130 | target = self.target_transform(target, width, height) 131 | if self.transform is not None: 132 | target = np.array(target) 133 | img, boxes, labels = self.transform(img, target[:, :4], 134 | target[:, 4]) 135 | # to rgb 136 | img = img[:, :, (2, 1, 0)] 137 | 138 | target = np.hstack((boxes, np.expand_dims(labels, axis=1))) 139 | return torch.from_numpy(img).permute(2, 0, 1), target, height, width 140 | 141 | def pull_image(self, index): 142 | '''Returns the original image object at index in PIL form 143 | 144 | Note: not using self.__getitem__(), as any transformations passed in 145 | could mess up this functionality. 146 | 147 | Argument: 148 | index (int): index of img to show 149 | Return: 150 | cv2 img 151 | ''' 152 | img_id = self.ids[index] 153 | path = self.coco.loadImgs(img_id)[0]['file_name'] 154 | return cv2.imread(osp.join(self.root, path), cv2.IMREAD_COLOR) 155 | 156 | def pull_anno(self, index): 157 | '''Returns the original annotation of image at index 158 | 159 | Note: not using self.__getitem__(), as any transformations passed in 160 | could mess up this functionality. 161 | 162 | Argument: 163 | index (int): index of img to get annotation of 164 | Return: 165 | list: [img_id, [(label, bbox coords),...]] 166 | eg: ('001718', [('dog', (96, 13, 438, 332))]) 167 | ''' 168 | img_id = self.ids[index] 169 | ann_ids = self.coco.getAnnIds(imgIds=img_id) 170 | return self.coco.loadAnns(ann_ids) 171 | 172 | def __repr__(self): 173 | fmt_str = 'Dataset ' + self.__class__.__name__ + '\n' 174 | fmt_str += ' Number of datapoints: {}\n'.format(self.__len__()) 175 | fmt_str += ' Root Location: {}\n'.format(self.root) 176 | tmp = ' Transforms (if any): ' 177 | fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 178 | tmp = ' Target Transforms (if any): ' 179 | fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 180 | return fmt_str 181 | -------------------------------------------------------------------------------- /data/coco_labels.txt: -------------------------------------------------------------------------------- 1 | 1,1,person 2 | 2,2,bicycle 3 | 3,3,car 4 | 4,4,motorcycle 5 | 5,5,airplane 6 | 6,6,bus 7 | 7,7,train 8 | 8,8,truck 9 | 9,9,boat 10 | 10,10,traffic light 11 | 11,11,fire hydrant 12 | 13,12,stop sign 13 | 14,13,parking meter 14 | 15,14,bench 15 | 16,15,bird 16 | 17,16,cat 17 | 18,17,dog 18 | 19,18,horse 19 | 20,19,sheep 20 | 21,20,cow 21 | 22,21,elephant 22 | 23,22,bear 23 | 24,23,zebra 24 | 25,24,giraffe 25 | 27,25,backpack 26 | 28,26,umbrella 27 | 31,27,handbag 28 | 32,28,tie 29 | 33,29,suitcase 30 | 34,30,frisbee 31 | 35,31,skis 32 | 36,32,snowboard 33 | 37,33,sports ball 34 | 38,34,kite 35 | 39,35,baseball bat 36 | 40,36,baseball glove 37 | 41,37,skateboard 38 | 42,38,surfboard 39 | 43,39,tennis racket 40 | 44,40,bottle 41 | 46,41,wine glass 42 | 47,42,cup 43 | 48,43,fork 44 | 49,44,knife 45 | 50,45,spoon 46 | 51,46,bowl 47 | 52,47,banana 48 | 53,48,apple 49 | 54,49,sandwich 50 | 55,50,orange 51 | 56,51,broccoli 52 | 57,52,carrot 53 | 58,53,hot dog 54 | 59,54,pizza 55 | 60,55,donut 56 | 61,56,cake 57 | 62,57,chair 58 | 63,58,couch 59 | 64,59,potted plant 60 | 65,60,bed 61 | 67,61,dining table 62 | 70,62,toilet 63 | 72,63,tv 64 | 73,64,laptop 65 | 74,65,mouse 66 | 75,66,remote 67 | 76,67,keyboard 68 | 77,68,cell phone 69 | 78,69,microwave 70 | 79,70,oven 71 | 80,71,toaster 72 | 81,72,sink 73 | 82,73,refrigerator 74 | 84,74,book 75 | 85,75,clock 76 | 86,76,vase 77 | 87,77,scissors 78 | 88,78,teddy bear 79 | 89,79,hair drier 80 | 90,80,toothbrush 81 | -------------------------------------------------------------------------------- /data/config.py: -------------------------------------------------------------------------------- 1 | # config.py 2 | import os.path 3 | 4 | # gets home dir cross platform 5 | # HOME = os.path.expanduser("~") 6 | HOME = '/data' 7 | 8 | # for making bounding boxes pretty 9 | COLORS = ((255, 0, 0, 128), (0, 255, 0, 128), (0, 0, 255, 128), 10 | (0, 255, 255, 128), (255, 0, 255, 128), (255, 255, 0, 128)) 11 | 12 | MEANS = (104, 117, 123) 13 | 14 | # SSD CONFIGS 15 | voc = { 16 | '300': { 17 | 'num_classes': 21, 18 | 'lr_steps': (80000, 100000, 120000), 19 | 'max_iter': 120000, 20 | 'feature_maps': [38, 19, 10, 5, 3, 1], 21 | 'min_dim': 300, 22 | 'steps': [8, 16, 32, 64, 100, 300], 23 | 'min_sizes': [30, 60, 111, 162, 213, 264], 24 | 'max_sizes': [60, 111, 162, 213, 264, 315], 25 | 'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]], 26 | 'variance': [0.1, 0.2], 27 | 'clip': True, 28 | 'name': 'VOC_300', 29 | }, 30 | '512': { 31 | 'num_classes': 21, 32 | 'lr_steps': (80000, 100000, 120000), 33 | 'max_iter': 120000, 34 | 'feature_maps': [64, 32, 16, 8, 4, 2, 1], 35 | 'min_dim': 512, 36 | 'steps': [8, 16, 32, 64, 128, 256, 512], 37 | 'min_sizes': [20, 51, 133, 215, 296, 378, 460], 38 | 'max_sizes': [51, 133, 215, 296, 378, 460, 542], 39 | 'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]], 40 | 'variance': [0.1, 0.2], 41 | 'clip': True, 42 | 'name': 'VOC_512', 43 | } 44 | } 45 | 46 | coco = { 47 | 'num_classes': 201, 48 | 'lr_steps': (280000, 360000, 400000), 49 | 'max_iter': 400000, 50 | 'feature_maps': [38, 19, 10, 5, 3, 1], 51 | 'min_dim': 300, 52 | 'steps': [8, 16, 32, 64, 100, 300], 53 | 'min_sizes': [21, 45, 99, 153, 207, 261], 54 | 'max_sizes': [45, 99, 153, 207, 261, 315], 55 | 'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]], 56 | 'variance': [0.1, 0.2], 57 | 'clip': True, 58 | 'name': 'COCO', 59 | } 60 | 61 | # RefineDet CONFIGS 62 | voc_refinedet = { 63 | '320': { 64 | 'num_classes': 21, 65 | 'lr_steps': (80000, 100000, 120000), 66 | 'max_iter': 120000, 67 | 'feature_maps': [40, 20, 10, 5], 68 | 'min_dim': 320, 69 | 'steps': [8, 16, 32, 64], 70 | 'min_sizes': [32, 64, 128, 256], 71 | 'max_sizes': [], 72 | 'aspect_ratios': [[2], [2], [2], [2]], 73 | 'variance': [0.1, 0.2], 74 | 'clip': True, 75 | 'name': 'RefineDet_VOC_320', 76 | }, 77 | '512': { 78 | 'num_classes': 21, 79 | 'lr_steps': (80000, 100000, 120000), 80 | 'max_iter': 120000, 81 | 'feature_maps': [64, 32, 16, 8], 82 | 'min_dim': 512, 83 | 'steps': [8, 16, 32, 64], 84 | 'min_sizes': [32, 64, 128, 256], 85 | 'max_sizes': [], 86 | 'aspect_ratios': [[2], [2], [2], [2]], 87 | 'variance': [0.1, 0.2], 88 | 'clip': True, 89 | 'name': 'RefineDet_VOC_320', 90 | } 91 | } 92 | 93 | coco_refinedet = { 94 | 'num_classes': 201, 95 | 'lr_steps': (280000, 360000, 400000), 96 | 'max_iter': 400000, 97 | 'feature_maps': [38, 19, 10, 5, 3, 1], 98 | 'min_dim': 300, 99 | 'steps': [8, 16, 32, 64, 100, 300], 100 | 'min_sizes': [21, 45, 99, 153, 207, 261], 101 | 'max_sizes': [45, 99, 153, 207, 261, 315], 102 | 'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]], 103 | 'variance': [0.1, 0.2], 104 | 'clip': True, 105 | 'name': 'COCO', 106 | } 107 | -------------------------------------------------------------------------------- /data/example.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/data/example.jpg -------------------------------------------------------------------------------- /data/scripts/COCO2014.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | start=`date +%s` 4 | 5 | # handle optional download dir 6 | if [ -z "$1" ] 7 | then 8 | # navigate to ~/data 9 | echo "navigating to ~/data/ ..." 10 | mkdir -p ~/data 11 | cd ~/data/ 12 | mkdir -p ./coco 13 | cd ./coco 14 | mkdir -p ./images 15 | mkdir -p ./annotations 16 | else 17 | # check if specified dir is valid 18 | if [ ! -d $1 ]; then 19 | echo $1 " is not a valid directory" 20 | exit 0 21 | fi 22 | echo "navigating to " $1 " ..." 23 | cd $1 24 | fi 25 | 26 | if [ ! -d images ] 27 | then 28 | mkdir -p ./images 29 | fi 30 | 31 | # Download the image data. 32 | cd ./images 33 | echo "Downloading MSCOCO train images ..." 34 | curl -LO http://images.cocodataset.org/zips/train2014.zip 35 | echo "Downloading MSCOCO val images ..." 36 | curl -LO http://images.cocodataset.org/zips/val2014.zip 37 | 38 | cd ../ 39 | if [ ! -d annotations] 40 | then 41 | mkdir -p ./annotations 42 | fi 43 | 44 | # Download the annotation data. 45 | cd ./annotations 46 | echo "Downloading MSCOCO train/val annotations ..." 47 | curl -LO http://images.cocodataset.org/annotations/annotations_trainval2014.zip 48 | echo "Finished downloading. Now extracting ..." 49 | 50 | # Unzip data 51 | echo "Extracting train images ..." 52 | unzip ../images/train2014.zip -d ../images 53 | echo "Extracting val images ..." 54 | unzip ../images/val2014.zip -d ../images 55 | echo "Extracting annotations ..." 56 | unzip ./annotations_trainval2014.zip 57 | 58 | echo "Removing zip files ..." 59 | rm ../images/train2014.zip 60 | rm ../images/val2014.zip 61 | rm ./annotations_trainval2014.zip 62 | 63 | echo "Creating trainval35k dataset..." 64 | 65 | # Download annotations json 66 | echo "Downloading trainval35k annotations from S3" 67 | curl -LO https://s3.amazonaws.com/amdegroot-datasets/instances_trainval35k.json.zip 68 | 69 | # combine train and val 70 | echo "Combining train and val images" 71 | mkdir ../images/trainval35k 72 | cd ../images/train2014 73 | find -maxdepth 1 -name '*.jpg' -exec cp -t ../trainval35k {} + # dir too large for cp 74 | cd ../val2014 75 | find -maxdepth 1 -name '*.jpg' -exec cp -t ../trainval35k {} + 76 | 77 | 78 | end=`date +%s` 79 | runtime=$((end-start)) 80 | 81 | echo "Completed in " $runtime " seconds" 82 | -------------------------------------------------------------------------------- /data/scripts/VOC2007.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Ellis Brown 3 | 4 | start=`date +%s` 5 | 6 | # handle optional download dir 7 | if [ -z "$1" ] 8 | then 9 | # navigate to ~/data 10 | echo "navigating to ~/data/ ..." 11 | mkdir -p ~/data 12 | cd ~/data/ 13 | else 14 | # check if is valid directory 15 | if [ ! -d $1 ]; then 16 | echo $1 "is not a valid directory" 17 | exit 0 18 | fi 19 | echo "navigating to" $1 "..." 20 | cd $1 21 | fi 22 | 23 | echo "Downloading VOC2007 trainval ..." 24 | # Download the data. 25 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar 26 | echo "Downloading VOC2007 test data ..." 27 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar 28 | echo "Done downloading." 29 | 30 | # Extract data 31 | echo "Extracting trainval ..." 32 | tar -xvf VOCtrainval_06-Nov-2007.tar 33 | echo "Extracting test ..." 34 | tar -xvf VOCtest_06-Nov-2007.tar 35 | echo "removing tars ..." 36 | rm VOCtrainval_06-Nov-2007.tar 37 | rm VOCtest_06-Nov-2007.tar 38 | 39 | end=`date +%s` 40 | runtime=$((end-start)) 41 | 42 | echo "Completed in" $runtime "seconds" -------------------------------------------------------------------------------- /data/scripts/VOC2012.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Ellis Brown 3 | 4 | start=`date +%s` 5 | 6 | # handle optional download dir 7 | if [ -z "$1" ] 8 | then 9 | # navigate to ~/data 10 | echo "navigating to ~/data/ ..." 11 | mkdir -p ~/data 12 | cd ~/data/ 13 | else 14 | # check if is valid directory 15 | if [ ! -d $1 ]; then 16 | echo $1 "is not a valid directory" 17 | exit 0 18 | fi 19 | echo "navigating to" $1 "..." 20 | cd $1 21 | fi 22 | 23 | echo "Downloading VOC2012 trainval ..." 24 | # Download the data. 25 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar 26 | echo "Done downloading." 27 | 28 | 29 | # Extract data 30 | echo "Extracting trainval ..." 31 | tar -xvf VOCtrainval_11-May-2012.tar 32 | echo "removing tar ..." 33 | rm VOCtrainval_11-May-2012.tar 34 | 35 | end=`date +%s` 36 | runtime=$((end-start)) 37 | 38 | echo "Completed in" $runtime "seconds" -------------------------------------------------------------------------------- /data/voc0712.py: -------------------------------------------------------------------------------- 1 | """VOC Dataset Classes 2 | 3 | Original author: Francisco Massa 4 | https://github.com/fmassa/vision/blob/voc_dataset/torchvision/datasets/voc.py 5 | 6 | Updated by: Ellis Brown, Max deGroot 7 | """ 8 | from .config import HOME 9 | import os.path as osp 10 | import sys 11 | import torch 12 | import torch.utils.data as data 13 | import cv2 14 | import numpy as np 15 | if sys.version_info[0] == 2: 16 | import xml.etree.cElementTree as ET 17 | else: 18 | import xml.etree.ElementTree as ET 19 | 20 | VOC_CLASSES = ( # always index 0 21 | 'aeroplane', 'bicycle', 'bird', 'boat', 22 | 'bottle', 'bus', 'car', 'cat', 'chair', 23 | 'cow', 'diningtable', 'dog', 'horse', 24 | 'motorbike', 'person', 'pottedplant', 25 | 'sheep', 'sofa', 'train', 'tvmonitor') 26 | 27 | # note: if you used our download scripts, this should be right 28 | VOC_ROOT = osp.join(HOME, "datasets/VOCdevkit/") 29 | 30 | 31 | class VOCAnnotationTransform(object): 32 | """Transforms a VOC annotation into a Tensor of bbox coords and label index 33 | Initilized with a dictionary lookup of classnames to indexes 34 | 35 | Arguments: 36 | class_to_ind (dict, optional): dictionary lookup of classnames -> indexes 37 | (default: alphabetic indexing of VOC's 20 classes) 38 | keep_difficult (bool, optional): keep difficult instances or not 39 | (default: False) 40 | height (int): height 41 | width (int): width 42 | """ 43 | 44 | def __init__(self, class_to_ind=None, keep_difficult=False): 45 | self.class_to_ind = class_to_ind or dict( 46 | zip(VOC_CLASSES, range(len(VOC_CLASSES)))) 47 | self.keep_difficult = keep_difficult 48 | 49 | def __call__(self, target, width, height): 50 | """ 51 | Arguments: 52 | target (annotation) : the target annotation to be made usable 53 | will be an ET.Element 54 | Returns: 55 | a list containing lists of bounding boxes [bbox coords, class name] 56 | """ 57 | res = [] 58 | for obj in target.iter('object'): 59 | difficult = int(obj.find('difficult').text) == 1 60 | if not self.keep_difficult and difficult: 61 | continue 62 | name = obj.find('name').text.lower().strip() 63 | bbox = obj.find('bndbox') 64 | 65 | pts = ['xmin', 'ymin', 'xmax', 'ymax'] 66 | bndbox = [] 67 | for i, pt in enumerate(pts): 68 | cur_pt = int(bbox.find(pt).text) - 1 69 | # scale height or width 70 | cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height 71 | bndbox.append(cur_pt) 72 | label_idx = self.class_to_ind[name] 73 | bndbox.append(label_idx) 74 | res += [bndbox] # [xmin, ymin, xmax, ymax, label_ind] 75 | # img_id = target.find('filename').text[:-4] 76 | 77 | return res # [[xmin, ymin, xmax, ymax, label_ind], ... ] 78 | 79 | 80 | class VOCDetection(data.Dataset): 81 | """VOC Detection Dataset Object 82 | 83 | input is image, target is annotation 84 | 85 | Arguments: 86 | root (string): filepath to VOCdevkit folder. 87 | image_set (string): imageset to use (eg. 'train', 'val', 'test') 88 | transform (callable, optional): transformation to perform on the 89 | input image 90 | target_transform (callable, optional): transformation to perform on the 91 | target `annotation` 92 | (eg: take in caption string, return tensor of word indices) 93 | dataset_name (string, optional): which dataset to load 94 | (default: 'VOC2007') 95 | """ 96 | 97 | def __init__(self, root, 98 | image_sets=[('2007', 'trainval'), ('2012', 'trainval')], 99 | transform=None, target_transform=VOCAnnotationTransform(), 100 | dataset_name='VOC0712'): 101 | self.root = root 102 | self.image_set = image_sets 103 | self.transform = transform 104 | self.target_transform = target_transform 105 | self.name = dataset_name 106 | self._annopath = osp.join('%s', 'Annotations', '%s.xml') 107 | self._imgpath = osp.join('%s', 'JPEGImages', '%s.jpg') 108 | self.ids = list() 109 | for (year, name) in image_sets: 110 | rootpath = osp.join(self.root, 'VOC' + year) 111 | for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')): 112 | self.ids.append((rootpath, line.strip())) 113 | 114 | def __getitem__(self, index): 115 | im, gt, h, w = self.pull_item(index) 116 | 117 | return im, gt 118 | 119 | def __len__(self): 120 | return len(self.ids) 121 | 122 | def pull_item(self, index): 123 | img_id = self.ids[index] 124 | 125 | target = ET.parse(self._annopath % img_id).getroot() 126 | img = cv2.imread(self._imgpath % img_id) 127 | height, width, channels = img.shape 128 | 129 | if self.target_transform is not None: 130 | target = self.target_transform(target, width, height) 131 | 132 | if self.transform is not None: 133 | target = np.array(target) 134 | img, boxes, labels = self.transform(img, target[:, :4], target[:, 4]) 135 | # to rgb 136 | img = img[:, :, (2, 1, 0)] 137 | # img = img.transpose(2, 0, 1) 138 | target = np.hstack((boxes, np.expand_dims(labels, axis=1))) 139 | return torch.from_numpy(img).permute(2, 0, 1), target, height, width 140 | # return torch.from_numpy(img), target, height, width 141 | 142 | def pull_image(self, index): 143 | '''Returns the original image object at index in PIL form 144 | 145 | Note: not using self.__getitem__(), as any transformations passed in 146 | could mess up this functionality. 147 | 148 | Argument: 149 | index (int): index of img to show 150 | Return: 151 | PIL img 152 | ''' 153 | img_id = self.ids[index] 154 | return cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR) 155 | 156 | def pull_anno(self, index): 157 | '''Returns the original annotation of image at index 158 | 159 | Note: not using self.__getitem__(), as any transformations passed in 160 | could mess up this functionality. 161 | 162 | Argument: 163 | index (int): index of img to get annotation of 164 | Return: 165 | list: [img_id, [(label, bbox coords),...]] 166 | eg: ('001718', [('dog', (96, 13, 438, 332))]) 167 | ''' 168 | img_id = self.ids[index] 169 | anno = ET.parse(self._annopath % img_id).getroot() 170 | gt = self.target_transform(anno, 1, 1) 171 | return img_id[1], gt 172 | 173 | def pull_tensor(self, index): 174 | '''Returns the original image at an index in tensor form 175 | 176 | Note: not using self.__getitem__(), as any transformations passed in 177 | could mess up this functionality. 178 | 179 | Argument: 180 | index (int): index of img to show 181 | Return: 182 | tensorized version of img, squeezed 183 | ''' 184 | return torch.Tensor(self.pull_image(index)).unsqueeze_(0) 185 | -------------------------------------------------------------------------------- /demo/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/demo/__init__.py -------------------------------------------------------------------------------- /demo/live.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | from torch.autograd import Variable 4 | import cv2 5 | import time 6 | from imutils.video import FPS, WebcamVideoStream 7 | import argparse 8 | 9 | parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection') 10 | parser.add_argument('--weights', default='weights/ssd_300_VOC0712.pth', 11 | type=str, help='Trained state_dict file path') 12 | parser.add_argument('--cuda', default=False, type=bool, 13 | help='Use cuda in live demo') 14 | args = parser.parse_args() 15 | 16 | COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] 17 | FONT = cv2.FONT_HERSHEY_SIMPLEX 18 | 19 | 20 | def cv2_demo(net, transform): 21 | def predict(frame): 22 | height, width = frame.shape[:2] 23 | x = torch.from_numpy(transform(frame)[0]).permute(2, 0, 1) 24 | x = Variable(x.unsqueeze(0)) 25 | y = net(x) # forward pass 26 | detections = y.data 27 | # scale each detection back up to the image 28 | scale = torch.Tensor([width, height, width, height]) 29 | for i in range(detections.size(1)): 30 | j = 0 31 | while detections[0, i, j, 0] >= 0.6: 32 | pt = (detections[0, i, j, 1:] * scale).cpu().numpy() 33 | cv2.rectangle(frame, 34 | (int(pt[0]), int(pt[1])), 35 | (int(pt[2]), int(pt[3])), 36 | COLORS[i % 3], 2) 37 | cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), 38 | FONT, 2, (255, 255, 255), 2, cv2.LINE_AA) 39 | j += 1 40 | return frame 41 | 42 | # start video stream thread, allow buffer to fill 43 | print("[INFO] starting threaded video stream...") 44 | stream = WebcamVideoStream(src=0).start() # default camera 45 | time.sleep(1.0) 46 | # start fps timer 47 | # loop over frames from the video file stream 48 | while True: 49 | # grab next frame 50 | frame = stream.read() 51 | key = cv2.waitKey(1) & 0xFF 52 | 53 | # update FPS counter 54 | fps.update() 55 | frame = predict(frame) 56 | 57 | # keybindings for display 58 | if key == ord('p'): # pause 59 | while True: 60 | key2 = cv2.waitKey(1) or 0xff 61 | cv2.imshow('frame', frame) 62 | if key2 == ord('p'): # resume 63 | break 64 | cv2.imshow('frame', frame) 65 | if key == 27: # exit 66 | break 67 | 68 | 69 | if __name__ == '__main__': 70 | import sys 71 | from os import path 72 | sys.path.append(path.dirname(path.dirname(path.abspath(__file__)))) 73 | 74 | from data import BaseTransform, VOC_CLASSES as labelmap 75 | from ssd import build_ssd 76 | 77 | net = build_ssd('test', 300, 21) # initialize SSD 78 | net.load_state_dict(torch.load(args.weights)) 79 | transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0)) 80 | 81 | fps = FPS().start() 82 | cv2_demo(net.eval(), transform) 83 | # stop the timer and display FPS information 84 | fps.stop() 85 | 86 | print("[INFO] elasped time: {:.2f}".format(fps.elapsed())) 87 | print("[INFO] approx. FPS: {:.2f}".format(fps.fps())) 88 | 89 | # cleanup 90 | cv2.destroyAllWindows() 91 | stream.stop() 92 | -------------------------------------------------------------------------------- /doc/SSD.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/doc/SSD.jpg -------------------------------------------------------------------------------- /doc/detection_example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/doc/detection_example.png -------------------------------------------------------------------------------- /doc/detection_example2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/doc/detection_example2.png -------------------------------------------------------------------------------- /doc/detection_examples.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/doc/detection_examples.png -------------------------------------------------------------------------------- /doc/ssd.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/doc/ssd.png -------------------------------------------------------------------------------- /eval_refinedet.py: -------------------------------------------------------------------------------- 1 | """Adapted from: 2 | @longcw faster_rcnn_pytorch: https://github.com/longcw/faster_rcnn_pytorch 3 | @rbgirshick py-faster-rcnn https://github.com/rbgirshick/py-faster-rcnn 4 | Licensed under The MIT License [see LICENSE for details] 5 | """ 6 | 7 | from __future__ import print_function 8 | import torch 9 | import torch.nn as nn 10 | import torch.backends.cudnn as cudnn 11 | from torch.autograd import Variable 12 | from data import VOC_ROOT, VOCAnnotationTransform, VOCDetection, BaseTransform 13 | from data import VOC_CLASSES as labelmap 14 | import torch.utils.data as data 15 | 16 | from models.refinedet import build_refinedet 17 | 18 | import sys 19 | import os 20 | import time 21 | import argparse 22 | import numpy as np 23 | import pickle 24 | import cv2 25 | 26 | if sys.version_info[0] == 2: 27 | import xml.etree.cElementTree as ET 28 | else: 29 | import xml.etree.ElementTree as ET 30 | 31 | 32 | def str2bool(v): 33 | return v.lower() in ("yes", "true", "t", "1") 34 | 35 | 36 | parser = argparse.ArgumentParser( 37 | description='Single Shot MultiBox Detector Evaluation') 38 | parser.add_argument('--trained_model', 39 | default='weights/ssd300_mAP_77.43_v2.pth', type=str, 40 | help='Trained state_dict file path to open') 41 | parser.add_argument('--save_folder', default='eval/', type=str, 42 | help='File path to save results') 43 | parser.add_argument('--confidence_threshold', default=0.01, type=float, 44 | help='Detection confidence threshold') 45 | parser.add_argument('--top_k', default=5, type=int, 46 | help='Further restrict the number of predictions to parse') 47 | parser.add_argument('--cuda', default=True, type=str2bool, 48 | help='Use cuda to train model') 49 | parser.add_argument('--voc_root', default=VOC_ROOT, 50 | help='Location of VOC root directory') 51 | parser.add_argument('--cleanup', default=True, type=str2bool, 52 | help='Cleanup and remove results files following eval') 53 | parser.add_argument('--input_size', default='320', choices=['320', '512'], 54 | type=str, help='RefineDet320 or RefineDet512') 55 | 56 | args = parser.parse_args() 57 | 58 | if not os.path.exists(args.save_folder): 59 | os.mkdir(args.save_folder) 60 | 61 | if torch.cuda.is_available(): 62 | if args.cuda: 63 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 64 | if not args.cuda: 65 | print("WARNING: It looks like you have a CUDA device, but aren't using \ 66 | CUDA. Run with --cuda for optimal eval speed.") 67 | torch.set_default_tensor_type('torch.FloatTensor') 68 | else: 69 | torch.set_default_tensor_type('torch.FloatTensor') 70 | 71 | annopath = os.path.join(args.voc_root, 'VOC2007', 'Annotations', '%s.xml') 72 | imgpath = os.path.join(args.voc_root, 'VOC2007', 'JPEGImages', '%s.jpg') 73 | imgsetpath = os.path.join(args.voc_root, 'VOC2007', 'ImageSets', 74 | 'Main', '{:s}.txt') 75 | YEAR = '2007' 76 | devkit_path = args.voc_root + 'VOC' + YEAR 77 | dataset_mean = (104, 117, 123) 78 | set_type = 'test' 79 | 80 | 81 | class Timer(object): 82 | """A simple timer.""" 83 | def __init__(self): 84 | self.total_time = 0. 85 | self.calls = 0 86 | self.start_time = 0. 87 | self.diff = 0. 88 | self.average_time = 0. 89 | 90 | def tic(self): 91 | # using time.time instead of time.clock because time time.clock 92 | # does not normalize for multithreading 93 | self.start_time = time.time() 94 | 95 | def toc(self, average=True): 96 | self.diff = time.time() - self.start_time 97 | self.total_time += self.diff 98 | self.calls += 1 99 | self.average_time = self.total_time / self.calls 100 | if average: 101 | return self.average_time 102 | else: 103 | return self.diff 104 | 105 | 106 | def parse_rec(filename): 107 | """ Parse a PASCAL VOC xml file """ 108 | tree = ET.parse(filename) 109 | objects = [] 110 | for obj in tree.findall('object'): 111 | obj_struct = {} 112 | obj_struct['name'] = obj.find('name').text 113 | obj_struct['pose'] = obj.find('pose').text 114 | obj_struct['truncated'] = int(obj.find('truncated').text) 115 | obj_struct['difficult'] = int(obj.find('difficult').text) 116 | bbox = obj.find('bndbox') 117 | obj_struct['bbox'] = [int(bbox.find('xmin').text) - 1, 118 | int(bbox.find('ymin').text) - 1, 119 | int(bbox.find('xmax').text) - 1, 120 | int(bbox.find('ymax').text) - 1] 121 | objects.append(obj_struct) 122 | 123 | return objects 124 | 125 | 126 | def get_output_dir(name, phase): 127 | """Return the directory where experimental artifacts are placed. 128 | If the directory does not exist, it is created. 129 | A canonical path is built using the name from an imdb and a network 130 | (if not None). 131 | """ 132 | filedir = os.path.join(name, phase) 133 | if not os.path.exists(filedir): 134 | os.makedirs(filedir) 135 | return filedir 136 | 137 | 138 | def get_voc_results_file_template(image_set, cls): 139 | # VOCdevkit/VOC2007/results/det_test_aeroplane.txt 140 | filename = 'det_' + image_set + '_%s.txt' % (cls) 141 | filedir = os.path.join(devkit_path, 'results') 142 | if not os.path.exists(filedir): 143 | os.makedirs(filedir) 144 | path = os.path.join(filedir, filename) 145 | return path 146 | 147 | 148 | def write_voc_results_file(all_boxes, dataset): 149 | for cls_ind, cls in enumerate(labelmap): 150 | print('Writing {:s} VOC results file'.format(cls)) 151 | filename = get_voc_results_file_template(set_type, cls) 152 | with open(filename, 'wt') as f: 153 | for im_ind, index in enumerate(dataset.ids): 154 | dets = all_boxes[cls_ind+1][im_ind] 155 | if dets == []: 156 | continue 157 | # the VOCdevkit expects 1-based indices 158 | for k in range(dets.shape[0]): 159 | f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'. 160 | format(index[1], dets[k, -1], 161 | dets[k, 0] + 1, dets[k, 1] + 1, 162 | dets[k, 2] + 1, dets[k, 3] + 1)) 163 | 164 | 165 | def do_python_eval(output_dir='output', use_07=True): 166 | cachedir = os.path.join(devkit_path, 'annotations_cache') 167 | aps = [] 168 | # The PASCAL VOC metric changed in 2010 169 | use_07_metric = use_07 170 | print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No')) 171 | if not os.path.isdir(output_dir): 172 | os.mkdir(output_dir) 173 | for i, cls in enumerate(labelmap): 174 | filename = get_voc_results_file_template(set_type, cls) 175 | rec, prec, ap = voc_eval( 176 | filename, annopath, imgsetpath.format(set_type), cls, cachedir, 177 | ovthresh=0.5, use_07_metric=use_07_metric) 178 | aps += [ap] 179 | print('AP for {} = {:.4f}'.format(cls, ap)) 180 | with open(os.path.join(output_dir, cls + '_pr.pkl'), 'wb') as f: 181 | pickle.dump({'rec': rec, 'prec': prec, 'ap': ap}, f) 182 | print('Mean AP = {:.4f}'.format(np.mean(aps))) 183 | print('~~~~~~~~') 184 | print('Results:') 185 | for ap in aps: 186 | print('{:.3f}'.format(ap)) 187 | print('{:.3f}'.format(np.mean(aps))) 188 | print('~~~~~~~~') 189 | print('') 190 | print('--------------------------------------------------------------') 191 | print('Results computed with the **unofficial** Python eval code.') 192 | print('Results should be very close to the official MATLAB eval code.') 193 | print('--------------------------------------------------------------') 194 | 195 | 196 | def voc_ap(rec, prec, use_07_metric=True): 197 | """ ap = voc_ap(rec, prec, [use_07_metric]) 198 | Compute VOC AP given precision and recall. 199 | If use_07_metric is true, uses the 200 | VOC 07 11 point method (default:True). 201 | """ 202 | if use_07_metric: 203 | # 11 point metric 204 | ap = 0. 205 | for t in np.arange(0., 1.1, 0.1): 206 | if np.sum(rec >= t) == 0: 207 | p = 0 208 | else: 209 | p = np.max(prec[rec >= t]) 210 | ap = ap + p / 11. 211 | else: 212 | # correct AP calculation 213 | # first append sentinel values at the end 214 | mrec = np.concatenate(([0.], rec, [1.])) 215 | mpre = np.concatenate(([0.], prec, [0.])) 216 | 217 | # compute the precision envelope 218 | for i in range(mpre.size - 1, 0, -1): 219 | mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i]) 220 | 221 | # to calculate area under PR curve, look for points 222 | # where X axis (recall) changes value 223 | i = np.where(mrec[1:] != mrec[:-1])[0] 224 | 225 | # and sum (\Delta recall) * prec 226 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) 227 | return ap 228 | 229 | 230 | def voc_eval(detpath, 231 | annopath, 232 | imagesetfile, 233 | classname, 234 | cachedir, 235 | ovthresh=0.5, 236 | use_07_metric=True): 237 | """rec, prec, ap = voc_eval(detpath, 238 | annopath, 239 | imagesetfile, 240 | classname, 241 | [ovthresh], 242 | [use_07_metric]) 243 | Top level function that does the PASCAL VOC evaluation. 244 | detpath: Path to detections 245 | detpath.format(classname) should produce the detection results file. 246 | annopath: Path to annotations 247 | annopath.format(imagename) should be the xml annotations file. 248 | imagesetfile: Text file containing the list of images, one image per line. 249 | classname: Category name (duh) 250 | cachedir: Directory for caching the annotations 251 | [ovthresh]: Overlap threshold (default = 0.5) 252 | [use_07_metric]: Whether to use VOC07's 11 point AP computation 253 | (default True) 254 | """ 255 | # assumes detections are in detpath.format(classname) 256 | # assumes annotations are in annopath.format(imagename) 257 | # assumes imagesetfile is a text file with each line an image name 258 | # cachedir caches the annotations in a pickle file 259 | # first load gt 260 | if not os.path.isdir(cachedir): 261 | os.mkdir(cachedir) 262 | cachefile = os.path.join(cachedir, 'annots.pkl') 263 | # read list of images 264 | with open(imagesetfile, 'r') as f: 265 | lines = f.readlines() 266 | imagenames = [x.strip() for x in lines] 267 | if not os.path.isfile(cachefile): 268 | # load annots 269 | recs = {} 270 | for i, imagename in enumerate(imagenames): 271 | recs[imagename] = parse_rec(annopath % (imagename)) 272 | if i % 100 == 0: 273 | print('Reading annotation for {:d}/{:d}'.format( 274 | i + 1, len(imagenames))) 275 | # save 276 | print('Saving cached annotations to {:s}'.format(cachefile)) 277 | with open(cachefile, 'wb') as f: 278 | pickle.dump(recs, f) 279 | else: 280 | # load 281 | with open(cachefile, 'rb') as f: 282 | recs = pickle.load(f) 283 | 284 | # extract gt objects for this class 285 | class_recs = {} 286 | npos = 0 287 | for imagename in imagenames: 288 | R = [obj for obj in recs[imagename] if obj['name'] == classname] 289 | bbox = np.array([x['bbox'] for x in R]) 290 | difficult = np.array([x['difficult'] for x in R]).astype(np.bool) 291 | det = [False] * len(R) 292 | npos = npos + sum(~difficult) 293 | class_recs[imagename] = {'bbox': bbox, 294 | 'difficult': difficult, 295 | 'det': det} 296 | 297 | # read dets 298 | detfile = detpath.format(classname) 299 | with open(detfile, 'r') as f: 300 | lines = f.readlines() 301 | if any(lines) == 1: 302 | 303 | splitlines = [x.strip().split(' ') for x in lines] 304 | image_ids = [x[0] for x in splitlines] 305 | confidence = np.array([float(x[1]) for x in splitlines]) 306 | BB = np.array([[float(z) for z in x[2:]] for x in splitlines]) 307 | 308 | # sort by confidence 309 | sorted_ind = np.argsort(-confidence) 310 | sorted_scores = np.sort(-confidence) 311 | BB = BB[sorted_ind, :] 312 | image_ids = [image_ids[x] for x in sorted_ind] 313 | 314 | # go down dets and mark TPs and FPs 315 | nd = len(image_ids) 316 | tp = np.zeros(nd) 317 | fp = np.zeros(nd) 318 | for d in range(nd): 319 | R = class_recs[image_ids[d]] 320 | bb = BB[d, :].astype(float) 321 | ovmax = -np.inf 322 | BBGT = R['bbox'].astype(float) 323 | if BBGT.size > 0: 324 | # compute overlaps 325 | # intersection 326 | ixmin = np.maximum(BBGT[:, 0], bb[0]) 327 | iymin = np.maximum(BBGT[:, 1], bb[1]) 328 | ixmax = np.minimum(BBGT[:, 2], bb[2]) 329 | iymax = np.minimum(BBGT[:, 3], bb[3]) 330 | iw = np.maximum(ixmax - ixmin, 0.) 331 | ih = np.maximum(iymax - iymin, 0.) 332 | inters = iw * ih 333 | uni = ((bb[2] - bb[0]) * (bb[3] - bb[1]) + 334 | (BBGT[:, 2] - BBGT[:, 0]) * 335 | (BBGT[:, 3] - BBGT[:, 1]) - inters) 336 | overlaps = inters / uni 337 | ovmax = np.max(overlaps) 338 | jmax = np.argmax(overlaps) 339 | 340 | if ovmax > ovthresh: 341 | if not R['difficult'][jmax]: 342 | if not R['det'][jmax]: 343 | tp[d] = 1. 344 | R['det'][jmax] = 1 345 | else: 346 | fp[d] = 1. 347 | else: 348 | fp[d] = 1. 349 | 350 | # compute precision recall 351 | fp = np.cumsum(fp) 352 | tp = np.cumsum(tp) 353 | rec = tp / float(npos) 354 | # avoid divide by zero in case the first detection matches a difficult 355 | # ground truth 356 | prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps) 357 | ap = voc_ap(rec, prec, use_07_metric) 358 | else: 359 | rec = -1. 360 | prec = -1. 361 | ap = -1. 362 | 363 | return rec, prec, ap 364 | 365 | 366 | def test_net(save_folder, net, cuda, dataset, transform, top_k, 367 | im_size=300, thresh=0.05): 368 | num_images = len(dataset) 369 | # all detections are collected into: 370 | # all_boxes[cls][image] = N x 5 array of detections in 371 | # (x1, y1, x2, y2, score) 372 | all_boxes = [[[] for _ in range(num_images)] 373 | for _ in range(len(labelmap)+1)] 374 | 375 | # timers 376 | _t = {'im_detect': Timer(), 'misc': Timer()} 377 | output_dir = get_output_dir('ssd300_120000', set_type) 378 | det_file = os.path.join(output_dir, 'detections.pkl') 379 | 380 | for i in range(num_images): 381 | im, gt, h, w = dataset.pull_item(i) 382 | 383 | x = Variable(im.unsqueeze(0)) 384 | if args.cuda: 385 | x = x.cuda() 386 | _t['im_detect'].tic() 387 | detections = net(x).data 388 | detect_time = _t['im_detect'].toc(average=False) 389 | 390 | # skip j = 0, because it's the background class 391 | for j in range(1, detections.size(1)): 392 | dets = detections[0, j, :] 393 | mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t() 394 | dets = torch.masked_select(dets, mask).view(-1, 5) 395 | if dets.size(0) == 0: 396 | continue 397 | boxes = dets[:, 1:] 398 | boxes[:, 0] *= w 399 | boxes[:, 2] *= w 400 | boxes[:, 1] *= h 401 | boxes[:, 3] *= h 402 | scores = dets[:, 0].cpu().numpy() 403 | cls_dets = np.hstack((boxes.cpu().numpy(), 404 | scores[:, np.newaxis])).astype(np.float32, 405 | copy=False) 406 | all_boxes[j][i] = cls_dets 407 | 408 | print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, 409 | num_images, detect_time)) 410 | 411 | with open(det_file, 'wb') as f: 412 | pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL) 413 | 414 | print('Evaluating detections') 415 | evaluate_detections(all_boxes, output_dir, dataset) 416 | 417 | 418 | def evaluate_detections(box_list, output_dir, dataset): 419 | write_voc_results_file(box_list, dataset) 420 | do_python_eval(output_dir) 421 | 422 | 423 | if __name__ == '__main__': 424 | # load net 425 | num_classes = len(labelmap) + 1 # +1 for background 426 | net = build_refinedet('test', int(args.input_size), num_classes) # initialize SSD 427 | net.load_state_dict(torch.load(args.trained_model)) 428 | net.eval() 429 | print('Finished loading model!') 430 | # load data 431 | dataset = VOCDetection(args.voc_root, [('2007', set_type)], 432 | BaseTransform(int(args.input_size), dataset_mean), 433 | VOCAnnotationTransform()) 434 | if args.cuda: 435 | net = net.cuda() 436 | cudnn.benchmark = True 437 | # evaluation 438 | test_net(args.save_folder, net, args.cuda, dataset, 439 | BaseTransform(net.size, dataset_mean), args.top_k, int(args.input_size), 440 | thresh=args.confidence_threshold) 441 | -------------------------------------------------------------------------------- /eval_refinedet.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=2 python eval_refinedet.py --trained_model weights/refinedet_testVOC.pth --save_folder eval_refinedet/ 2 | -------------------------------------------------------------------------------- /layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .functions import * 2 | from .modules import * 3 | -------------------------------------------------------------------------------- /layers/box_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import torch 3 | 4 | 5 | def point_form(boxes): 6 | """ Convert prior_boxes to (xmin, ymin, xmax, ymax) 7 | representation for comparison to point form ground truth data. 8 | Args: 9 | boxes: (tensor) center-size default boxes from priorbox layers. 10 | Return: 11 | boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. 12 | """ 13 | return torch.cat((boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin 14 | boxes[:, :2] + boxes[:, 2:]/2), 1) # xmax, ymax 15 | 16 | 17 | def center_size(boxes): 18 | """ Convert prior_boxes to (cx, cy, w, h) 19 | representation for comparison to center-size form ground truth data. 20 | Args: 21 | boxes: (tensor) point_form boxes 22 | Return: 23 | boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. 24 | """ 25 | return torch.cat([(boxes[:, 2:] + boxes[:, :2])/2, # cx, cy 26 | boxes[:, 2:] - boxes[:, :2]], 1) # w, h 27 | 28 | 29 | def intersect(box_a, box_b): 30 | """ We resize both tensors to [A,B,2] without new malloc: 31 | [A,2] -> [A,1,2] -> [A,B,2] 32 | [B,2] -> [1,B,2] -> [A,B,2] 33 | Then we compute the area of intersect between box_a and box_b. 34 | Args: 35 | box_a: (tensor) bounding boxes, Shape: [A,4]. 36 | box_b: (tensor) bounding boxes, Shape: [B,4]. 37 | Return: 38 | (tensor) intersection area, Shape: [A,B]. 39 | """ 40 | A = box_a.size(0) 41 | B = box_b.size(0) 42 | max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), 43 | box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) 44 | min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), 45 | box_b[:, :2].unsqueeze(0).expand(A, B, 2)) 46 | inter = torch.clamp((max_xy - min_xy), min=0) 47 | return inter[:, :, 0] * inter[:, :, 1] 48 | 49 | 50 | def jaccard(box_a, box_b): 51 | """Compute the jaccard overlap of two sets of boxes. The jaccard overlap 52 | is simply the intersection over union of two boxes. Here we operate on 53 | ground truth boxes and default boxes. 54 | E.g.: 55 | A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) 56 | Args: 57 | box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4] 58 | box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4] 59 | Return: 60 | jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)] 61 | """ 62 | inter = intersect(box_a, box_b) 63 | area_a = ((box_a[:, 2]-box_a[:, 0]) * 64 | (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B] 65 | area_b = ((box_b[:, 2]-box_b[:, 0]) * 66 | (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B] 67 | union = area_a + area_b - inter 68 | return inter / union # [A,B] 69 | 70 | def match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx): 71 | """Match each prior box with the ground truth box of the highest jaccard 72 | overlap, encode the bounding boxes, then return the matched indices 73 | corresponding to both confidence and location preds. 74 | Args: 75 | threshold: (float) The overlap threshold used when mathing boxes. 76 | truths: (tensor) Ground truth boxes, Shape: [num_obj, num_priors]. 77 | priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4]. 78 | variances: (tensor) Variances corresponding to each prior coord, 79 | Shape: [num_priors, 4]. 80 | labels: (tensor) All the class labels for the image, Shape: [num_obj]. 81 | loc_t: (tensor) Tensor to be filled w/ endcoded location targets. 82 | conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. 83 | idx: (int) current batch index 84 | Return: 85 | The matched indices corresponding to 1)location and 2)confidence preds. 86 | """ 87 | # jaccard index 88 | overlaps = jaccard( 89 | truths, 90 | point_form(priors) 91 | ) 92 | # (Bipartite Matching) 93 | # [1,num_objects] best prior for each ground truth 94 | best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) 95 | # [1,num_priors] best ground truth for each prior 96 | best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) 97 | best_truth_idx.squeeze_(0) 98 | best_truth_overlap.squeeze_(0) 99 | best_prior_idx.squeeze_(1) 100 | best_prior_overlap.squeeze_(1) 101 | best_truth_overlap.index_fill_(0, best_prior_idx, 2) # ensure best prior 102 | # TODO refactor: index best_prior_idx with long tensor 103 | # ensure every gt matches with its prior of max overlap 104 | for j in range(best_prior_idx.size(0)): 105 | best_truth_idx[best_prior_idx[j]] = j 106 | matches = truths[best_truth_idx] # Shape: [num_priors,4] 107 | conf = labels[best_truth_idx] + 1 # Shape: [num_priors] 108 | conf[best_truth_overlap < threshold] = 0 # label as background 109 | loc = encode(matches, priors, variances) 110 | loc_t[idx] = loc # [num_priors,4] encoded offsets to learn 111 | conf_t[idx] = conf # [num_priors] top class label for each prior 112 | 113 | def refine_match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx, arm_loc=None): 114 | """Match each arm bbox with the ground truth box of the highest jaccard 115 | overlap, encode the bounding boxes, then return the matched indices 116 | corresponding to both confidence and location preds. 117 | Args: 118 | threshold: (float) The overlap threshold used when mathing boxes. 119 | truths: (tensor) Ground truth boxes, Shape: [num_obj, num_priors]. 120 | priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4]. 121 | variances: (tensor) Variances corresponding to each prior coord, 122 | Shape: [num_priors, 4]. 123 | labels: (tensor) All the class labels for the image, Shape: [num_obj]. 124 | loc_t: (tensor) Tensor to be filled w/ endcoded location targets. 125 | conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. 126 | idx: (int) current batch index 127 | Return: 128 | The matched indices corresponding to 1)location and 2)confidence preds. 129 | """ 130 | 131 | # jaccard index 132 | if arm_loc is None: 133 | overlaps = jaccard(truths, point_form(priors)) 134 | else: 135 | decode_arm = decode(arm_loc, priors=priors, variances=variances) 136 | overlaps = jaccard(truths, decode_arm) 137 | # (Bipartite Matching) 138 | # [1,num_objects] best prior for each ground truth 139 | best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) 140 | # [1,num_priors] best ground truth for each prior 141 | best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) 142 | best_truth_idx.squeeze_(0) 143 | best_truth_overlap.squeeze_(0) 144 | best_prior_idx.squeeze_(1) 145 | best_prior_overlap.squeeze_(1) 146 | best_truth_overlap.index_fill_(0, best_prior_idx, 2) # ensure best prior 147 | # TODO refactor: index best_prior_idx with long tensor 148 | # ensure every gt matches with its prior of max overlap 149 | for j in range(best_prior_idx.size(0)): 150 | best_truth_idx[best_prior_idx[j]] = j 151 | matches = truths[best_truth_idx] # Shape: [num_priors,4] 152 | if arm_loc is None: 153 | conf = labels[best_truth_idx] # Shape: [num_priors] 154 | loc = encode(matches, priors, variances) 155 | else: 156 | conf = labels[best_truth_idx] + 1 # Shape: [num_priors] 157 | loc = encode(matches, center_size(decode_arm), variances) 158 | conf[best_truth_overlap < threshold] = 0 # label as background 159 | loc_t[idx] = loc # [num_priors,4] encoded offsets to learn 160 | conf_t[idx] = conf # [num_priors] top class label for each prior 161 | 162 | def encode(matched, priors, variances): 163 | """Encode the variances from the priorbox layers into the ground truth boxes 164 | we have matched (based on jaccard overlap) with the prior boxes. 165 | Args: 166 | matched: (tensor) Coords of ground truth for each prior in point-form 167 | Shape: [num_priors, 4]. 168 | priors: (tensor) Prior boxes in center-offset form 169 | Shape: [num_priors,4]. 170 | variances: (list[float]) Variances of priorboxes 171 | Return: 172 | encoded boxes (tensor), Shape: [num_priors, 4] 173 | """ 174 | 175 | # dist b/t match center and prior's center 176 | g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] 177 | # encode variance 178 | g_cxcy /= (variances[0] * priors[:, 2:]) 179 | # match wh / prior wh 180 | g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] 181 | g_wh = torch.log(g_wh + 1e-5) / variances[1] 182 | # return target for smooth_l1_loss 183 | return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4] 184 | 185 | 186 | # Adapted from https://github.com/Hakuyume/chainer-ssd 187 | def decode(loc, priors, variances): 188 | """Decode locations from predictions using priors to undo 189 | the encoding we did for offset regression at train time. 190 | Args: 191 | loc (tensor): location predictions for loc layers, 192 | Shape: [num_priors,4] 193 | priors (tensor): Prior boxes in center-offset form. 194 | Shape: [num_priors,4]. 195 | variances: (list[float]) Variances of priorboxes 196 | Return: 197 | decoded bounding box predictions 198 | """ 199 | 200 | boxes = torch.cat(( 201 | priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], 202 | priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) 203 | boxes[:, :2] -= boxes[:, 2:] / 2 204 | boxes[:, 2:] += boxes[:, :2] 205 | return boxes 206 | 207 | 208 | def log_sum_exp(x): 209 | """Utility function for computing log_sum_exp while determining 210 | This will be used to determine unaveraged confidence loss across 211 | all examples in a batch. 212 | Args: 213 | x (Variable(tensor)): conf_preds from conf layers 214 | """ 215 | x_max = x.data.max() 216 | return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max 217 | 218 | 219 | # Original author: Francisco Massa: 220 | # https://github.com/fmassa/object-detection.torch 221 | # Ported to PyTorch by Max deGroot (02/01/2017) 222 | def nms(boxes, scores, overlap=0.5, top_k=200): 223 | """Apply non-maximum suppression at test time to avoid detecting too many 224 | overlapping bounding boxes for a given object. 225 | Args: 226 | boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. 227 | scores: (tensor) The class predscores for the img, Shape:[num_priors]. 228 | overlap: (float) The overlap thresh for suppressing unnecessary boxes. 229 | top_k: (int) The Maximum number of box preds to consider. 230 | Return: 231 | The indices of the kept boxes with respect to num_priors. 232 | """ 233 | 234 | keep = scores.new(scores.size(0)).zero_().long() 235 | if boxes.numel() == 0: 236 | return keep 237 | x1 = boxes[:, 0] 238 | y1 = boxes[:, 1] 239 | x2 = boxes[:, 2] 240 | y2 = boxes[:, 3] 241 | area = torch.mul(x2 - x1, y2 - y1) 242 | v, idx = scores.sort(0) # sort in ascending order 243 | # I = I[v >= 0.01] 244 | idx = idx[-top_k:] # indices of the top-k largest vals 245 | xx1 = boxes.new() 246 | yy1 = boxes.new() 247 | xx2 = boxes.new() 248 | yy2 = boxes.new() 249 | w = boxes.new() 250 | h = boxes.new() 251 | 252 | # keep = torch.Tensor() 253 | count = 0 254 | while idx.numel() > 0: 255 | i = idx[-1] # index of current largest val 256 | # keep.append(i) 257 | keep[count] = i 258 | count += 1 259 | if idx.size(0) == 1: 260 | break 261 | idx = idx[:-1] # remove kept element from view 262 | # load bboxes of next highest vals 263 | torch.index_select(x1, 0, idx, out=xx1) 264 | torch.index_select(y1, 0, idx, out=yy1) 265 | torch.index_select(x2, 0, idx, out=xx2) 266 | torch.index_select(y2, 0, idx, out=yy2) 267 | # store element-wise max with next highest score 268 | xx1 = torch.clamp(xx1, min=x1[i]) 269 | yy1 = torch.clamp(yy1, min=y1[i]) 270 | xx2 = torch.clamp(xx2, max=x2[i]) 271 | yy2 = torch.clamp(yy2, max=y2[i]) 272 | w.resize_as_(xx2) 273 | h.resize_as_(yy2) 274 | w = xx2 - xx1 275 | h = yy2 - yy1 276 | # check sizes of xx1 and xx2.. after each iteration 277 | w = torch.clamp(w, min=0.0) 278 | h = torch.clamp(h, min=0.0) 279 | inter = w*h 280 | # IoU = i / (area(a) + area(b) - i) 281 | rem_areas = torch.index_select(area, 0, idx) # load remaining areas) 282 | union = (rem_areas - inter) + area[i] 283 | IoU = inter/union # store result in iou 284 | # keep only elements with an IoU <= overlap 285 | idx = idx[IoU.le(overlap)] 286 | return keep, count 287 | -------------------------------------------------------------------------------- /layers/functions/__init__.py: -------------------------------------------------------------------------------- 1 | from .detection import Detect 2 | from .prior_box import PriorBox 3 | from .detection_refinedet import Detect_RefineDet 4 | 5 | 6 | __all__ = ['Detect', 'PriorBox', 'Detect_RefineDet'] 7 | -------------------------------------------------------------------------------- /layers/functions/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/functions/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /layers/functions/__pycache__/detection.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/functions/__pycache__/detection.cpython-36.pyc -------------------------------------------------------------------------------- /layers/functions/__pycache__/detection_refinedet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/functions/__pycache__/detection_refinedet.cpython-36.pyc -------------------------------------------------------------------------------- /layers/functions/__pycache__/prior_box.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/functions/__pycache__/prior_box.cpython-36.pyc -------------------------------------------------------------------------------- /layers/functions/detection.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from ..box_utils import decode, nms 4 | from data import voc as cfg 5 | 6 | 7 | class Detect(Function): 8 | """At test time, Detect is the final layer of SSD. Decode location preds, 9 | apply non-maximum suppression to location predictions based on conf 10 | scores and threshold to a top_k number of output predictions for both 11 | confidence score and locations. 12 | """ 13 | def __init__(self, num_classes, size, bkg_label, top_k, conf_thresh, nms_thresh): 14 | self.num_classes = num_classes 15 | self.background_label = bkg_label 16 | self.top_k = top_k 17 | # Parameters used in nms. 18 | self.nms_thresh = nms_thresh 19 | if nms_thresh <= 0: 20 | raise ValueError('nms_threshold must be non negative.') 21 | self.conf_thresh = conf_thresh 22 | self.variance = cfg[str(size)]['variance'] 23 | 24 | def forward(self, loc_data, conf_data, prior_data): 25 | """ 26 | Args: 27 | loc_data: (tensor) Loc preds from loc layers 28 | Shape: [batch,num_priors*4] 29 | conf_data: (tensor) Shape: Conf preds from conf layers 30 | Shape: [batch*num_priors,num_classes] 31 | prior_data: (tensor) Prior boxes and variances from priorbox layers 32 | Shape: [1,num_priors,4] 33 | """ 34 | num = loc_data.size(0) # batch size 35 | num_priors = prior_data.size(0) 36 | output = torch.zeros(num, self.num_classes, self.top_k, 5) 37 | conf_preds = conf_data.view(num, num_priors, 38 | self.num_classes).transpose(2, 1) 39 | 40 | # Decode predictions into bboxes. 41 | for i in range(num): 42 | decoded_boxes = decode(loc_data[i], prior_data, self.variance) 43 | # For each class, perform nms 44 | conf_scores = conf_preds[i].clone() 45 | #print(decoded_boxes, conf_scores) 46 | for cl in range(1, self.num_classes): 47 | c_mask = conf_scores[cl].gt(self.conf_thresh) 48 | scores = conf_scores[cl][c_mask] 49 | #print(scores.dim()) 50 | if scores.size(0) == 0: 51 | continue 52 | l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes) 53 | boxes = decoded_boxes[l_mask].view(-1, 4) 54 | # idx of highest scoring and non-overlapping boxes per class 55 | #print(boxes, scores) 56 | ids, count = nms(boxes, scores, self.nms_thresh, self.top_k) 57 | output[i, cl, :count] = \ 58 | torch.cat((scores[ids[:count]].unsqueeze(1), 59 | boxes[ids[:count]]), 1) 60 | flt = output.contiguous().view(num, -1, 5) 61 | _, idx = flt[:, :, 0].sort(1, descending=True) 62 | _, rank = idx.sort(1) 63 | flt[(rank < self.top_k).unsqueeze(-1).expand_as(flt)].fill_(0) 64 | return output 65 | -------------------------------------------------------------------------------- /layers/functions/detection_refinedet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from ..box_utils import decode, nms, center_size 4 | from data import voc_refinedet as cfg 5 | 6 | 7 | class Detect_RefineDet(Function): 8 | """At test time, Detect is the final layer of SSD. Decode location preds, 9 | apply non-maximum suppression to location predictions based on conf 10 | scores and threshold to a top_k number of output predictions for both 11 | confidence score and locations. 12 | """ 13 | def __init__(self, num_classes, size, bkg_label, top_k, conf_thresh, nms_thresh, 14 | objectness_thre, keep_top_k): 15 | self.num_classes = num_classes 16 | self.background_label = bkg_label 17 | self.top_k = top_k 18 | self.keep_top_k = keep_top_k 19 | # Parameters used in nms. 20 | self.nms_thresh = nms_thresh 21 | if nms_thresh <= 0: 22 | raise ValueError('nms_threshold must be non negative.') 23 | self.conf_thresh = conf_thresh 24 | self.objectness_thre = objectness_thre 25 | self.variance = cfg[str(size)]['variance'] 26 | 27 | def forward(self, arm_loc_data, arm_conf_data, odm_loc_data, odm_conf_data, prior_data): 28 | """ 29 | Args: 30 | loc_data: (tensor) Loc preds from loc layers 31 | Shape: [batch,num_priors*4] 32 | conf_data: (tensor) Shape: Conf preds from conf layers 33 | Shape: [batch*num_priors,num_classes] 34 | prior_data: (tensor) Prior boxes and variances from priorbox layers 35 | Shape: [1,num_priors,4] 36 | """ 37 | loc_data = odm_loc_data 38 | conf_data = odm_conf_data 39 | 40 | arm_object_conf = arm_conf_data.data[:, :, 1:] 41 | no_object_index = arm_object_conf <= self.objectness_thre 42 | conf_data[no_object_index.expand_as(conf_data)] = 0 43 | 44 | num = loc_data.size(0) # batch size 45 | num_priors = prior_data.size(0) 46 | output = torch.zeros(num, self.num_classes, self.top_k, 5) 47 | conf_preds = conf_data.view(num, num_priors, 48 | self.num_classes).transpose(2, 1) 49 | 50 | # Decode predictions into bboxes. 51 | for i in range(num): 52 | default = decode(arm_loc_data[i], prior_data, self.variance) 53 | default = center_size(default) 54 | decoded_boxes = decode(loc_data[i], default, self.variance) 55 | # For each class, perform nms 56 | conf_scores = conf_preds[i].clone() 57 | #print(decoded_boxes, conf_scores) 58 | for cl in range(1, self.num_classes): 59 | c_mask = conf_scores[cl].gt(self.conf_thresh) 60 | scores = conf_scores[cl][c_mask] 61 | #print(scores.dim()) 62 | if scores.size(0) == 0: 63 | continue 64 | l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes) 65 | boxes = decoded_boxes[l_mask].view(-1, 4) 66 | # idx of highest scoring and non-overlapping boxes per class 67 | #print(boxes, scores) 68 | ids, count = nms(boxes, scores, self.nms_thresh, self.top_k) 69 | output[i, cl, :count] = \ 70 | torch.cat((scores[ids[:count]].unsqueeze(1), 71 | boxes[ids[:count]]), 1) 72 | flt = output.contiguous().view(num, -1, 5) 73 | _, idx = flt[:, :, 0].sort(1, descending=True) 74 | _, rank = idx.sort(1) 75 | flt[(rank < self.keep_top_k).unsqueeze(-1).expand_as(flt)].fill_(0) 76 | return output 77 | -------------------------------------------------------------------------------- /layers/functions/prior_box.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | from math import sqrt as sqrt 3 | from itertools import product as product 4 | import torch 5 | 6 | 7 | class PriorBox(object): 8 | """Compute priorbox coordinates in center-offset form for each source 9 | feature map. 10 | """ 11 | def __init__(self, cfg): 12 | super(PriorBox, self).__init__() 13 | self.image_size = cfg['min_dim'] 14 | # number of priors for feature map location (either 4 or 6) 15 | self.num_priors = len(cfg['aspect_ratios']) 16 | self.variance = cfg['variance'] or [0.1] 17 | self.feature_maps = cfg['feature_maps'] 18 | self.min_sizes = cfg['min_sizes'] 19 | self.max_sizes = cfg['max_sizes'] 20 | self.steps = cfg['steps'] 21 | self.aspect_ratios = cfg['aspect_ratios'] 22 | self.clip = cfg['clip'] 23 | self.version = cfg['name'] 24 | for v in self.variance: 25 | if v <= 0: 26 | raise ValueError('Variances must be greater than 0') 27 | 28 | def forward(self): 29 | mean = [] 30 | for k, f in enumerate(self.feature_maps): 31 | for i, j in product(range(f), repeat=2): 32 | f_k = self.image_size / self.steps[k] 33 | # unit center x,y 34 | cx = (j + 0.5) / f_k 35 | cy = (i + 0.5) / f_k 36 | 37 | # aspect_ratio: 1 38 | # rel size: min_size 39 | s_k = self.min_sizes[k]/self.image_size 40 | mean += [cx, cy, s_k, s_k] 41 | 42 | # aspect_ratio: 1 43 | # rel size: sqrt(s_k * s_(k+1)) 44 | if self.max_sizes: 45 | s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size)) 46 | mean += [cx, cy, s_k_prime, s_k_prime] 47 | 48 | # rest of aspect ratios 49 | for ar in self.aspect_ratios[k]: 50 | mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)] 51 | mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)] 52 | # back to torch land 53 | output = torch.Tensor(mean).view(-1, 4) 54 | if self.clip: 55 | output.clamp_(max=1, min=0) 56 | return output 57 | -------------------------------------------------------------------------------- /layers/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .l2norm import L2Norm 2 | from .multibox_loss import MultiBoxLoss 3 | from .refinedet_multibox_loss import RefineDetMultiBoxLoss 4 | 5 | __all__ = ['L2Norm', 'MultiBoxLoss', 'RefineDetMultiBoxLoss'] 6 | -------------------------------------------------------------------------------- /layers/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /layers/modules/__pycache__/l2norm.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/modules/__pycache__/l2norm.cpython-36.pyc -------------------------------------------------------------------------------- /layers/modules/__pycache__/multibox_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/modules/__pycache__/multibox_loss.cpython-36.pyc -------------------------------------------------------------------------------- /layers/modules/__pycache__/refinedet_multibox_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luuuyi/RefineDet.PyTorch/0e4b24ce07245fcb8c48292326a731729cc5746a/layers/modules/__pycache__/refinedet_multibox_loss.cpython-36.pyc -------------------------------------------------------------------------------- /layers/modules/l2norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Function 4 | #from torch.autograd import Variable 5 | import torch.nn.init as init 6 | 7 | class L2Norm(nn.Module): 8 | def __init__(self,n_channels, scale): 9 | super(L2Norm,self).__init__() 10 | self.n_channels = n_channels 11 | self.gamma = scale or None 12 | self.eps = 1e-10 13 | self.weight = nn.Parameter(torch.Tensor(self.n_channels)) 14 | self.reset_parameters() 15 | 16 | def reset_parameters(self): 17 | init.constant_(self.weight,self.gamma) 18 | 19 | def forward(self, x): 20 | norm = x.pow(2).sum(dim=1, keepdim=True).sqrt()+self.eps 21 | #x /= norm 22 | x = torch.div(x,norm) 23 | out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x 24 | return out 25 | -------------------------------------------------------------------------------- /layers/modules/multibox_loss.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.autograd import Variable 6 | from data import coco as cfg 7 | from ..box_utils import match, log_sum_exp 8 | 9 | 10 | class MultiBoxLoss(nn.Module): 11 | """SSD Weighted Loss Function 12 | Compute Targets: 13 | 1) Produce Confidence Target Indices by matching ground truth boxes 14 | with (default) 'priorboxes' that have jaccard index > threshold parameter 15 | (default threshold: 0.5). 16 | 2) Produce localization target by 'encoding' variance into offsets of ground 17 | truth boxes and their matched 'priorboxes'. 18 | 3) Hard negative mining to filter the excessive number of negative examples 19 | that comes with using a large number of default bounding boxes. 20 | (default negative:positive ratio 3:1) 21 | Objective Loss: 22 | L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 23 | Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss 24 | weighted by α which is set to 1 by cross val. 25 | Args: 26 | c: class confidences, 27 | l: predicted boxes, 28 | g: ground truth boxes 29 | N: number of matched default boxes 30 | See: https://arxiv.org/pdf/1512.02325.pdf for more details. 31 | """ 32 | 33 | def __init__(self, num_classes, overlap_thresh, prior_for_matching, 34 | bkg_label, neg_mining, neg_pos, neg_overlap, encode_target, 35 | use_gpu=True): 36 | super(MultiBoxLoss, self).__init__() 37 | self.use_gpu = use_gpu 38 | self.num_classes = num_classes 39 | self.threshold = overlap_thresh 40 | self.background_label = bkg_label 41 | self.encode_target = encode_target 42 | self.use_prior_for_matching = prior_for_matching 43 | self.do_neg_mining = neg_mining 44 | self.negpos_ratio = neg_pos 45 | self.neg_overlap = neg_overlap 46 | self.variance = cfg['variance'] 47 | 48 | def forward(self, predictions, targets): 49 | """Multibox Loss 50 | Args: 51 | predictions (tuple): A tuple containing loc preds, conf preds, 52 | and prior boxes from SSD net. 53 | conf shape: torch.size(batch_size,num_priors,num_classes) 54 | loc shape: torch.size(batch_size,num_priors,4) 55 | priors shape: torch.size(num_priors,4) 56 | 57 | targets (tensor): Ground truth boxes and labels for a batch, 58 | shape: [batch_size,num_objs,5] (last idx is the label). 59 | """ 60 | loc_data, conf_data, priors = predictions 61 | num = loc_data.size(0) 62 | priors = priors[:loc_data.size(1), :] 63 | num_priors = (priors.size(0)) 64 | num_classes = self.num_classes 65 | print(loc_data.size(), conf_data.size(), priors.size()) 66 | 67 | # match priors (default boxes) and ground truth boxes 68 | loc_t = torch.Tensor(num, num_priors, 4) 69 | conf_t = torch.LongTensor(num, num_priors) 70 | for idx in range(num): 71 | truths = targets[idx][:, :-1].data 72 | labels = targets[idx][:, -1].data 73 | defaults = priors.data 74 | match(self.threshold, truths, defaults, self.variance, labels, 75 | loc_t, conf_t, idx) 76 | if self.use_gpu: 77 | loc_t = loc_t.cuda() 78 | conf_t = conf_t.cuda() 79 | # wrap targets 80 | #loc_t = Variable(loc_t, requires_grad=False) 81 | #conf_t = Variable(conf_t, requires_grad=False) 82 | loc_t.requires_grad = False 83 | conf_t.requires_grad = False 84 | print(loc_t.size(), conf_t.size()) 85 | 86 | pos = conf_t > 0 87 | print(pos.size()) 88 | #num_pos = pos.sum(dim=1, keepdim=True) 89 | 90 | # Localization Loss (Smooth L1) 91 | # Shape: [batch,num_priors,4] 92 | pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data) 93 | loc_p = loc_data[pos_idx].view(-1, 4) 94 | loc_t = loc_t[pos_idx].view(-1, 4) 95 | loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum') 96 | 97 | # Compute max conf across batch for hard negative mining 98 | batch_conf = conf_data.view(-1, self.num_classes) 99 | loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1)) 100 | print(loss_c.size()) 101 | 102 | # Hard Negative Mining 103 | loss_c[pos.view(-1,1)] = 0 # filter out pos boxes for now 104 | loss_c = loss_c.view(num, -1) 105 | _, loss_idx = loss_c.sort(1, descending=True) 106 | _, idx_rank = loss_idx.sort(1) 107 | num_pos = pos.long().sum(1, keepdim=True) 108 | num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1) 109 | neg = idx_rank < num_neg.expand_as(idx_rank) 110 | print(num_pos.size(), num_neg.size(), neg.size()) 111 | 112 | # Confidence Loss Including Positive and Negative Examples 113 | pos_idx = pos.unsqueeze(2).expand_as(conf_data) 114 | neg_idx = neg.unsqueeze(2).expand_as(conf_data) 115 | conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes) 116 | targets_weighted = conf_t[(pos+neg).gt(0)] 117 | print(pos_idx.size(), neg_idx.size(), conf_p.size(), targets_weighted.size()) 118 | loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum') 119 | 120 | # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 121 | 122 | N = num_pos.data.sum().float() 123 | #N = max(num_pos.data.sum().float(), 1) 124 | loss_l /= N 125 | loss_c /= N 126 | #print(N, loss_l, loss_c) 127 | return loss_l, loss_c 128 | -------------------------------------------------------------------------------- /layers/modules/refinedet_multibox_loss.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.autograd import Variable 6 | from data import coco as cfg 7 | from ..box_utils import match, log_sum_exp, refine_match 8 | 9 | 10 | class RefineDetMultiBoxLoss(nn.Module): 11 | """SSD Weighted Loss Function 12 | Compute Targets: 13 | 1) Produce Confidence Target Indices by matching ground truth boxes 14 | with (default) 'priorboxes' that have jaccard index > threshold parameter 15 | (default threshold: 0.5). 16 | 2) Produce localization target by 'encoding' variance into offsets of ground 17 | truth boxes and their matched 'priorboxes'. 18 | 3) Hard negative mining to filter the excessive number of negative examples 19 | that comes with using a large number of default bounding boxes. 20 | (default negative:positive ratio 3:1) 21 | Objective Loss: 22 | L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 23 | Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss 24 | weighted by α which is set to 1 by cross val. 25 | Args: 26 | c: class confidences, 27 | l: predicted boxes, 28 | g: ground truth boxes 29 | N: number of matched default boxes 30 | See: https://arxiv.org/pdf/1512.02325.pdf for more details. 31 | """ 32 | 33 | def __init__(self, num_classes, overlap_thresh, prior_for_matching, 34 | bkg_label, neg_mining, neg_pos, neg_overlap, encode_target, 35 | use_gpu=True, theta=0.01, use_ARM=False): 36 | super(RefineDetMultiBoxLoss, self).__init__() 37 | self.use_gpu = use_gpu 38 | self.num_classes = num_classes 39 | self.threshold = overlap_thresh 40 | self.background_label = bkg_label 41 | self.encode_target = encode_target 42 | self.use_prior_for_matching = prior_for_matching 43 | self.do_neg_mining = neg_mining 44 | self.negpos_ratio = neg_pos 45 | self.neg_overlap = neg_overlap 46 | self.variance = cfg['variance'] 47 | self.theta = theta 48 | self.use_ARM = use_ARM 49 | 50 | def forward(self, predictions, targets): 51 | """Multibox Loss 52 | Args: 53 | predictions (tuple): A tuple containing loc preds, conf preds, 54 | and prior boxes from SSD net. 55 | conf shape: torch.size(batch_size,num_priors,num_classes) 56 | loc shape: torch.size(batch_size,num_priors,4) 57 | priors shape: torch.size(num_priors,4) 58 | 59 | targets (tensor): Ground truth boxes and labels for a batch, 60 | shape: [batch_size,num_objs,5] (last idx is the label). 61 | """ 62 | arm_loc_data, arm_conf_data, odm_loc_data, odm_conf_data, priors = predictions 63 | #print(arm_loc_data.size(), arm_conf_data.size(), 64 | # odm_loc_data.size(), odm_conf_data.size(), priors.size()) 65 | #input() 66 | if self.use_ARM: 67 | loc_data, conf_data = odm_loc_data, odm_conf_data 68 | else: 69 | loc_data, conf_data = arm_loc_data, arm_conf_data 70 | num = loc_data.size(0) 71 | priors = priors[:loc_data.size(1), :] 72 | num_priors = (priors.size(0)) 73 | num_classes = self.num_classes 74 | #print(loc_data.size(), conf_data.size(), priors.size()) 75 | 76 | # match priors (default boxes) and ground truth boxes 77 | loc_t = torch.Tensor(num, num_priors, 4) 78 | conf_t = torch.LongTensor(num, num_priors) 79 | for idx in range(num): 80 | truths = targets[idx][:, :-1].data 81 | labels = targets[idx][:, -1].data 82 | if num_classes == 2: 83 | labels = labels >= 0 84 | defaults = priors.data 85 | if self.use_ARM: 86 | refine_match(self.threshold, truths, defaults, self.variance, labels, 87 | loc_t, conf_t, idx, arm_loc_data[idx].data) 88 | else: 89 | refine_match(self.threshold, truths, defaults, self.variance, labels, 90 | loc_t, conf_t, idx) 91 | if self.use_gpu: 92 | loc_t = loc_t.cuda() 93 | conf_t = conf_t.cuda() 94 | # wrap targets 95 | #loc_t = Variable(loc_t, requires_grad=False) 96 | #conf_t = Variable(conf_t, requires_grad=False) 97 | loc_t.requires_grad = False 98 | conf_t.requires_grad = False 99 | #print(loc_t.size(), conf_t.size()) 100 | 101 | if self.use_ARM: 102 | P = F.softmax(arm_conf_data, 2) 103 | arm_conf_tmp = P[:,:,1] 104 | object_score_index = arm_conf_tmp <= self.theta 105 | pos = conf_t > 0 106 | pos[object_score_index.data] = 0 107 | else: 108 | pos = conf_t > 0 109 | #print(pos.size()) 110 | #num_pos = pos.sum(dim=1, keepdim=True) 111 | 112 | # Localization Loss (Smooth L1) 113 | # Shape: [batch,num_priors,4] 114 | pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data) 115 | loc_p = loc_data[pos_idx].view(-1, 4) 116 | loc_t = loc_t[pos_idx].view(-1, 4) 117 | loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum') 118 | 119 | # Compute max conf across batch for hard negative mining 120 | batch_conf = conf_data.view(-1, self.num_classes) 121 | loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1)) 122 | #print(loss_c.size()) 123 | 124 | # Hard Negative Mining 125 | loss_c[pos.view(-1,1)] = 0 # filter out pos boxes for now 126 | loss_c = loss_c.view(num, -1) 127 | _, loss_idx = loss_c.sort(1, descending=True) 128 | _, idx_rank = loss_idx.sort(1) 129 | num_pos = pos.long().sum(1, keepdim=True) 130 | num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1) 131 | neg = idx_rank < num_neg.expand_as(idx_rank) 132 | #print(num_pos.size(), num_neg.size(), neg.size()) 133 | 134 | # Confidence Loss Including Positive and Negative Examples 135 | pos_idx = pos.unsqueeze(2).expand_as(conf_data) 136 | neg_idx = neg.unsqueeze(2).expand_as(conf_data) 137 | conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes) 138 | targets_weighted = conf_t[(pos+neg).gt(0)] 139 | #print(pos_idx.size(), neg_idx.size(), conf_p.size(), targets_weighted.size()) 140 | loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum') 141 | 142 | # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 143 | 144 | N = num_pos.data.sum().float() 145 | #N = max(num_pos.data.sum().float(), 1) 146 | loss_l /= N 147 | loss_c /= N 148 | #print(N, loss_l, loss_c) 149 | return loss_l, loss_c 150 | -------------------------------------------------------------------------------- /models/refinedet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from layers import * 6 | from data import voc_refinedet, coco_refinedet 7 | import os 8 | 9 | 10 | class RefineDet(nn.Module): 11 | """Single Shot Multibox Architecture 12 | The network is composed of a base VGG network followed by the 13 | added multibox conv layers. Each multibox layer branches into 14 | 1) conv2d for class conf scores 15 | 2) conv2d for localization predictions 16 | 3) associated priorbox layer to produce default bounding 17 | boxes specific to the layer's feature map size. 18 | See: https://arxiv.org/pdf/1512.02325.pdf for more details. 19 | 20 | Args: 21 | phase: (string) Can be "test" or "train" 22 | size: input image size 23 | base: VGG16 layers for input, size of either 300 or 500 24 | extras: extra layers that feed to multibox loc and conf layers 25 | head: "multibox head" consists of loc and conf conv layers 26 | """ 27 | 28 | def __init__(self, phase, size, base, extras, ARM, ODM, TCB, num_classes): 29 | super(RefineDet, self).__init__() 30 | self.phase = phase 31 | self.num_classes = num_classes 32 | self.cfg = (coco_refinedet, voc_refinedet)[num_classes == 21] 33 | self.priorbox = PriorBox(self.cfg[str(size)]) 34 | with torch.no_grad(): 35 | self.priors = self.priorbox.forward() 36 | self.size = size 37 | 38 | # SSD network 39 | self.vgg = nn.ModuleList(base) 40 | # Layer learns to scale the l2 normalized features from conv4_3 41 | self.conv4_3_L2Norm = L2Norm(512, 10) 42 | self.conv5_3_L2Norm = L2Norm(512, 8) 43 | self.extras = nn.ModuleList(extras) 44 | 45 | self.arm_loc = nn.ModuleList(ARM[0]) 46 | self.arm_conf = nn.ModuleList(ARM[1]) 47 | self.odm_loc = nn.ModuleList(ODM[0]) 48 | self.odm_conf = nn.ModuleList(ODM[1]) 49 | #self.tcb = nn.ModuleList(TCB) 50 | self.tcb0 = nn.ModuleList(TCB[0]) 51 | self.tcb1 = nn.ModuleList(TCB[1]) 52 | self.tcb2 = nn.ModuleList(TCB[2]) 53 | 54 | if phase == 'test': 55 | self.softmax = nn.Softmax(dim=-1) 56 | self.detect = Detect_RefineDet(num_classes, self.size, 0, 1000, 0.01, 0.45, 0.01, 500) 57 | 58 | def forward(self, x): 59 | """Applies network layers and ops on input image(s) x. 60 | 61 | Args: 62 | x: input image or batch of images. Shape: [batch,3,300,300]. 63 | 64 | Return: 65 | Depending on phase: 66 | test: 67 | Variable(tensor) of output class label predictions, 68 | confidence score, and corresponding location predictions for 69 | each object detected. Shape: [batch,topk,7] 70 | 71 | train: 72 | list of concat outputs from: 73 | 1: confidence layers, Shape: [batch*num_priors,num_classes] 74 | 2: localization layers, Shape: [batch,num_priors*4] 75 | 3: priorbox layers, Shape: [2,num_priors*4] 76 | """ 77 | sources = list() 78 | tcb_source = list() 79 | arm_loc = list() 80 | arm_conf = list() 81 | odm_loc = list() 82 | odm_conf = list() 83 | 84 | # apply vgg up to conv4_3 relu and conv5_3 relu 85 | for k in range(30): 86 | x = self.vgg[k](x) 87 | if 22 == k: 88 | s = self.conv4_3_L2Norm(x) 89 | sources.append(s) 90 | elif 29 == k: 91 | s = self.conv5_3_L2Norm(x) 92 | sources.append(s) 93 | 94 | # apply vgg up to fc7 95 | for k in range(30, len(self.vgg)): 96 | x = self.vgg[k](x) 97 | sources.append(x) 98 | 99 | # apply extra layers and cache source layer outputs 100 | for k, v in enumerate(self.extras): 101 | x = F.relu(v(x), inplace=True) 102 | if k % 2 == 1: 103 | sources.append(x) 104 | 105 | # apply ARM and ODM to source layers 106 | for (x, l, c) in zip(sources, self.arm_loc, self.arm_conf): 107 | arm_loc.append(l(x).permute(0, 2, 3, 1).contiguous()) 108 | arm_conf.append(c(x).permute(0, 2, 3, 1).contiguous()) 109 | arm_loc = torch.cat([o.view(o.size(0), -1) for o in arm_loc], 1) 110 | arm_conf = torch.cat([o.view(o.size(0), -1) for o in arm_conf], 1) 111 | #print([x.size() for x in sources]) 112 | # calculate TCB features 113 | #print([x.size() for x in sources]) 114 | p = None 115 | for k, v in enumerate(sources[::-1]): 116 | s = v 117 | for i in range(3): 118 | s = self.tcb0[(3-k)*3 + i](s) 119 | #print(s.size()) 120 | if k != 0: 121 | u = p 122 | u = self.tcb1[3-k](u) 123 | s += u 124 | for i in range(3): 125 | s = self.tcb2[(3-k)*3 + i](s) 126 | p = s 127 | tcb_source.append(s) 128 | #print([x.size() for x in tcb_source]) 129 | tcb_source.reverse() 130 | 131 | # apply ODM to source layers 132 | for (x, l, c) in zip(tcb_source, self.odm_loc, self.odm_conf): 133 | odm_loc.append(l(x).permute(0, 2, 3, 1).contiguous()) 134 | odm_conf.append(c(x).permute(0, 2, 3, 1).contiguous()) 135 | odm_loc = torch.cat([o.view(o.size(0), -1) for o in odm_loc], 1) 136 | odm_conf = torch.cat([o.view(o.size(0), -1) for o in odm_conf], 1) 137 | #print(arm_loc.size(), arm_conf.size(), odm_loc.size(), odm_conf.size()) 138 | 139 | if self.phase == "test": 140 | #print(loc, conf) 141 | output = self.detect( 142 | arm_loc.view(arm_loc.size(0), -1, 4), # arm loc preds 143 | self.softmax(arm_conf.view(arm_conf.size(0), -1, 144 | 2)), # arm conf preds 145 | odm_loc.view(odm_loc.size(0), -1, 4), # odm loc preds 146 | self.softmax(odm_conf.view(odm_conf.size(0), -1, 147 | self.num_classes)), # odm conf preds 148 | self.priors.type(type(x.data)) # default boxes 149 | ) 150 | else: 151 | output = ( 152 | arm_loc.view(arm_loc.size(0), -1, 4), 153 | arm_conf.view(arm_conf.size(0), -1, 2), 154 | odm_loc.view(odm_loc.size(0), -1, 4), 155 | odm_conf.view(odm_conf.size(0), -1, self.num_classes), 156 | self.priors 157 | ) 158 | return output 159 | 160 | def load_weights(self, base_file): 161 | other, ext = os.path.splitext(base_file) 162 | if ext == '.pkl' or '.pth': 163 | print('Loading weights into state dict...') 164 | self.load_state_dict(torch.load(base_file, 165 | map_location=lambda storage, loc: storage)) 166 | print('Finished!') 167 | else: 168 | print('Sorry only .pth and .pkl files supported.') 169 | 170 | 171 | # This function is derived from torchvision VGG make_layers() 172 | # https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py 173 | def vgg(cfg, i, batch_norm=False): 174 | layers = [] 175 | in_channels = i 176 | for v in cfg: 177 | if v == 'M': 178 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 179 | elif v == 'C': 180 | layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)] 181 | else: 182 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 183 | if batch_norm: 184 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 185 | else: 186 | layers += [conv2d, nn.ReLU(inplace=True)] 187 | in_channels = v 188 | pool5 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) 189 | conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=3, dilation=3) 190 | conv7 = nn.Conv2d(1024, 1024, kernel_size=1) 191 | layers += [pool5, conv6, 192 | nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)] 193 | return layers 194 | 195 | 196 | def add_extras(cfg, size, i, batch_norm=False): 197 | # Extra layers added to VGG for feature scaling 198 | layers = [] 199 | in_channels = i 200 | flag = False 201 | for k, v in enumerate(cfg): 202 | if in_channels != 'S': 203 | if v == 'S': 204 | layers += [nn.Conv2d(in_channels, cfg[k + 1], 205 | kernel_size=(1, 3)[flag], stride=2, padding=1)] 206 | else: 207 | layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])] 208 | flag = not flag 209 | in_channels = v 210 | return layers 211 | 212 | def arm_multibox(vgg, extra_layers, cfg): 213 | arm_loc_layers = [] 214 | arm_conf_layers = [] 215 | vgg_source = [21, 28, -2] 216 | for k, v in enumerate(vgg_source): 217 | arm_loc_layers += [nn.Conv2d(vgg[v].out_channels, 218 | cfg[k] * 4, kernel_size=3, padding=1)] 219 | arm_conf_layers += [nn.Conv2d(vgg[v].out_channels, 220 | cfg[k] * 2, kernel_size=3, padding=1)] 221 | for k, v in enumerate(extra_layers[1::2], 3): 222 | arm_loc_layers += [nn.Conv2d(v.out_channels, cfg[k] 223 | * 4, kernel_size=3, padding=1)] 224 | arm_conf_layers += [nn.Conv2d(v.out_channels, cfg[k] 225 | * 2, kernel_size=3, padding=1)] 226 | return (arm_loc_layers, arm_conf_layers) 227 | 228 | def odm_multibox(vgg, extra_layers, cfg, num_classes): 229 | odm_loc_layers = [] 230 | odm_conf_layers = [] 231 | vgg_source = [21, 28, -2] 232 | for k, v in enumerate(vgg_source): 233 | odm_loc_layers += [nn.Conv2d(256, cfg[k] * 4, kernel_size=3, padding=1)] 234 | odm_conf_layers += [nn.Conv2d(256, cfg[k] * num_classes, kernel_size=3, padding=1)] 235 | for k, v in enumerate(extra_layers[1::2], 3): 236 | odm_loc_layers += [nn.Conv2d(256, cfg[k] * 4, kernel_size=3, padding=1)] 237 | odm_conf_layers += [nn.Conv2d(256, cfg[k] * num_classes, kernel_size=3, padding=1)] 238 | return (odm_loc_layers, odm_conf_layers) 239 | 240 | def add_tcb(cfg): 241 | feature_scale_layers = [] 242 | feature_upsample_layers = [] 243 | feature_pred_layers = [] 244 | for k, v in enumerate(cfg): 245 | feature_scale_layers += [nn.Conv2d(cfg[k], 256, 3, padding=1), 246 | nn.ReLU(inplace=True), 247 | nn.Conv2d(256, 256, 3, padding=1) 248 | ] 249 | feature_pred_layers += [nn.ReLU(inplace=True), 250 | nn.Conv2d(256, 256, 3, padding=1), 251 | nn.ReLU(inplace=True) 252 | ] 253 | if k != len(cfg) - 1: 254 | feature_upsample_layers += [nn.ConvTranspose2d(256, 256, 2, 2)] 255 | return (feature_scale_layers, feature_upsample_layers, feature_pred_layers) 256 | 257 | base = { 258 | '320': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 259 | 512, 512, 512], 260 | '512': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 261 | 512, 512, 512], 262 | } 263 | extras = { 264 | '320': [256, 'S', 512], 265 | '512': [256, 'S', 512], 266 | } 267 | mbox = { 268 | '320': [3, 3, 3, 3], # number of boxes per feature map location 269 | '512': [3, 3, 3, 3], # number of boxes per feature map location 270 | } 271 | 272 | tcb = { 273 | '320': [512, 512, 1024, 512], 274 | '512': [512, 512, 1024, 512], 275 | } 276 | 277 | 278 | def build_refinedet(phase, size=320, num_classes=21): 279 | if phase != "test" and phase != "train": 280 | print("ERROR: Phase: " + phase + " not recognized") 281 | return 282 | if size != 320 and size != 512: 283 | print("ERROR: You specified size " + repr(size) + ". However, " + 284 | "currently only RefineDet320 and RefineDet512 is supported!") 285 | return 286 | base_ = vgg(base[str(size)], 3) 287 | extras_ = add_extras(extras[str(size)], size, 1024) 288 | ARM_ = arm_multibox(base_, extras_, mbox[str(size)]) 289 | ODM_ = odm_multibox(base_, extras_, mbox[str(size)], num_classes) 290 | TCB_ = add_tcb(tcb[str(size)]) 291 | return RefineDet(phase, size, base_, extras_, ARM_, ODM_, TCB_, num_classes) 292 | -------------------------------------------------------------------------------- /train_refinedet.py: -------------------------------------------------------------------------------- 1 | from data import * 2 | from utils.augmentations import SSDAugmentation 3 | from layers.modules import RefineDetMultiBoxLoss 4 | #from ssd import build_ssd 5 | from models.refinedet import build_refinedet 6 | import os 7 | import sys 8 | import time 9 | import torch 10 | import torch.nn as nn 11 | import torch.optim as optim 12 | import torch.backends.cudnn as cudnn 13 | import torch.nn.init as init 14 | import torch.utils.data as data 15 | import numpy as np 16 | import argparse 17 | from utils.logging import Logger 18 | 19 | def str2bool(v): 20 | return v.lower() in ("yes", "true", "t", "1") 21 | 22 | 23 | parser = argparse.ArgumentParser( 24 | description='Single Shot MultiBox Detector Training With Pytorch') 25 | train_set = parser.add_mutually_exclusive_group() 26 | parser.add_argument('--dataset', default='VOC', choices=['VOC', 'COCO'], 27 | type=str, help='VOC or COCO') 28 | parser.add_argument('--input_size', default='320', choices=['320', '512'], 29 | type=str, help='RefineDet320 or RefineDet512') 30 | parser.add_argument('--dataset_root', default=VOC_ROOT, 31 | help='Dataset root directory path') 32 | parser.add_argument('--basenet', default='./weights/vgg16_reducedfc.pth', 33 | help='Pretrained base model') 34 | parser.add_argument('--batch_size', default=32, type=int, 35 | help='Batch size for training') 36 | parser.add_argument('--resume', default=None, type=str, 37 | help='Checkpoint state_dict file to resume training from') 38 | parser.add_argument('--start_iter', default=0, type=int, 39 | help='Resume training at this iter') 40 | parser.add_argument('--num_workers', default=8, type=int, 41 | help='Number of workers used in dataloading') 42 | parser.add_argument('--cuda', default=True, type=str2bool, 43 | help='Use CUDA to train model') 44 | parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float, 45 | help='initial learning rate') 46 | parser.add_argument('--momentum', default=0.9, type=float, 47 | help='Momentum value for optim') 48 | parser.add_argument('--weight_decay', default=5e-4, type=float, 49 | help='Weight decay for SGD') 50 | parser.add_argument('--gamma', default=0.1, type=float, 51 | help='Gamma update for SGD') 52 | parser.add_argument('--visdom', default=False, type=str2bool, 53 | help='Use visdom for loss visualization') 54 | parser.add_argument('--save_folder', default='weights/', 55 | help='Directory for saving checkpoint models') 56 | args = parser.parse_args() 57 | 58 | 59 | if torch.cuda.is_available(): 60 | if args.cuda: 61 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 62 | if not args.cuda: 63 | print("WARNING: It looks like you have a CUDA device, but aren't " + 64 | "using CUDA.\nRun with --cuda for optimal training speed.") 65 | torch.set_default_tensor_type('torch.FloatTensor') 66 | else: 67 | torch.set_default_tensor_type('torch.FloatTensor') 68 | 69 | if not os.path.exists(args.save_folder): 70 | os.mkdir(args.save_folder) 71 | 72 | sys.stdout = Logger(os.path.join(args.save_folder, 'log.txt')) 73 | 74 | def train(): 75 | if args.dataset == 'COCO': 76 | '''if args.dataset_root == VOC_ROOT: 77 | if not os.path.exists(COCO_ROOT): 78 | parser.error('Must specify dataset_root if specifying dataset') 79 | print("WARNING: Using default COCO dataset_root because " + 80 | "--dataset_root was not specified.") 81 | args.dataset_root = COCO_ROOT 82 | cfg = coco 83 | dataset = COCODetection(root=args.dataset_root, 84 | transform=SSDAugmentation(cfg['min_dim'], 85 | MEANS))''' 86 | elif args.dataset == 'VOC': 87 | '''if args.dataset_root == COCO_ROOT: 88 | parser.error('Must specify dataset if specifying dataset_root')''' 89 | cfg = voc_refinedet[args.input_size] 90 | dataset = VOCDetection(root=args.dataset_root, 91 | transform=SSDAugmentation(cfg['min_dim'], 92 | MEANS)) 93 | 94 | if args.visdom: 95 | import visdom 96 | viz = visdom.Visdom() 97 | 98 | refinedet_net = build_refinedet('train', cfg['min_dim'], cfg['num_classes']) 99 | net = refinedet_net 100 | print(net) 101 | #input() 102 | 103 | if args.cuda: 104 | net = torch.nn.DataParallel(refinedet_net) 105 | cudnn.benchmark = True 106 | 107 | if args.resume: 108 | print('Resuming training, loading {}...'.format(args.resume)) 109 | refinedet_net.load_weights(args.resume) 110 | else: 111 | #vgg_weights = torch.load(args.save_folder + args.basenet) 112 | vgg_weights = torch.load(args.basenet) 113 | print('Loading base network...') 114 | refinedet_net.vgg.load_state_dict(vgg_weights) 115 | 116 | if args.cuda: 117 | net = net.cuda() 118 | 119 | if not args.resume: 120 | print('Initializing weights...') 121 | # initialize newly added layers' weights with xavier method 122 | refinedet_net.extras.apply(weights_init) 123 | refinedet_net.arm_loc.apply(weights_init) 124 | refinedet_net.arm_conf.apply(weights_init) 125 | refinedet_net.odm_loc.apply(weights_init) 126 | refinedet_net.odm_conf.apply(weights_init) 127 | #refinedet_net.tcb.apply(weights_init) 128 | refinedet_net.tcb0.apply(weights_init) 129 | refinedet_net.tcb1.apply(weights_init) 130 | refinedet_net.tcb2.apply(weights_init) 131 | 132 | optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, 133 | weight_decay=args.weight_decay) 134 | arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, 135 | False, args.cuda) 136 | odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, 137 | False, args.cuda, use_ARM=True) 138 | 139 | net.train() 140 | # loss counters 141 | arm_loc_loss = 0 142 | arm_conf_loss = 0 143 | odm_loc_loss = 0 144 | odm_conf_loss = 0 145 | epoch = 0 146 | print('Loading the dataset...') 147 | 148 | epoch_size = len(dataset) // args.batch_size 149 | print('Training RefineDet on:', dataset.name) 150 | print('Using the specified args:') 151 | print(args) 152 | 153 | step_index = 0 154 | 155 | if args.visdom: 156 | vis_title = 'RefineDet.PyTorch on ' + dataset.name 157 | vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss'] 158 | iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend) 159 | epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend) 160 | 161 | data_loader = data.DataLoader(dataset, args.batch_size, 162 | num_workers=args.num_workers, 163 | shuffle=True, collate_fn=detection_collate, 164 | pin_memory=True) 165 | # create batch iterator 166 | batch_iterator = iter(data_loader) 167 | for iteration in range(args.start_iter, cfg['max_iter']): 168 | if args.visdom and iteration != 0 and (iteration % epoch_size == 0): 169 | update_vis_plot(epoch, arm_loc_loss, arm_conf_loss, epoch_plot, None, 170 | 'append', epoch_size) 171 | # reset epoch loss counters 172 | arm_loc_loss = 0 173 | arm_conf_loss = 0 174 | odm_loc_loss = 0 175 | odm_conf_loss = 0 176 | epoch += 1 177 | 178 | if iteration in cfg['lr_steps']: 179 | step_index += 1 180 | adjust_learning_rate(optimizer, args.gamma, step_index) 181 | 182 | # load train data 183 | try: 184 | images, targets = next(batch_iterator) 185 | except StopIteration: 186 | batch_iterator = iter(data_loader) 187 | images, targets = next(batch_iterator) 188 | 189 | if args.cuda: 190 | images = images.cuda() 191 | targets = [ann.cuda() for ann in targets] 192 | else: 193 | images = images 194 | targets = [ann for ann in targets] 195 | # forward 196 | t0 = time.time() 197 | out = net(images) 198 | # backprop 199 | optimizer.zero_grad() 200 | arm_loss_l, arm_loss_c = arm_criterion(out, targets) 201 | odm_loss_l, odm_loss_c = odm_criterion(out, targets) 202 | #input() 203 | arm_loss = arm_loss_l + arm_loss_c 204 | odm_loss = odm_loss_l + odm_loss_c 205 | loss = arm_loss + odm_loss 206 | loss.backward() 207 | optimizer.step() 208 | t1 = time.time() 209 | arm_loc_loss += arm_loss_l.item() 210 | arm_conf_loss += arm_loss_c.item() 211 | odm_loc_loss += odm_loss_l.item() 212 | odm_conf_loss += odm_loss_c.item() 213 | 214 | if iteration % 10 == 0: 215 | print('timer: %.4f sec.' % (t1 - t0)) 216 | print('iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \ 217 | % (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ') 218 | 219 | if args.visdom: 220 | update_vis_plot(iteration, arm_loss_l.data[0], arm_loss_c.data[0], 221 | iter_plot, epoch_plot, 'append') 222 | 223 | if iteration != 0 and iteration % 5000 == 0: 224 | print('Saving state, iter:', iteration) 225 | torch.save(refinedet_net.state_dict(), args.save_folder 226 | + '/RefineDet{}_{}_{}.pth'.format(args.input_size, args.dataset, 227 | repr(iteration))) 228 | torch.save(refinedet_net.state_dict(), args.save_folder 229 | + '/RefineDet{}_{}_final.pth'.format(args.input_size, args.dataset)) 230 | 231 | 232 | def adjust_learning_rate(optimizer, gamma, step): 233 | """Sets the learning rate to the initial LR decayed by 10 at every 234 | specified step 235 | # Adapted from PyTorch Imagenet example: 236 | # https://github.com/pytorch/examples/blob/master/imagenet/main.py 237 | """ 238 | lr = args.lr * (gamma ** (step)) 239 | for param_group in optimizer.param_groups: 240 | param_group['lr'] = lr 241 | 242 | 243 | def xavier(param): 244 | init.xavier_uniform_(param) 245 | 246 | 247 | def weights_init(m): 248 | if isinstance(m, nn.Conv2d): 249 | xavier(m.weight.data) 250 | m.bias.data.zero_() 251 | elif isinstance(m, nn.ConvTranspose2d): 252 | xavier(m.weight.data) 253 | m.bias.data.zero_() 254 | 255 | 256 | def create_vis_plot(_xlabel, _ylabel, _title, _legend): 257 | return viz.line( 258 | X=torch.zeros((1,)).cpu(), 259 | Y=torch.zeros((1, 3)).cpu(), 260 | opts=dict( 261 | xlabel=_xlabel, 262 | ylabel=_ylabel, 263 | title=_title, 264 | legend=_legend 265 | ) 266 | ) 267 | 268 | 269 | def update_vis_plot(iteration, loc, conf, window1, window2, update_type, 270 | epoch_size=1): 271 | viz.line( 272 | X=torch.ones((1, 3)).cpu() * iteration, 273 | Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu() / epoch_size, 274 | win=window1, 275 | update=update_type 276 | ) 277 | # initialize epoch plot on first iteration 278 | if iteration == 0: 279 | viz.line( 280 | X=torch.zeros((1, 3)).cpu(), 281 | Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu(), 282 | win=window2, 283 | update=True 284 | ) 285 | 286 | 287 | if __name__ == '__main__': 288 | train() 289 | -------------------------------------------------------------------------------- /train_refinedet320.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train_refinedet.py --save_folder weights/RefineDet320/ --input_size 320 2 | -------------------------------------------------------------------------------- /train_refinedet512.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train_refinedet.py --save_folder weights/RefineDet512/ --input_size 512 2 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .augmentations import SSDAugmentation -------------------------------------------------------------------------------- /utils/augmentations.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import transforms 3 | import cv2 4 | import numpy as np 5 | import types 6 | from numpy import random 7 | 8 | 9 | def intersect(box_a, box_b): 10 | max_xy = np.minimum(box_a[:, 2:], box_b[2:]) 11 | min_xy = np.maximum(box_a[:, :2], box_b[:2]) 12 | inter = np.clip((max_xy - min_xy), a_min=0, a_max=np.inf) 13 | return inter[:, 0] * inter[:, 1] 14 | 15 | 16 | def jaccard_numpy(box_a, box_b): 17 | """Compute the jaccard overlap of two sets of boxes. The jaccard overlap 18 | is simply the intersection over union of two boxes. 19 | E.g.: 20 | A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) 21 | Args: 22 | box_a: Multiple bounding boxes, Shape: [num_boxes,4] 23 | box_b: Single bounding box, Shape: [4] 24 | Return: 25 | jaccard overlap: Shape: [box_a.shape[0], box_a.shape[1]] 26 | """ 27 | inter = intersect(box_a, box_b) 28 | area_a = ((box_a[:, 2]-box_a[:, 0]) * 29 | (box_a[:, 3]-box_a[:, 1])) # [A,B] 30 | area_b = ((box_b[2]-box_b[0]) * 31 | (box_b[3]-box_b[1])) # [A,B] 32 | union = area_a + area_b - inter 33 | return inter / union # [A,B] 34 | 35 | 36 | class Compose(object): 37 | """Composes several augmentations together. 38 | Args: 39 | transforms (List[Transform]): list of transforms to compose. 40 | Example: 41 | >>> augmentations.Compose([ 42 | >>> transforms.CenterCrop(10), 43 | >>> transforms.ToTensor(), 44 | >>> ]) 45 | """ 46 | 47 | def __init__(self, transforms): 48 | self.transforms = transforms 49 | 50 | def __call__(self, img, boxes=None, labels=None): 51 | for t in self.transforms: 52 | img, boxes, labels = t(img, boxes, labels) 53 | return img, boxes, labels 54 | 55 | 56 | class Lambda(object): 57 | """Applies a lambda as a transform.""" 58 | 59 | def __init__(self, lambd): 60 | assert isinstance(lambd, types.LambdaType) 61 | self.lambd = lambd 62 | 63 | def __call__(self, img, boxes=None, labels=None): 64 | return self.lambd(img, boxes, labels) 65 | 66 | 67 | class ConvertFromInts(object): 68 | def __call__(self, image, boxes=None, labels=None): 69 | return image.astype(np.float32), boxes, labels 70 | 71 | 72 | class SubtractMeans(object): 73 | def __init__(self, mean): 74 | self.mean = np.array(mean, dtype=np.float32) 75 | 76 | def __call__(self, image, boxes=None, labels=None): 77 | image = image.astype(np.float32) 78 | image -= self.mean 79 | return image.astype(np.float32), boxes, labels 80 | 81 | 82 | class ToAbsoluteCoords(object): 83 | def __call__(self, image, boxes=None, labels=None): 84 | height, width, channels = image.shape 85 | boxes[:, 0] *= width 86 | boxes[:, 2] *= width 87 | boxes[:, 1] *= height 88 | boxes[:, 3] *= height 89 | 90 | return image, boxes, labels 91 | 92 | 93 | class ToPercentCoords(object): 94 | def __call__(self, image, boxes=None, labels=None): 95 | height, width, channels = image.shape 96 | boxes[:, 0] /= width 97 | boxes[:, 2] /= width 98 | boxes[:, 1] /= height 99 | boxes[:, 3] /= height 100 | 101 | return image, boxes, labels 102 | 103 | 104 | class Resize(object): 105 | def __init__(self, size=300): 106 | self.size = size 107 | 108 | def __call__(self, image, boxes=None, labels=None): 109 | image = cv2.resize(image, (self.size, 110 | self.size)) 111 | return image, boxes, labels 112 | 113 | 114 | class RandomSaturation(object): 115 | def __init__(self, lower=0.5, upper=1.5): 116 | self.lower = lower 117 | self.upper = upper 118 | assert self.upper >= self.lower, "contrast upper must be >= lower." 119 | assert self.lower >= 0, "contrast lower must be non-negative." 120 | 121 | def __call__(self, image, boxes=None, labels=None): 122 | if random.randint(2): 123 | image[:, :, 1] *= random.uniform(self.lower, self.upper) 124 | 125 | return image, boxes, labels 126 | 127 | 128 | class RandomHue(object): 129 | def __init__(self, delta=18.0): 130 | assert delta >= 0.0 and delta <= 360.0 131 | self.delta = delta 132 | 133 | def __call__(self, image, boxes=None, labels=None): 134 | if random.randint(2): 135 | image[:, :, 0] += random.uniform(-self.delta, self.delta) 136 | image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0 137 | image[:, :, 0][image[:, :, 0] < 0.0] += 360.0 138 | return image, boxes, labels 139 | 140 | 141 | class RandomLightingNoise(object): 142 | def __init__(self): 143 | self.perms = ((0, 1, 2), (0, 2, 1), 144 | (1, 0, 2), (1, 2, 0), 145 | (2, 0, 1), (2, 1, 0)) 146 | 147 | def __call__(self, image, boxes=None, labels=None): 148 | if random.randint(2): 149 | swap = self.perms[random.randint(len(self.perms))] 150 | shuffle = SwapChannels(swap) # shuffle channels 151 | image = shuffle(image) 152 | return image, boxes, labels 153 | 154 | 155 | class ConvertColor(object): 156 | def __init__(self, current='BGR', transform='HSV'): 157 | self.transform = transform 158 | self.current = current 159 | 160 | def __call__(self, image, boxes=None, labels=None): 161 | if self.current == 'BGR' and self.transform == 'HSV': 162 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 163 | elif self.current == 'HSV' and self.transform == 'BGR': 164 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 165 | else: 166 | raise NotImplementedError 167 | return image, boxes, labels 168 | 169 | 170 | class RandomContrast(object): 171 | def __init__(self, lower=0.5, upper=1.5): 172 | self.lower = lower 173 | self.upper = upper 174 | assert self.upper >= self.lower, "contrast upper must be >= lower." 175 | assert self.lower >= 0, "contrast lower must be non-negative." 176 | 177 | # expects float image 178 | def __call__(self, image, boxes=None, labels=None): 179 | if random.randint(2): 180 | alpha = random.uniform(self.lower, self.upper) 181 | image *= alpha 182 | return image, boxes, labels 183 | 184 | 185 | class RandomBrightness(object): 186 | def __init__(self, delta=32): 187 | assert delta >= 0.0 188 | assert delta <= 255.0 189 | self.delta = delta 190 | 191 | def __call__(self, image, boxes=None, labels=None): 192 | if random.randint(2): 193 | delta = random.uniform(-self.delta, self.delta) 194 | image += delta 195 | return image, boxes, labels 196 | 197 | 198 | class ToCV2Image(object): 199 | def __call__(self, tensor, boxes=None, labels=None): 200 | return tensor.cpu().numpy().astype(np.float32).transpose((1, 2, 0)), boxes, labels 201 | 202 | 203 | class ToTensor(object): 204 | def __call__(self, cvimage, boxes=None, labels=None): 205 | return torch.from_numpy(cvimage.astype(np.float32)).permute(2, 0, 1), boxes, labels 206 | 207 | 208 | class RandomSampleCrop(object): 209 | """Crop 210 | Arguments: 211 | img (Image): the image being input during training 212 | boxes (Tensor): the original bounding boxes in pt form 213 | labels (Tensor): the class labels for each bbox 214 | mode (float tuple): the min and max jaccard overlaps 215 | Return: 216 | (img, boxes, classes) 217 | img (Image): the cropped image 218 | boxes (Tensor): the adjusted bounding boxes in pt form 219 | labels (Tensor): the class labels for each bbox 220 | """ 221 | def __init__(self): 222 | self.sample_options = ( 223 | # using entire original input image 224 | None, 225 | # sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9 226 | (0.1, None), 227 | (0.3, None), 228 | (0.7, None), 229 | (0.9, None), 230 | # randomly sample a patch 231 | (None, None), 232 | ) 233 | 234 | def __call__(self, image, boxes=None, labels=None): 235 | height, width, _ = image.shape 236 | while True: 237 | # randomly choose a mode 238 | mode = random.choice(self.sample_options) 239 | if mode is None: 240 | return image, boxes, labels 241 | 242 | min_iou, max_iou = mode 243 | if min_iou is None: 244 | min_iou = float('-inf') 245 | if max_iou is None: 246 | max_iou = float('inf') 247 | 248 | # max trails (50) 249 | for _ in range(50): 250 | current_image = image 251 | 252 | w = random.uniform(0.3 * width, width) 253 | h = random.uniform(0.3 * height, height) 254 | 255 | # aspect ratio constraint b/t .5 & 2 256 | if h / w < 0.5 or h / w > 2: 257 | continue 258 | 259 | left = random.uniform(width - w) 260 | top = random.uniform(height - h) 261 | 262 | # convert to integer rect x1,y1,x2,y2 263 | rect = np.array([int(left), int(top), int(left+w), int(top+h)]) 264 | 265 | # calculate IoU (jaccard overlap) b/t the cropped and gt boxes 266 | overlap = jaccard_numpy(boxes, rect) 267 | 268 | # is min and max overlap constraint satisfied? if not try again 269 | if overlap.min() < min_iou and max_iou < overlap.max(): 270 | continue 271 | 272 | # cut the crop from the image 273 | current_image = current_image[rect[1]:rect[3], rect[0]:rect[2], 274 | :] 275 | 276 | # keep overlap with gt box IF center in sampled patch 277 | centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0 278 | 279 | # mask in all gt boxes that above and to the left of centers 280 | m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1]) 281 | 282 | # mask in all gt boxes that under and to the right of centers 283 | m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1]) 284 | 285 | # mask in that both m1 and m2 are true 286 | mask = m1 * m2 287 | 288 | # have any valid boxes? try again if not 289 | if not mask.any(): 290 | continue 291 | 292 | # take only matching gt boxes 293 | current_boxes = boxes[mask, :].copy() 294 | 295 | # take only matching gt labels 296 | current_labels = labels[mask] 297 | 298 | # should we use the box left and top corner or the crop's 299 | current_boxes[:, :2] = np.maximum(current_boxes[:, :2], 300 | rect[:2]) 301 | # adjust to crop (by substracting crop's left,top) 302 | current_boxes[:, :2] -= rect[:2] 303 | 304 | current_boxes[:, 2:] = np.minimum(current_boxes[:, 2:], 305 | rect[2:]) 306 | # adjust to crop (by substracting crop's left,top) 307 | current_boxes[:, 2:] -= rect[:2] 308 | 309 | return current_image, current_boxes, current_labels 310 | 311 | 312 | class Expand(object): 313 | def __init__(self, mean): 314 | self.mean = mean 315 | 316 | def __call__(self, image, boxes, labels): 317 | if random.randint(2): 318 | return image, boxes, labels 319 | 320 | height, width, depth = image.shape 321 | ratio = random.uniform(1, 4) 322 | left = random.uniform(0, width*ratio - width) 323 | top = random.uniform(0, height*ratio - height) 324 | 325 | expand_image = np.zeros( 326 | (int(height*ratio), int(width*ratio), depth), 327 | dtype=image.dtype) 328 | expand_image[:, :, :] = self.mean 329 | expand_image[int(top):int(top + height), 330 | int(left):int(left + width)] = image 331 | image = expand_image 332 | 333 | boxes = boxes.copy() 334 | boxes[:, :2] += (int(left), int(top)) 335 | boxes[:, 2:] += (int(left), int(top)) 336 | 337 | return image, boxes, labels 338 | 339 | 340 | class RandomMirror(object): 341 | def __call__(self, image, boxes, classes): 342 | _, width, _ = image.shape 343 | if random.randint(2): 344 | image = image[:, ::-1] 345 | boxes = boxes.copy() 346 | boxes[:, 0::2] = width - boxes[:, 2::-2] 347 | return image, boxes, classes 348 | 349 | 350 | class SwapChannels(object): 351 | """Transforms a tensorized image by swapping the channels in the order 352 | specified in the swap tuple. 353 | Args: 354 | swaps (int triple): final order of channels 355 | eg: (2, 1, 0) 356 | """ 357 | 358 | def __init__(self, swaps): 359 | self.swaps = swaps 360 | 361 | def __call__(self, image): 362 | """ 363 | Args: 364 | image (Tensor): image tensor to be transformed 365 | Return: 366 | a tensor with channels swapped according to swap 367 | """ 368 | # if torch.is_tensor(image): 369 | # image = image.data.cpu().numpy() 370 | # else: 371 | # image = np.array(image) 372 | image = image[:, :, self.swaps] 373 | return image 374 | 375 | 376 | class PhotometricDistort(object): 377 | def __init__(self): 378 | self.pd = [ 379 | RandomContrast(), 380 | ConvertColor(transform='HSV'), 381 | RandomSaturation(), 382 | RandomHue(), 383 | ConvertColor(current='HSV', transform='BGR'), 384 | RandomContrast() 385 | ] 386 | self.rand_brightness = RandomBrightness() 387 | self.rand_light_noise = RandomLightingNoise() 388 | 389 | def __call__(self, image, boxes, labels): 390 | im = image.copy() 391 | im, boxes, labels = self.rand_brightness(im, boxes, labels) 392 | if random.randint(2): 393 | distort = Compose(self.pd[:-1]) 394 | else: 395 | distort = Compose(self.pd[1:]) 396 | im, boxes, labels = distort(im, boxes, labels) 397 | return self.rand_light_noise(im, boxes, labels) 398 | 399 | 400 | class SSDAugmentation(object): 401 | def __init__(self, size=300, mean=(104, 117, 123)): 402 | self.mean = mean 403 | self.size = size 404 | self.augment = Compose([ 405 | ConvertFromInts(), 406 | ToAbsoluteCoords(), 407 | PhotometricDistort(), 408 | Expand(self.mean), 409 | RandomSampleCrop(), 410 | RandomMirror(), 411 | ToPercentCoords(), 412 | Resize(self.size), 413 | SubtractMeans(self.mean) 414 | ]) 415 | 416 | def __call__(self, img, boxes, labels): 417 | return self.augment(img, boxes, labels) 418 | -------------------------------------------------------------------------------- /utils/logging.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import os 3 | import sys 4 | 5 | from .osutils import mkdir_if_missing 6 | 7 | 8 | class Logger(object): 9 | def __init__(self, fpath=None): 10 | self.console = sys.stdout 11 | self.file = None 12 | if fpath is not None: 13 | mkdir_if_missing(os.path.dirname(fpath)) 14 | self.file = open(fpath, 'w') 15 | 16 | def __del__(self): 17 | self.close() 18 | 19 | def __enter__(self): 20 | pass 21 | 22 | def __exit__(self, *args): 23 | self.close() 24 | 25 | def write(self, msg): 26 | self.console.write(msg) 27 | if self.file is not None: 28 | self.file.write(msg) 29 | 30 | def flush(self): 31 | self.console.flush() 32 | if self.file is not None: 33 | self.file.flush() 34 | os.fsync(self.file.fileno()) 35 | 36 | def close(self): 37 | self.console.close() 38 | if self.file is not None: 39 | self.file.close() 40 | -------------------------------------------------------------------------------- /utils/osutils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import os 3 | import errno 4 | 5 | 6 | def mkdir_if_missing(dir_path): 7 | try: 8 | os.makedirs(dir_path) 9 | except OSError as e: 10 | if e.errno != errno.EEXIST: 11 | raise 12 | --------------------------------------------------------------------------------