├── MANIFEST.in ├── README.md ├── __pycache__ └── save_json.cpython-36.pyc ├── classify_rule.json ├── config.json ├── configs ├── efficient_net_b3_ssd300_voc0712.yaml ├── mobilenet_v2_ssd320_voc0712.yaml ├── resnet50_224.yaml ├── resnet50_300.yaml ├── resnet50_512.yaml ├── vgg_ssd224_voc0712.yaml ├── vgg_ssd300_coco_trainval35k.yaml ├── vgg_ssd300_voc0712.yaml ├── vgg_ssd512_coco_trainval35k.yaml └── vgg_ssd512_voc0712.yaml ├── customize_service.py ├── data_aug ├── README.md ├── color.py ├── mirror.py ├── padd_rotated_crop.py └── rotated.py ├── demo.py ├── draw_rectangle.py ├── ext ├── __init__.py ├── build.py ├── cpu │ ├── nms_cpu.cpp │ └── vision.h ├── cuda │ ├── nms.cu │ └── vision.h ├── nms.h └── vision.cpp ├── figures ├── 004545.jpg ├── losses.png ├── lr.png └── metrics.png ├── get_ratio.py ├── get_test.py ├── infer.py ├── make_imagesets.py ├── save_json.py ├── setup.py ├── ssd ├── __init__.py ├── __pycache__ │ └── __init__.cpython-36.pyc ├── config │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── defaults.cpython-36.pyc │ │ └── path_catlog.cpython-36.pyc │ ├── defaults.py │ └── path_catlog.py ├── data │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── build.cpython-36.pyc │ ├── build.py │ ├── datasets │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── coco.cpython-36.pyc │ │ │ └── voc.cpython-36.pyc │ │ ├── coco.py │ │ ├── evaluation │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ │ └── __init__.cpython-36.pyc │ │ │ ├── coco │ │ │ │ ├── __init__.py │ │ │ │ └── __pycache__ │ │ │ │ │ └── __init__.cpython-36.pyc │ │ │ └── voc │ │ │ │ ├── __init__.py │ │ │ │ ├── __pycache__ │ │ │ │ ├── __init__.cpython-36.pyc │ │ │ │ └── eval_detection_voc.cpython-36.pyc │ │ │ │ └── eval_detection_voc.py │ │ └── voc.py │ ├── samplers │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── distributed.cpython-36.pyc │ │ │ └── iteration_based_batch_sampler.cpython-36.pyc │ │ ├── distributed.py │ │ └── iteration_based_batch_sampler.py │ └── transforms │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── target_transform.cpython-36.pyc │ │ └── transforms.cpython-36.pyc │ │ ├── target_transform.py │ │ └── transforms.py ├── engine │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── inference.cpython-36.pyc │ │ └── trainer.cpython-36.pyc │ ├── inference.py │ └── trainer.py ├── layers │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── separable_conv.cpython-36.pyc │ └── separable_conv.py ├── modeling │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── registry.cpython-36.pyc │ ├── anchors │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── prior_box.cpython-36.pyc │ │ └── prior_box.py │ ├── backbone │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── mobilenet.cpython-36.pyc │ │ │ ├── resnet.cpython-36.pyc │ │ │ └── vgg.cpython-36.pyc │ │ ├── efficient_net │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ │ ├── __init__.cpython-36.pyc │ │ │ │ ├── efficient_net.cpython-36.pyc │ │ │ │ └── utils.cpython-36.pyc │ │ │ ├── efficient_net.py │ │ │ └── utils.py │ │ ├── mobilenet.py │ │ ├── resnet.py │ │ ├── utils.py │ │ └── vgg.py │ ├── box │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-35.pyc │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── box_utils.cpython-35.pyc │ │ │ ├── box_utils.cpython-36.pyc │ │ │ ├── prior_box.cpython-35.pyc │ │ │ └── prior_box.cpython-36.pyc │ │ ├── box_utils.py │ │ └── prior_box.py │ ├── box_head │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── box_head.cpython-36.pyc │ │ │ ├── box_predictor.cpython-36.pyc │ │ │ ├── inference.cpython-36.pyc │ │ │ ├── loss.cpython-36.pyc │ │ │ └── loss_Iou.cpython-36.pyc │ │ ├── box_head.py │ │ ├── box_predictor.py │ │ ├── box_utils.py │ │ ├── inference.py │ │ ├── loss.py │ │ └── loss_Iou.py │ ├── detector │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── ssd_detector.cpython-36.pyc │ │ └── ssd_detector.py │ └── registry.py ├── solver │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── build.cpython-36.pyc │ │ └── lr_scheduler.cpython-36.pyc │ ├── build.py │ └── lr_scheduler.py ├── structures │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── container.cpython-36.pyc │ └── container.py └── utils │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── box_utils.cpython-36.pyc │ ├── checkpoint.cpython-36.pyc │ ├── dist_util.cpython-36.pyc │ ├── logger.cpython-36.pyc │ ├── metric_logger.cpython-36.pyc │ ├── misc.cpython-36.pyc │ ├── model_zoo.cpython-36.pyc │ ├── nms.cpython-36.pyc │ └── registry.cpython-36.pyc │ ├── box_utils.py │ ├── checkpoint.py │ ├── dist_util.py │ ├── logger.py │ ├── metric_logger.py │ ├── misc.py │ ├── model_zoo.py │ ├── nms.py │ └── registry.py ├── test.py └── train.py /MANIFEST.in: -------------------------------------------------------------------------------- 1 | recursive-include configs *.yaml 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GarbageDetection 2 | “华为云杯”2020深圳开放数据应用创新大赛·生活垃圾图片分类.排名:50/4388;方案:SSD-efficientd3-DiouLoss 3 | 4 | # “华为云杯”2020深圳开放数据应用创新大赛·生活垃圾图片分类. 5 | 6 | - 排名:50/4388; 7 | - 方案:SSD-efficientd3-DiouLoss 8 | 本次比赛由于需要在modelart上面线上部署，而且有推理时间限制，因此我采取的是单阶段模型SSD.主干网络efficientNetd3/ResNet50. 9 | loss 由smoothL1更改为了CiouLoss. 10 | 11 | ## 对仓库代码的说明 12 | 13 | 所用的SSD是pytoch版本，来源于该仓库：[SSD](https://github.com/lufficc/SSD). 14 | 15 | 针对本次垃圾检测分类我所做的更改如下： 16 | 17 | - BackBone的替换:增加了ResNet系列(torchvision的官方实现)，提取了resnet四层feature 18 | - 关键在于提取resnet的四层feature,得到每个特征图的channel大小,对应输入尺度的相应尺寸，写好相关配置文件 19 | - loss的替换:由SmoothL1替换为IOU/Diou/Giou/Ciou-loss系列.分类loss有采取过focalLoss,但表现差强人意 20 | - 在SSD中，网络的输出是偏置，原来仓库的实现计算的偏置的SmoothL1 21 | - 算iouloss需要四个点的坐标.因此需要将网络的输出(偏置)利用先验框解码为x1,y1,x2,y2 22 | ### ssd/modeling/backbone/resnet.py 23 | 24 | 这个是新增加的以resnet作为banckbone的脚本(提取出4层featureMaps) 25 | 26 | ### ssd/modeling/box_head/loss_iou.py 27 | 28 | iou/giou/diou/ciou的实现 29 | 30 | ### make_imagesets.py 31 | 32 | 分割voc的训练集/验证集/测试集合 33 | 34 | ### customize_service.py 35 | 36 | 华为云modelart需要的推理脚本 37 | 38 | ### get_ratio.py 39 | 40 | 计算类别比例 41 | -------------------------------------------------------------------------------- /__pycache__/save_json.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/__pycache__/save_json.cpython-36.pyc -------------------------------------------------------------------------------- /classify_rule.json: -------------------------------------------------------------------------------- 1 | { 2 | "可回收物": [ 3 | "充电宝", 4 | "包", 5 | "洗护用品", 6 | "塑料玩具", 7 | "塑料器皿", 8 | "塑料衣架", 9 | "玻璃器皿", 10 | "金属器皿", 11 | "快递纸袋", 12 | "插头电线", 13 | "旧衣服", 14 | "易拉罐", 15 | "枕头", 16 | "毛绒玩具", 17 | "鞋", 18 | "砧板", 19 | "纸盒纸箱", 20 | "调料瓶", 21 | "酒瓶", 22 | "金属食品罐", 23 | "金属厨具", 24 | "锅", 25 | "食用油桶", 26 | "饮料瓶", 27 | "书籍纸张", 28 | "垃圾桶" 29 | ], 30 | "厨余垃圾": [ 31 | "剩饭剩菜", 32 | "大骨头", 33 | "果皮果肉", 34 | "茶叶渣", 35 | "菜帮菜叶", 36 | "蛋壳", 37 | "鱼骨" 38 | ], 39 | "有害垃圾": [ 40 | "干电池", 41 | "软膏", 42 | "过期药物" 43 | ], 44 | "其他垃圾": [ 45 | "一次性快餐盒", 46 | "污损塑料", 47 | "烟蒂", 48 | "牙签", 49 | "花盆", 50 | "陶瓷器皿", 51 | "筷子", 52 | "污损用纸" 53 | ] 54 | } -------------------------------------------------------------------------------- /config.json: -------------------------------------------------------------------------------- 1 | { 2 | 3 | "model_algorithm": "object_detection", 4 | "model_type": "PyTorch", 5 | "runtime": "python3.6", 6 | "apis": [ 7 | { 8 | "procotol": "http", 9 | "url": "/", 10 | "method": "post", 11 | 12 | "request": { 13 | "Content-type": "multipart/form-data", 14 | "data": { 15 | "type": "object", 16 | "properties": { 17 | "images": { 18 | "type": "file" 19 | } 20 | } 21 | } 22 | }, 23 | 24 | 25 | "response": { 26 | "Content-type": "multipart/form-data", 27 | "data": { 28 | "type": "object", 29 | "properties": { 30 | "detection_classes": { 31 | "type": "list", 32 | "items": [{ 33 | "type": "string" 34 | }] 35 | }, 36 | "detection_scores": { 37 | "type": "list", 38 | "items": [{ 39 | "type": "number" 40 | }] 41 | }, 42 | "detection_boxes": { 43 | "type": "list", 44 | "items": [{ 45 | "type": "list", 46 | "minItems": 4, 47 | "maxItems": 4, 48 | "items": [{ 49 | "type": "number" 50 | }] 51 | }] 52 | }//, 53 | // "latency_time": { 54 | // "type": "string" 55 | // 56 | // } 57 | } 58 | } 59 | } 60 | 61 | } 62 | 63 | ], 64 | 65 | 66 | "metrics": { 67 | "f1": 0.0, 68 | "recall": 0.0, 69 | "precision": 0.0, 70 | "accuracy": 0.0 71 | }, 72 | 73 | // "dependencies": [{ 74 | // "installer": "pip", 75 | // "packages": [ 76 | // { 77 | // "restraint": "EXACT", 78 | // "package_version": "1.4.0", 79 | // "package_name": "torch" 80 | // }, 81 | // { 82 | // "restraint": "EXACT", 83 | // "package_version": "0.5.0", 84 | // "package_name": "torchvision" 85 | // } 86 | // ] 87 | // }] 88 | 89 | 90 | } 91 | -------------------------------------------------------------------------------- /configs/efficient_net_b3_ssd300_voc0712.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 45 3 | BACKBONE: 4 | NAME: 'efficient_net-b3' 5 | OUT_CHANNELS: (48, 136, 384, 256, 256, 256) 6 | INPUT: 7 | IMAGE_SIZE: 300 8 | DATASETS: 9 | TRAIN: ("voc_2007_trainval", ) 10 | TEST: ("voc_2007_test", ) 11 | SOLVER: 12 | MAX_ITER: 140000 13 | LR_STEPS: [80000, 100000] 14 | GAMMA: 0.1 15 | BATCH_SIZE: 8 16 | LR: 1e-3 17 | 18 | OUTPUT_DIR: 'outputs/efficient_net_b3_ssd300_voc0712' -------------------------------------------------------------------------------- /configs/mobilenet_v2_ssd320_voc0712.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 21 3 | BOX_HEAD: 4 | PREDICTOR: 'SSDLiteBoxPredictor' 5 | BACKBONE: 6 | NAME: 'mobilenet_v2' 7 | OUT_CHANNELS: (96, 1280, 512, 256, 256, 64) 8 | PRIORS: 9 | FEATURE_MAPS: [20, 10, 5, 3, 2, 1] 10 | STRIDES: [16, 32, 64, 107, 160, 320] 11 | MIN_SIZES: [60, 105, 150, 195, 240, 285] 12 | MAX_SIZES: [105, 150, 195, 240, 285, 330] 13 | ASPECT_RATIOS: [[2, 3], [2, 3], [2, 3], [2, 3], [2, 3], [2, 3]] 14 | BOXES_PER_LOCATION: [6, 6, 6, 6, 6, 6] 15 | INPUT: 16 | IMAGE_SIZE: 320 17 | DATASETS: 18 | TRAIN: ("voc_2007_trainval", "voc_2012_trainval") 19 | TEST: ("voc_2007_test", ) 20 | SOLVER: 21 | MAX_ITER: 120000 22 | LR_STEPS: [80000, 100000] 23 | GAMMA: 0.1 24 | BATCH_SIZE: 32 25 | LR: 1e-3 26 | 27 | OUTPUT_DIR: 'outputs/mobilenet_v2_ssd320_voc0712' -------------------------------------------------------------------------------- /configs/resnet50_224.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | BACKBONE: 3 | NAME: 'res_backbone' 4 | OUT_CHANNELS: (256, 512, 1024, 2048) # should match feature1 - feature4's out_channels in MyBackbone 5 | PRIORS: 6 | FEATURE_MAPS: [56, 28, 14, 7] # feature1 - feature4's size 7 | STRIDES: [4, 8, 16, 30] # feature1 - feature4's output stride 8 | MIN_SIZES: [21, 45, 99, 153] # your custom anchor settings 9 | MAX_SIZES: [45, 99, 153, 224] 10 | ASPECT_RATIOS: [[2, 3], [3, 4], [3, 2], [1, 1]] 11 | BOXES_PER_LOCATION: [6, 6, 6, 6] 12 | INPUT: 13 | IMAGE_SIZE: 224 14 | DATASETS: 15 | TRAIN: ("voc_2007_trainval", ) 16 | TEST: ("voc_2007_test", ) 17 | SOLVER: 18 | MAX_ITER: 120000 19 | LR_STEPS: [80000, 100000] 20 | GAMMA: 0.1 21 | BATCH_SIZE: 32 22 | LR: 1e-3 23 | 24 | OUTPUT_DIR: 'outputs/resnet50_ssd224' -------------------------------------------------------------------------------- /configs/resnet50_300.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | BACKBONE: 3 | NAME: 'res_backbone' 4 | OUT_CHANNELS: (256, 512, 1024, 2048) # should match feature1 - feature4's out_channels in MyBackbone 5 | PRIORS: 6 | FEATURE_MAPS: [75, 38, 19, 10] # feature1 - feature4's size 7 | STRIDES: [4, 8, 16, 30] # feature1 - feature4's output stride 8 | MIN_SIZES: [36, 100, 159, 253] # your custom anchor settings 9 | MAX_SIZES: [100, 159, 253, 300] 10 | ASPECT_RATIOS: [[2, 3], [4, 3], [3, 2], [1, 1]] 11 | BOXES_PER_LOCATION: [6, 6, 6, 6] 12 | INPUT: 13 | IMAGE_SIZE: 300 14 | DATASETS: 15 | TRAIN: ("voc_2007_trainval", ) 16 | TEST: ("voc_2007_test", ) 17 | SOLVER: 18 | MAX_ITER: 120000 19 | LR_STEPS: [10000,50000, 60000] 20 | GAMMA: 0.1 21 | BATCH_SIZE: 8 22 | LR: 1e-3 23 | 24 | OUTPUT_DIR: 'outputs/resnet101_ssd300_voc0712' -------------------------------------------------------------------------------- /configs/resnet50_512.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | BACKBONE: 3 | NAME: 'res_backbone' 4 | OUT_CHANNELS: (256, 512, 1024, 2048) # should match feature1 - feature4's out_channels in MyBackbone 5 | PRIORS: 6 | FEATURE_MAPS: [128, 64, 32, 16] # feature1 - feature4's size 7 | STRIDES: [4, 8, 16, 30] # feature1 - feature4's output stride 8 | MIN_SIZES: [50, 102, 226, 349] # your custom anchor settings 9 | MAX_SIZES: [102, 226, 349, 512] 10 | ASPECT_RATIOS: [[2, 3], [3, 4], [3, 2], [1, 1]] 11 | BOXES_PER_LOCATION: [6, 6, 6, 6] 12 | INPUT: 13 | IMAGE_SIZE: 512 14 | DATASETS: 15 | TRAIN: ("voc_2007_trainval", ) 16 | TEST: ("voc_2007_test", ) 17 | SOLVER: 18 | MAX_ITER: 120000 19 | LR_STEPS: [80000, 100000] 20 | GAMMA: 0.1 21 | BATCH_SIZE: 4 22 | LR: 1e-3 23 | 24 | OUTPUT_DIR: 'outputs/resnet50_ssd512' -------------------------------------------------------------------------------- /configs/vgg_ssd224_voc0712.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 45 3 | INPUT: 4 | IMAGE_SIZE: 300 5 | DATASETS: 6 | TRAIN: ("voc_2007_trainval", ) 7 | TEST: ("voc_2007_test", ) 8 | SOLVER: 9 | MAX_ITER: 120000 10 | LR_STEPS: [80000, 100000] 11 | GAMMA: 0.1 12 | BATCH_SIZE: 8 13 | LR: 1e-3 14 | 15 | OUTPUT_DIR: 'outputs/vgg_ssd300_voc0712' -------------------------------------------------------------------------------- /configs/vgg_ssd300_coco_trainval35k.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 81 3 | PRIORS: 4 | FEATURE_MAPS: [38, 19, 10, 5, 3, 1] 5 | STRIDES: [8, 16, 32, 64, 100, 300] 6 | MIN_SIZES: [21, 45, 99, 153, 207, 261] 7 | MAX_SIZES: [45, 99, 153, 207, 261, 315] 8 | ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2], [2]] 9 | BOXES_PER_LOCATION: [4, 6, 6, 6, 4, 4] 10 | INPUT: 11 | IMAGE_SIZE: 300 12 | DATASETS: 13 | TRAIN: ("coco_2014_train", "coco_2014_valminusminival") 14 | TEST: ("coco_2014_minival", ) 15 | SOLVER: 16 | MAX_ITER: 400000 17 | LR_STEPS: [280000, 360000] 18 | GAMMA: 0.1 19 | BATCH_SIZE: 32 20 | LR: 1e-3 21 | 22 | OUTPUT_DIR: 'outputs/vgg_ssd300_coco_trainval35k' -------------------------------------------------------------------------------- /configs/vgg_ssd300_voc0712.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 45 3 | INPUT: 4 | IMAGE_SIZE: 300 5 | DATASETS: 6 | TRAIN: ("voc_2007_trainval", ) 7 | TEST: ("voc_2007_test", ) 8 | SOLVER: 9 | MAX_ITER: 120000 10 | LR_STEPS: [80000, 100000] 11 | GAMMA: 0.1 12 | BATCH_SIZE: 8 13 | LR: 1e-3 14 | 15 | OUTPUT_DIR: 'outputs/vgg_ssd300_voc0712' -------------------------------------------------------------------------------- /configs/vgg_ssd512_coco_trainval35k.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 45 3 | BACKBONE: 4 | OUT_CHANNELS: (512, 1024, 512, 256, 256, 256, 256) 5 | PRIORS: 6 | FEATURE_MAPS: [64, 32, 16, 8, 4, 2, 1] 7 | STRIDES: [8, 16, 32, 64, 128, 256, 512] 8 | MIN_SIZES: [20.48, 51.2, 133.12, 215.04, 296.96, 378.88, 460.8] 9 | MAX_SIZES: [51.2, 133.12, 215.04, 296.96, 378.88, 460.8, 542.72] 10 | ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]] 11 | BOXES_PER_LOCATION: [4, 6, 6, 6, 6, 4, 4] 12 | INPUT: 13 | IMAGE_SIZE: 512 14 | DATASETS: 15 | TRAIN: ("voc_2007_trainval", ) 16 | TEST: ("voc_2007_test", ) 17 | SOLVER: 18 | MAX_ITER: 150000 19 | LR_STEPS: [100000, 120000] 20 | GAMMA: 0.1 21 | BATCH_SIZE: 4 22 | LR: 1e-3 23 | 24 | OUTPUT_DIR: 'outputs/vgg_ssd512_coco_trainval35k' -------------------------------------------------------------------------------- /configs/vgg_ssd512_voc0712.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | NUM_CLASSES: 45 3 | BACKBONE: 4 | OUT_CHANNELS: (512, 1024, 512, 256, 256, 256, 256) 5 | PRIORS: 6 | FEATURE_MAPS: [64, 32, 16, 8, 4, 2, 1] 7 | STRIDES: [8, 16, 32, 64, 128, 256, 512] 8 | MIN_SIZES: [35.84, 76.8, 153.6, 230.4, 307.2, 384.0, 460.8] 9 | MAX_SIZES: [76.8, 153.6, 230.4, 307.2, 384.0, 460.8, 537.65] 10 | ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]] 11 | BOXES_PER_LOCATION: [4, 6, 6, 6, 6, 4, 4] 12 | INPUT: 13 | IMAGE_SIZE: 512 14 | DATASETS: 15 | TRAIN: ("voc_2007_trainval", ) 16 | TEST: ("voc_2007_test", ) 17 | SOLVER: 18 | MAX_ITER: 120000 19 | LR_STEPS: [80000, 100000] 20 | GAMMA: 0.1 21 | BATCH_SIZE: 8 22 | LR: 1e-3 23 | 24 | OUTPUT_DIR: 'outputs/vgg_ssd512_voc0712' -------------------------------------------------------------------------------- /customize_service.py: -------------------------------------------------------------------------------- 1 | # -*- encoding: utf-8 -*- 2 | from PIL import Image 3 | import logging as log 4 | from model_service.pytorch_model_service import PTServingBaseService 5 | from metric.metrics_manager import MetricsManager 6 | import torch.nn.functional as F 7 | 8 | import torch.nn as nn 9 | import torch 10 | import json 11 | import numpy as np 12 | import torchvision 13 | import time 14 | import os 15 | import copy 16 | 17 | import sys 18 | import cv2 19 | from ssd.config import cfg 20 | from ssd.modeling.detector import build_detection_model 21 | import torch 22 | from torch import nn 23 | 24 | import codecs 25 | sys.stdout = codecs.getwriter("utf-8")(sys.stdout.detach()) 26 | 27 | from torch.utils.data import Dataset, DataLoader 28 | from torchvision import datasets, models, transforms 29 | 30 | from save_json import label_list, create_class_dict, get_classes_name, save_result_as_json 31 | 32 | print('CUDA available: {}'.format(torch.cuda.is_available())) 33 | 34 | 35 | logger = log.getLogger(__name__) 36 | 37 | IMAGES_KEY = 'images' 38 | MODEL_INPUT_KEY = 'images' 39 | 40 | def Net(model_path): 41 | cfg.merge_from_file('configs/efficient_net_b3_ssd300_voc0712.yaml') 42 | model = build_detection_model(cfg) 43 | state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['model'] 44 | model.load_state_dict(state_dict) 45 | #model.eval() 46 | return model 47 | class Resize(object): 48 | def __init__(self, size=300): 49 | self.size = size 50 | 51 | def __call__(self, image, boxes=None, labels=None): 52 | image = cv2.resize(image, (self.size, 53 | self.size)) 54 | return image, boxes, labels 55 | class SubtractMeans(object): 56 | def __init__(self, mean): 57 | self.mean = np.array(mean, dtype=np.float32) 58 | 59 | def __call__(self, image, boxes=None, labels=None): 60 | image = image.astype(np.float32) 61 | image -= self.mean 62 | return image.astype(np.float32), boxes, labels 63 | class ToTensor(object): 64 | def __call__(self, cvimage, boxes=None, labels=None): 65 | return torch.from_numpy(cvimage.astype(np.float32)).permute(2, 0, 1), boxes, labels 66 | def remove_empty_boxes(boxes, labels): 67 | """Removes bounding boxes of W or H equal to 0 and its labels 68 | 69 | Args: 70 | boxes (ndarray): NP Array with bounding boxes as lines 71 | * BBOX[x1, y1, x2, y2] 72 | labels (labels): Corresponding labels with boxes 73 | 74 | Returns: 75 | ndarray: Valid bounding boxes 76 | ndarray: Corresponding labels 77 | """ 78 | del_boxes = [] 79 | for idx, box in enumerate(boxes): 80 | if box[0] == box[2] or box[1] == box[3]: 81 | del_boxes.append(idx) 82 | 83 | return np.delete(boxes, del_boxes, 0), np.delete(labels, del_boxes) 84 | class Compose(object): 85 | """Composes several augmentations together. 86 | Args: 87 | transforms (List[Transform]): list of transforms to compose. 88 | Example: 89 | >>> augmentations.Compose([ 90 | >>> transforms.CenterCrop(10), 91 | >>> transforms.ToTensor(), 92 | >>> ]) 93 | """ 94 | 95 | def __init__(self, transforms): 96 | self.transforms = transforms 97 | 98 | def __call__(self, img, boxes=None, labels=None): 99 | for t in self.transforms: 100 | img, boxes, labels = t(img, boxes, labels) 101 | if boxes is not None: 102 | boxes, labels = remove_empty_boxes(boxes, labels) 103 | return img, boxes, labels 104 | def decode_image(file_content): 105 | """ 106 | Decode bytes to a single image 107 | :param file_content: bytes 108 | :return: ndarray with rank=3 109 | """ 110 | image = Image.open(file_content) 111 | image = image.convert('RGB') 112 | # print(image.shape) 113 | image = np.array(image) 114 | return image 115 | 116 | def build_transforms(): 117 | transform = [ 118 | Resize(300), 119 | SubtractMeans([123, 117, 104]), 120 | ToTensor() 121 | ] 122 | transform = Compose(transform) 123 | return transform 124 | 125 | 126 | 127 | class PTVisionService(PTServingBaseService): 128 | 129 | def __init__(self, model_path): 130 | super(PTVisionService, self).__init__(model_path) 131 | 132 | self.dir_path = os.path.dirname(os.path.realpath(model_path)) 133 | # load label name 134 | self.label = ('__background__','一次性快餐盒','书籍纸张','充电宝', 135 | '剩饭剩菜' ,'包','垃圾桶','塑料器皿','塑料玩具','塑料衣架', 136 | '大骨头','干电池','快递纸袋','插头电线','旧衣服','易拉罐', 137 | '枕头','果皮果肉','毛绒玩具','污损塑料','污损用纸','洗护用品', 138 | '烟蒂','牙签','玻璃器皿','砧板','筷子','纸盒纸箱','花盆', 139 | '茶叶渣','菜帮菜叶','蛋壳','调料瓶','软膏','过期药物', 140 | '酒瓶','金属厨具','金属器皿','金属食品罐','锅','陶瓷器皿', 141 | '鞋','食用油桶','饮料瓶','鱼骨') 142 | self.classify={'充电宝': '可回收物', '包': '可回收物', '洗护用品': '可回收物', '塑料玩具': '可回收物', 143 | '塑料器皿': '可回收物', '塑料衣架': '可回收物', '玻璃器皿': '可回收物', '金属器皿': '可回收物', 144 | '快递纸袋': '可回收物', '插头电线': '可回收物', '旧衣服': '可回收物', '易拉罐': '可回收物', 145 | '枕头': '可回收物', '毛绒玩具': '可回收物', '鞋': '可回收物', '砧板': '可回收物', '纸盒纸箱': '可回收物', '调料瓶': '可回收物', 146 | '酒瓶': '可回收物', '金属食品罐': '可回收物', '金属厨具': '可回收物', '锅': '可回收物', 147 | '食用油桶': '可回收物', '饮料瓶': '可回收物', '书籍纸张': '可回收物', '垃圾桶': '可回收物', 148 | '剩饭剩菜': '厨余垃圾', '大骨头': '厨余垃圾', '果皮果肉': '厨余垃圾', '茶叶渣': '厨余垃圾', 149 | '菜帮菜叶': '厨余垃圾', '蛋壳': '厨余垃圾', '鱼骨': '厨余垃圾', 150 | '干电池': '有害垃圾', '软膏': '有害垃圾', '过期药物': '有害垃圾', 151 | '一次性快餐盒': '其他垃圾', '污损塑料': '其他垃圾', '烟蒂': '其他垃圾', '牙签': '其他垃圾', 152 | '花盆': '其他垃圾', '陶瓷器皿': '其他垃圾', '筷子': '其他垃圾', '污损用纸': '其他垃圾' 153 | } 154 | self.num_class = len(self.label) 155 | self.transform = build_transforms() 156 | self.score_threshold=0.5 157 | 158 | # Load your model 159 | self.model = Net(model_path) 160 | 161 | def _preprocess(self, data): 162 | 163 | preprocessed_data = {} 164 | 165 | pre_st = time.time() 166 | 167 | for k, v in data.items(): 168 | for file_name, file_content in v.items(): 169 | # print('\tAppending image: %s' % file_name) 170 | image = decode_image(file_content) 171 | height, width = image.shape[:2] 172 | image = self.transform(image)[0].unsqueeze(0) 173 | sample = {'img': image, 'img_name': file_name,'scale':[height,width]} 174 | preprocessed_data[k] = sample 175 | 176 | pre_et = time.time() 177 | self.pre_time = pre_st-pre_et 178 | return preprocessed_data 179 | 180 | def _inference(self, data): 181 | sample = data[IMAGES_KEY] # img, img_name, scale 182 | images=sample['img'] 183 | height, width=sample['scale'] 184 | st = time.time() 185 | if torch.cuda.is_available(): 186 | device = torch.device('cuda') 187 | cpu_device = torch.device('cpu') 188 | else: 189 | device = torch.device('cpu') 190 | self.model = self.model.to(device) 191 | self.model = self.model.eval() 192 | with torch.no_grad(): 193 | result = self.model(images.to(device))[0] 194 | result = result.resize((width, height)).to(cpu_device).numpy() 195 | boxes, labels, scores = result['boxes'], result['labels'], result['scores'] 196 | indices = scores > self.score_threshold 197 | boxes = boxes[indices] 198 | labels = labels[indices] 199 | scores = scores[indices] 200 | et = time.time() 201 | # print('Elapsed time: {} or {}(et-st)'.format(time.time() - st, et - st)) 202 | result_i = {'img_name': sample['img_name'], 'scale': sample['scale'], 203 | 'class': labels, 'pred_boxes': boxes, 204 | 'score': scores, 'time': et - st} 205 | 206 | return result_i 207 | 208 | def _postprocess(self, data): 209 | class_name = self.label 210 | labels = label_list(os.path.join(self.dir_path, 'data/class_name.csv')) # ('data/class_name.csv')# 211 | 212 | post_st = time.time() 213 | result = data 214 | img_name = result['img_name'] 215 | scale = result['scale'] 216 | classification = result['class'] 217 | predict_bboxes = result['pred_boxes'] 218 | scores = result['score'] 219 | lantecy_time = result['time'] 220 | detection_bboxes = [] 221 | detection_classes = [] 222 | for c in range(len(scores)): 223 | text1 = self.label[classification[c]] 224 | text2=self.classify[text1] 225 | text=text1+'/'+text2 226 | detection_classes.append(text) 227 | detection_bboxes.append(predict_bboxes[c]) 228 | 229 | post_et = time.time() 230 | self.post_time = post_st - post_et 231 | all_run_time = lantecy_time + self.pre_time + self.post_time 232 | all_run_time *=1000 # ms 233 | json_file = save_result_as_json(img_name, detection_classes, np.array(scores), 234 | np.array(detection_bboxes), all_run_time) 235 | 236 | return json_file 237 | 238 | def inference(self, data): 239 | pre_start_time = time.time() 240 | data = self._preprocess(data) 241 | infer_start_time = time.time() 242 | # Update preprocess latency metric 243 | pre_time_in_ms = (infer_start_time - pre_start_time) * 1000 244 | logger.info('preprocess time: ' + str(pre_time_in_ms) + 'ms') 245 | 246 | if self.model_name + '_LatencyPreprocess' in MetricsManager.metrics: 247 | MetricsManager.metrics[self.model_name + '_LatencyPreprocess'].update(pre_time_in_ms) 248 | 249 | data = self._inference(data) 250 | infer_end_time = time.time() 251 | infer_in_ms = (infer_end_time - infer_start_time) * 1000 252 | 253 | logger.info('infer time: ' + str(infer_in_ms) + 'ms') 254 | data = self._postprocess(data) 255 | 256 | # Update inference latency metric 257 | post_time_in_ms = (time.time() - infer_end_time) * 1000 258 | logger.info('postprocess time: ' + str(post_time_in_ms) + 'ms') 259 | if self.model_name + '_LatencyInference' in MetricsManager.metrics: 260 | MetricsManager.metrics[self.model_name + '_LatencyInference'].update(post_time_in_ms) 261 | 262 | # Update overall latency metric 263 | if self.model_name + '_LatencyOverall' in MetricsManager.metrics: 264 | MetricsManager.metrics[self.model_name + '_LatencyOverall'].update(pre_time_in_ms + post_time_in_ms) 265 | 266 | logger.info('latency: ' + str(pre_time_in_ms + infer_in_ms + post_time_in_ms) + 'ms') 267 | data['latency_time'] = str(round(pre_time_in_ms + infer_in_ms + post_time_in_ms, 1)) + ' ms' 268 | return data 269 | 270 | 271 | -------------------------------------------------------------------------------- /data_aug/README.md: -------------------------------------------------------------------------------- 1 | # DataAugmentation_ForObjectDetect 2 | 本仓库主要包含了针对目标检测数据集的增强手段和源码：图像的旋转，镜像，裁剪，亮度/对比度的变换等采取的实验数据格式为VOC的格式,标签存储在xml文件中.代码中涉及一些对.xml文件的基本操作 3 | ## rotated.py包含了旋转图像以及对应的标签的旋转代码 4 | 对图像的旋转可以调用cv2的库函数，但是对标签的操作需要搞清楚一个点绕着另一个点旋转后的坐标与原坐标的对应关系，计算公式见代码！ 5 | 6 | ## padd_rotated_crop.py 7 | > 由于图像的旋转会使得一部分信息损失，而且不是很容易的判断旋转后的图像是否还包含我们的完整目标。 8 | 因此可以尝试把原图像先安装长短边的长度填充为一个正方形，这样可以确保填充后的图像在旋转过程中，原图中的目标信息不会丢失。 9 | 然后可以旋转作一个裁剪（根据目标的标签坐标，只要避开旋转后的标签坐标就能很好的裁剪出图像），当然也可以不裁剪，但是这样会降低模型 10 | 训练的效率 11 | 12 | ## color.py 目前只是对图像的亮度/对比度的改变 13 | > 该部分代码还应该有对图像的颜色变换，下次更新！对亮度和对比度的调节利用的是像素的简单线性变换。 14 | 15 | ## mirror.py 16 | > 该部分代码主要是对图像进行水平，竖直，对角的镜像变换 17 | 对图像的操作cv2.flip()；对标签的变化其实就是利用原图的长宽来减去原像素，三种情况有所不同，但很好理解，只是有些细节需要注意，见代码！ 18 | 注：对原图的对角镜像其实等价于对原图旋转180° 19 | -------------------------------------------------------------------------------- /data_aug/color.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import math 3 | import numpy as np 4 | import os 5 | import glob 6 | import json 7 | import shutil 8 | import xml.etree.ElementTree as ET 9 | from xml.etree.ElementTree import ElementTree, Element 10 | need=['塑料衣架', '金属厨具', '快递纸袋', '污损用纸', '花盆', '书籍纸张', '牙签', '垃圾桶'] 11 | def getColorImg(alpha,beta,img_path,img_write_path): 12 | img = cv2.imread(img_path)# 13 | colored_img = np.uint8(np.clip((alpha * img + beta), 0, 255)) 14 | cv2.imwrite(img_write_path,colored_img) 15 | 16 | def getColorAnno(anno_path,anno_write_path): 17 | tree = ET.parse(anno_path) 18 | tree.write(anno_write_path) # 保存修改后的XML文件 19 | 20 | def color(alpha,beta,img_dir,anno_dir,img_write_dir,anno_write_dir): 21 | if not os.path.exists(img_write_dir): 22 | os.makedirs(img_write_dir) 23 | 24 | if not os.path.exists(anno_write_dir): 25 | os.makedirs(anno_write_dir) 26 | img_names=os.listdir(img_dir) 27 | for img_name in img_names: 28 | img_path=os.path.join(img_dir,img_name) 29 | img_write_path=os.path.join(img_write_dir,img_name[:-4]+'color'+str(int(alpha*10))+'.jpg') 30 | # 31 | anno_path=os.path.join(anno_dir,img_name[:-4]+'.xml') 32 | objects = ET.parse(anno_path).findall("object") 33 | flag = False 34 | for obj in objects: 35 | class_name = obj.find('name').text.lower().strip() 36 | if class_name in need: 37 | flag = True 38 | if flag==True: 39 | anno_write_path = os.path.join(anno_write_dir, img_name[:-4]+'color'+str(int(alpha*10))+'.xml') 40 | # 41 | getColorImg(alpha,beta,img_path,img_write_path) 42 | getColorAnno(anno_path,anno_write_path) 43 | 44 | alphas=[0.5,1.5]#[0.3,0.5,1.2,1.6] 45 | beta=10 46 | img_dir='../datasets/VOC2007/JPEGImages' 47 | anno_dir='../datasets/VOC2007//Annotations' 48 | img_write_dir='../datasets/VOC2007/Mirror/mirrored_JPEGImages' 49 | anno_write_dir='../datasets/VOC2007/Mirror/mirrored_Annotations' 50 | for alpha in alphas: 51 | color(alpha,beta,img_dir,anno_dir,img_write_dir,anno_write_dir) -------------------------------------------------------------------------------- /data_aug/mirror.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import math 3 | import numpy as np 4 | import os 5 | import glob 6 | import json 7 | import shutil 8 | import xml.etree.ElementTree as ET 9 | from xml.etree.ElementTree import ElementTree, Element 10 | ext=['陶瓷器皿', '剩饭剩菜', '插头电线', '污损塑料', '果皮果肉', '洗护用品'] 11 | need1=['砧板', '鞋', '蛋壳', '食用油桶', '干电池', '包', '塑料玩具', '烟蒂', '软膏', '易拉罐', '大骨头', 12 | '充电宝', '茶叶渣', '金属食品罐', '旧衣服', '枕头', '饮料瓶', '酒瓶', '金属器皿', '一次性快餐盒', 13 | '塑料衣架', '金属厨具', '快递纸袋', '污损用纸', '花盆', '书籍纸张', '牙签', '垃圾桶'] 14 | need2=['玻璃器皿', '筷子', '过期药物', '塑料器皿', '毛绒玩具', '纸盒纸箱', '调料瓶', '锅', '鱼骨'] 15 | need3=['花盆', '牙签', '金属器皿', '书籍纸张', '金属厨具', '污损用纸', '垃圾桶'] 16 | def h_MirrorImg(img_path,img_write_path): 17 | img = cv2.imread(img_path) 18 | mirror_img = cv2.flip(img, 1) # 19 | cv2.imwrite(img_write_path,mirror_img) 20 | def v_MirrorImg(img_path,img_write_path): 21 | img = cv2.imread(img_path) 22 | mirror_img = cv2.flip(img, 0) # 23 | cv2.imwrite(img_write_path,mirror_img) 24 | def a_MirrorImg(img_path,img_write_path): 25 | img = cv2.imread(img_path) 26 | mirror_img = cv2.flip(img, -1) # 27 | cv2.imwrite(img_write_path,mirror_img) 28 | 29 | def h_MirrorAnno(anno_path,anno_write_path): 30 | tree = ET.parse(anno_path) 31 | root = tree.getroot() 32 | size=root.find('size') 33 | w=int(size.find('width').text) 34 | objects = root.findall("object") 35 | for obj in objects: 36 | bbox = obj.find('bndbox') 37 | x1 = float(bbox.find('xmin').text) 38 | x2 = float(bbox.find('xmax').text) 39 | x1=w-x1+1 40 | x2=w-x2+1 41 | 42 | assert x1>0 43 | assert x2>0 44 | 45 | bbox.find('xmin').text=str(int(x2)) 46 | bbox.find('xmax').text=str(int(x1)) 47 | 48 | tree.write(anno_write_path) # 保存修改后的XML文件 49 | def v_MirrorAnno(anno_path,anno_write_path): 50 | tree = ET.parse(anno_path) 51 | root = tree.getroot() 52 | size = root.find('size') 53 | h=int(size.find('height').text) 54 | objects = root.findall("object") 55 | for obj in objects: 56 | bbox = obj.find('bndbox') 57 | y1 = float(bbox.find('ymin').text) 58 | y2 = float(bbox.find('ymax').text) 59 | 60 | y1=h-y1+1 61 | y2=h-y2+1 62 | 63 | assert y1>0 64 | assert y2>0 65 | 66 | bbox.find('ymin').text=str(int(y2)) 67 | bbox.find('ymax').text=str(int(y1)) 68 | 69 | tree.write(anno_write_path) # 保存修改后的XML文件 70 | def a_MirrorAnno(anno_path,anno_write_path): 71 | tree = ET.parse(anno_path) 72 | root = tree.getroot() 73 | size = root.find('size') 74 | w=int(size.find('width').text) 75 | h = int(size.find('height').text) 76 | objects = root.findall("object") 77 | for obj in objects: 78 | bbox = obj.find('bndbox') 79 | x1 = float(bbox.find('xmin').text) 80 | y1 = float(bbox.find('ymin').text) 81 | x2 = float(bbox.find('xmax').text) 82 | y2 = float(bbox.find('ymax').text) 83 | 84 | x1=w-x1+1 85 | x2=w-x2+1 86 | 87 | y1 = h - y1+1 88 | y2 = h - y2+1 89 | 90 | assert x1 > 0 91 | assert x2 > 0 92 | assert y1 > 0 93 | assert y2 > 0 94 | 95 | bbox.find('xmin').text=str(int(x2)) 96 | bbox.find('xmax').text=str(int(x1)) 97 | bbox.find('ymin').text=str(int(y2)) 98 | bbox.find('ymax').text=str(int(y1)) 99 | 100 | tree.write(anno_write_path) # 保存修改后的XML文件 101 | def mirror(img_dir,anno_dir,img_write_dir,anno_write_dir): 102 | if not os.path.exists(img_write_dir): 103 | os.makedirs(img_write_dir) 104 | 105 | if not os.path.exists(anno_write_dir): 106 | os.makedirs(anno_write_dir) 107 | img_names=os.listdir(img_dir) 108 | cnt=0 109 | aug_cnt=0 110 | for img_name in img_names: 111 | if cnt%1000==0: 112 | print(cnt,aug_cnt) 113 | cnt=cnt+1 114 | img_path=os.path.join(img_dir,img_name) 115 | anno_path=os.path.join(anno_dir,img_name[:-4]+'.xml') 116 | objects = ET.parse(anno_path).findall("object") 117 | flag=False 118 | t=0 119 | r=0 120 | for obj in objects: 121 | t+=1 122 | class_name = obj.find('name').text.lower().strip() 123 | if class_name in need3 : 124 | flag=True 125 | r+=1 126 | #if r!=t: 127 | #flag=False 128 | if flag==True: 129 | aug_cnt += 1 130 | h_img_write_path = os.path.join(img_write_dir, img_name[:-4] + 'h' + '.jpg') 131 | h_anno_write_path = os.path.join(anno_write_dir, img_name[:-4]+'h'+'.xml') 132 | # 133 | v_img_write_path = os.path.join(img_write_dir, img_name[:-4] + 'v' + '.jpg') 134 | v_anno_write_path = os.path.join(anno_write_dir, img_name[:-4] + 'v' + '.xml') 135 | # 136 | a_img_write_path = os.path.join(img_write_dir, img_name[:-4] + 'a' + '.jpg') 137 | a_anno_write_path = os.path.join(anno_write_dir, img_name[:-4] + 'a' + '.xml') 138 | # 139 | h_MirrorImg(img_path,h_img_write_path) 140 | v_MirrorImg(img_path,v_img_write_path) 141 | a_MirrorImg(img_path, a_img_write_path) 142 | h_MirrorAnno(anno_path,h_anno_write_path) 143 | v_MirrorAnno(anno_path, v_anno_write_path) 144 | a_MirrorAnno(anno_path, a_anno_write_path) 145 | 146 | 147 | img_dir='../datasets/VOC2007/JPEGImages_raw' 148 | anno_dir='../datasets/VOC2007/Annotations_raw' 149 | img_write_dir='../datasets/VOC2007/Mirror3/mirrored_JPEGImages' 150 | anno_write_dir='../datasets/VOC2007/Mirror3/mirrored_Annotations' 151 | if os.path.exists(img_write_dir): 152 | os.makedirs(img_write_dir) 153 | if os.path.exists(anno_write_dir): 154 | os.makedirs(anno_write_dir) 155 | mirror(img_dir,anno_dir,img_write_dir,anno_write_dir) 156 | ''' 157 | imga=cv2.imread('14a.jpg') 158 | imgh=cv2.imread('14h.jpg') 159 | imgv=cv2.imread('14v.jpg') 160 | 161 | 162 | cv2.rectangle(imga,(606,898),(855,1432),(0,255,0),4) 163 | cv2.rectangle(imgh,(606,489),(855,1023),(0,255,0),4) 164 | cv2.rectangle(imgv,(32,898),(281,1432),(0,255,0),4) 165 | cv2.imwrite('a.jpg',imga) 166 | cv2.imwrite('v.jpg',imgv) 167 | cv2.imwrite('h.jpg',imgh) 168 | ''' 169 | -------------------------------------------------------------------------------- /data_aug/padd_rotated_crop.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import math 3 | img=cv2.imread('14h.jpg') 4 | rows, cols = img.shape[:2] 5 | ##填充图像为正方形，而且要能保证填充后的图像在0到360°旋转的时候，原图像的像素不会损失 6 | re=cv2.copyMakeBorder(img,int(cols/2),int(cols/2),int(rows/2),int(rows/2),cv2.BORDER_CONSTANT) 7 | 8 | def getRotatedImg(Pi_angle,img_path,img_write_path): 9 | img = cv2.imread(img_path) 10 | rows, cols = img.shape[:2] 11 | a, b = cols / 2, rows / 2 12 | M = cv2.getRotationMatrix2D((a, b), angle, 1) 13 | rotated_img = cv2.warpAffine(img, M, (cols, rows)) # 旋转后的图像保持大小不变 14 | cv2.imwrite(img_write_path,rotated_img) 15 | 16 | for angle in range(0,180,30): 17 | Pi_angle = -angle * math.pi / 180.0 18 | img_path='re.jpg' 19 | img_write_path=str(angle)+'.jpg' 20 | getRotatedImg(Pi_angle, img_path, img_write_path) 21 | 22 | #验证是否标签被正确的改变 23 | #for origin image: xmin:606 ymin:489 xmax:855 ymax:1023 24 | cv2.rectangle(re,(606+int(rows/2),489+int(cols/2)),(855+int(rows/2),1023+int(cols/2)),(0,255,0),4) 25 | # 26 | def crop(): 27 | pass -------------------------------------------------------------------------------- /data_aug/rotated.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import math 3 | import numpy as np 4 | import os 5 | import glob 6 | import json 7 | import shutil 8 | import xml.etree.ElementTree as ET 9 | from xml.etree.ElementTree import ElementTree, Element 10 | 11 | def getRotatedImg(Pi_angle,img_path,img_write_path): 12 | img = cv2.imread(img_path) 13 | rows, cols = img.shape[:2] 14 | a, b = cols / 2, rows / 2 15 | M = cv2.getRotationMatrix2D((a, b), angle, 1) 16 | rotated_img = cv2.warpAffine(img, M, (cols, rows)) # 旋转后的图像保持大小不变 17 | cv2.imwrite(img_write_path,rotated_img) 18 | return a,b 19 | 20 | def getRotatedAnno(Pi_angle,a,b,anno_path,anno_write_path): 21 | tree = ET.parse(anno_path) 22 | root = tree.getroot() 23 | objects = root.findall("object") 24 | for obj in objects: 25 | bbox = obj.find('bndbox') 26 | x1 = float(bbox.find('xmin').text) - 1 27 | y1 = float(bbox.find('ymin').text) - 1 28 | x2 = float(bbox.find('xmax').text) - 1 29 | y2 = float(bbox.find('ymax').text) - 1 30 | 31 | x3=x1 32 | y3=y2 33 | x4=x2 34 | y4=y1 35 | 36 | X1 = (x1 - a) * math.cos(Pi_angle) - (y1 - b) * math.sin(Pi_angle) + a 37 | Y1 = (x1 - a) * math.sin(Pi_angle) + (y1 - b) * math.cos(Pi_angle) + b 38 | 39 | X2 = (x2 - a) * math.cos(Pi_angle) - (y2 - b) * math.sin(Pi_angle) + a 40 | Y2 = (x2 - a) * math.sin(Pi_angle) + (y2 - b) * math.cos(Pi_angle) + b 41 | 42 | X3 = (x3 - a) * math.cos(Pi_angle) - (y3 - b) * math.sin(Pi_angle) + a 43 | Y3 = (x3 - a) * math.sin(Pi_angle) + (y3 - b) * math.cos(Pi_angle) + b 44 | 45 | X4 = (x4 - a) * math.cos(Pi_angle) - (y4 - b) * math.sin(Pi_angle) + a 46 | Y4 = (x4 - a) * math.sin(Pi_angle) + (y4 - b) * math.cos(Pi_angle) + b 47 | 48 | X_MIN=min(X1,X2,X3,X4) 49 | X_MAX = max(X1, X2, X3, X4) 50 | Y_MIN = min(Y1, Y2, Y3, Y4) 51 | Y_MAX = max(Y1, Y2, Y3, Y4) 52 | 53 | bbox.find('xmin').text=str(int(X_MIN)) 54 | bbox.find('ymin').text=str(int(Y_MIN)) 55 | bbox.find('xmax').text=str(int(X_MAX)) 56 | bbox.find('ymax').text=str(int(Y_MAX)) 57 | 58 | tree.write(anno_write_path) # 保存修改后的XML文件 59 | 60 | def rotate(angle,img_dir,anno_dir,img_write_dir,anno_write_dir): 61 | if not os.path.exists(img_write_dir): 62 | os.makedirs(img_write_dir) 63 | 64 | if not os.path.exists(anno_write_dir): 65 | os.makedirs(anno_write_dir) 66 | 67 | Pi_angle = -angle * math.pi / 180.0 # 弧度制，后面旋转坐标需要用到，注意负号！！！ 68 | img_names=os.listdir(img_dir) 69 | for img_name in img_names: 70 | img_path=os.path.join(img_dir,img_name) 71 | img_write_path=os.path.join(img_write_dir,img_name[:-4]+'R'+str(angle)+'.jpg') 72 | # 73 | anno_path=os.path.join(anno_dir,img_name[:-4]+'.xml') 74 | anno_write_path = os.path.join(anno_write_dir, img_name[:-4]+'R'+str(angle)+'.xml') 75 | # 76 | a,b=getRotatedImg(Pi_angle,img_path,img_write_path) 77 | getRotatedAnno(Pi_angle,a,b,anno_path,anno_write_path) 78 | 79 | angle=180 80 | img_dir='several/JPEGImages' 81 | anno_dir='several/Annotations' 82 | img_write_dir='Rotated/rotated_JPEGImages' 83 | anno_write_dir='Rotated/rotated_Annotations' 84 | 85 | rotate(angle,img_dir,anno_dir,img_write_dir,anno_write_dir) -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | import time 4 | 5 | import torch 6 | from PIL import Image 7 | from vizer.draw import draw_boxes 8 | 9 | from ssd.config import cfg 10 | from ssd.data.datasets import COCODataset, VOCDataset 11 | import argparse 12 | import numpy as np 13 | 14 | from ssd.data.transforms import build_transforms 15 | from ssd.modeling.detector import build_detection_model 16 | from ssd.utils import mkdir 17 | from ssd.utils.checkpoint import CheckPointer 18 | 19 | 20 | @torch.no_grad() 21 | def run_demo(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type): 22 | if dataset_type == "voc": 23 | class_names = VOCDataset.class_names 24 | elif dataset_type == 'coco': 25 | class_names = COCODataset.class_names 26 | else: 27 | raise NotImplementedError('Not implemented now.') 28 | device = torch.device(cfg.MODEL.DEVICE) 29 | 30 | model = build_detection_model(cfg) 31 | model = model.to(device) 32 | checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR) 33 | checkpointer.load(ckpt, use_latest=ckpt is None) 34 | weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file() 35 | print('Loaded weights from {}'.format(weight_file)) 36 | 37 | image_paths = glob.glob(os.path.join(images_dir, '*.jpg')) 38 | mkdir(output_dir) 39 | 40 | cpu_device = torch.device("cpu") 41 | transforms = build_transforms(cfg, is_train=False) 42 | model.eval() 43 | for i, image_path in enumerate(image_paths): 44 | start = time.time() 45 | image_name = os.path.basename(image_path) 46 | 47 | image = np.array(Image.open(image_path).convert("RGB")) 48 | height, width = image.shape[:2] 49 | images = transforms(image)[0].unsqueeze(0) 50 | load_time = time.time() - start 51 | 52 | start = time.time() 53 | result = model(images.to(device))[0] 54 | inference_time = time.time() - start 55 | 56 | result = result.resize((width, height)).to(cpu_device).numpy() 57 | boxes, labels, scores = result['boxes'], result['labels'], result['scores'] 58 | 59 | indices = scores > score_threshold 60 | boxes = boxes[indices] 61 | labels = labels[indices] 62 | scores = scores[indices] 63 | meters = ' | '.join( 64 | [ 65 | 'objects {:02d}'.format(len(boxes)), 66 | 'load {:03d}ms'.format(round(load_time * 1000)), 67 | 'inference {:03d}ms'.format(round(inference_time * 1000)), 68 | 'FPS {}'.format(round(1.0 / inference_time)) 69 | ] 70 | ) 71 | print('({:04d}/{:04d}) {}: {}'.format(i + 1, len(image_paths), image_name, meters)) 72 | 73 | drawn_image = draw_boxes(image, boxes, labels, scores, class_names).astype(np.uint8) 74 | Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name)) 75 | 76 | 77 | def main(): 78 | parser = argparse.ArgumentParser(description="SSD Demo.") 79 | parser.add_argument( 80 | "--config-file", 81 | default="configs/efficient_net_b3_ssd300_voc0712.yaml", 82 | metavar="FILE", 83 | help="path to config file", 84 | type=str, 85 | ) 86 | parser.add_argument("--ckpt", type=str, default=None, help="Trained weights.") 87 | parser.add_argument("--score_threshold", type=float, default=0.7) 88 | parser.add_argument("--images_dir", default='demo/test_trash', type=str, help='Specify a image dir to do prediction.') 89 | parser.add_argument("--output_dir", default='demo/result', type=str, help='Specify a image dir to save predicted images.') 90 | parser.add_argument("--dataset_type", default="voc", type=str, help='Specify dataset type. Currently support voc and coco.') 91 | 92 | parser.add_argument( 93 | "opts", 94 | help="Modify config options using the command-line", 95 | default=None, 96 | nargs=argparse.REMAINDER, 97 | ) 98 | args = parser.parse_args() 99 | print(args) 100 | 101 | cfg.merge_from_file(args.config_file) 102 | cfg.merge_from_list(args.opts) 103 | cfg.freeze() 104 | 105 | print("Loaded configuration file {}".format(args.config_file)) 106 | with open(args.config_file, "r") as cf: 107 | config_str = "\n" + cf.read() 108 | print(config_str) 109 | print("Running with config:\n{}".format(cfg)) 110 | 111 | run_demo(cfg=cfg, 112 | ckpt=args.ckpt, 113 | score_threshold=args.score_threshold, 114 | images_dir=args.images_dir, 115 | output_dir=args.output_dir, 116 | dataset_type=args.dataset_type) 117 | 118 | 119 | if __name__ == '__main__': 120 | main() 121 | -------------------------------------------------------------------------------- /draw_rectangle.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | cv_image=cv2.imread('1.jpg') 4 | outs=[(156, 3, 976, 1445), (320, -21, 846, 1451)] 5 | for c in outs: 6 | cv2.rectangle(cv_image, c, 7 | (0, 0, 255), 4) 8 | cv2.imwrite('1x.jpg', cv_image) 9 | -------------------------------------------------------------------------------- /ext/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ext/__init__.py -------------------------------------------------------------------------------- /ext/build.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | 4 | import torch 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import CUDA_HOME 7 | from torch.utils.cpp_extension import CppExtension 8 | from torch.utils.cpp_extension import CUDAExtension 9 | 10 | requirements = ["torch"] 11 | 12 | 13 | def get_extensions(): 14 | extensions_dir = os.path.dirname(os.path.abspath(__file__)) 15 | 16 | main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) 17 | source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp")) 18 | source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu")) 19 | 20 | sources = main_file + source_cpu 21 | extension = CppExtension 22 | 23 | extra_compile_args = {"cxx": []} 24 | define_macros = [] 25 | 26 | if torch.cuda.is_available() and CUDA_HOME is not None: 27 | extension = CUDAExtension 28 | sources += source_cuda 29 | define_macros += [("WITH_CUDA", None)] 30 | extra_compile_args["nvcc"] = [ 31 | "-DCUDA_HAS_FP16=1", 32 | "-D__CUDA_NO_HALF_OPERATORS__", 33 | "-D__CUDA_NO_HALF_CONVERSIONS__", 34 | "-D__CUDA_NO_HALF2_OPERATORS__", 35 | ] 36 | 37 | sources = [os.path.join(extensions_dir, s) for s in sources] 38 | 39 | include_dirs = [extensions_dir] 40 | 41 | ext_modules = [ 42 | extension( 43 | "torch_extension", 44 | sources, 45 | include_dirs=include_dirs, 46 | define_macros=define_macros, 47 | extra_compile_args=extra_compile_args, 48 | ) 49 | ] 50 | 51 | return ext_modules 52 | 53 | 54 | setup( 55 | name="torch_extension", 56 | version="0.1", 57 | ext_modules=get_extensions(), 58 | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}) 59 | -------------------------------------------------------------------------------- /ext/cpu/nms_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include "cpu/vision.h" 3 | 4 | 5 | template 6 | at::Tensor nms_cpu_kernel(const at::Tensor& dets, 7 | const at::Tensor& scores, 8 | const float threshold) { 9 | AT_ASSERTM(!dets.type().is_cuda(), "dets must be a CPU tensor"); 10 | AT_ASSERTM(!scores.type().is_cuda(), "scores must be a CPU tensor"); 11 | AT_ASSERTM(dets.type() == scores.type(), "dets should have the same type as scores"); 12 | 13 | if (dets.numel() == 0) { 14 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 15 | } 16 | 17 | auto x1_t = dets.select(1, 0).contiguous(); 18 | auto y1_t = dets.select(1, 1).contiguous(); 19 | auto x2_t = dets.select(1, 2).contiguous(); 20 | auto y2_t = dets.select(1, 3).contiguous(); 21 | 22 | at::Tensor areas_t = (x2_t - x1_t) * (y2_t - y1_t); 23 | 24 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 25 | 26 | auto ndets = dets.size(0); 27 | at::Tensor suppressed_t = at::zeros({ndets}, dets.options().dtype(at::kByte).device(at::kCPU)); 28 | 29 | auto suppressed = suppressed_t.data(); 30 | auto order = order_t.data(); 31 | auto x1 = x1_t.data(); 32 | auto y1 = y1_t.data(); 33 | auto x2 = x2_t.data(); 34 | auto y2 = y2_t.data(); 35 | auto areas = areas_t.data(); 36 | 37 | for (int64_t _i = 0; _i < ndets; _i++) { 38 | auto i = order[_i]; 39 | if (suppressed[i] == 1) 40 | continue; 41 | auto ix1 = x1[i]; 42 | auto iy1 = y1[i]; 43 | auto ix2 = x2[i]; 44 | auto iy2 = y2[i]; 45 | auto iarea = areas[i]; 46 | 47 | for (int64_t _j = _i + 1; _j < ndets; _j++) { 48 | auto j = order[_j]; 49 | if (suppressed[j] == 1) 50 | continue; 51 | auto xx1 = std::max(ix1, x1[j]); 52 | auto yy1 = std::max(iy1, y1[j]); 53 | auto xx2 = std::min(ix2, x2[j]); 54 | auto yy2 = std::min(iy2, y2[j]); 55 | 56 | auto w = std::max(static_cast(0), xx2 - xx1); 57 | auto h = std::max(static_cast(0), yy2 - yy1); 58 | auto inter = w * h; 59 | auto ovr = inter / (iarea + areas[j] - inter); 60 | if (ovr >= threshold) 61 | suppressed[j] = 1; 62 | } 63 | } 64 | return at::nonzero(suppressed_t == 0).squeeze(1); 65 | } 66 | 67 | at::Tensor nms_cpu(const at::Tensor& dets, 68 | const at::Tensor& scores, 69 | const float threshold) { 70 | at::Tensor result; 71 | AT_DISPATCH_FLOATING_TYPES(dets.type(), "nms", [&] { 72 | result = nms_cpu_kernel(dets, scores, threshold); 73 | }); 74 | return result; 75 | } -------------------------------------------------------------------------------- /ext/cpu/vision.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include 4 | 5 | at::Tensor nms_cpu(const at::Tensor& dets, 6 | const at::Tensor& scores, 7 | const float threshold); 8 | -------------------------------------------------------------------------------- /ext/cuda/nms.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | int const threadsPerBlock = sizeof(unsigned long long) * 8; 12 | 13 | __device__ inline float devIoU(float const * const a, float const * const b) { 14 | float left = max(a[0], b[0]), right = min(a[2], b[2]); 15 | float top = max(a[1], b[1]), bottom = min(a[3], b[3]); 16 | float width = max(right - left, 0.f), height = max(bottom - top, 0.f); 17 | float interS = width * height; 18 | float Sa = (a[2] - a[0]) * (a[3] - a[1]); 19 | float Sb = (b[2] - b[0]) * (b[3] - b[1]); 20 | return interS / (Sa + Sb - interS); 21 | } 22 | 23 | __global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh, 24 | const float *dev_boxes, unsigned long long *dev_mask) { 25 | const int row_start = blockIdx.y; 26 | const int col_start = blockIdx.x; 27 | 28 | // if (row_start > col_start) return; 29 | 30 | const int row_size = 31 | min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); 32 | const int col_size = 33 | min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); 34 | 35 | __shared__ float block_boxes[threadsPerBlock * 5]; 36 | if (threadIdx.x < col_size) { 37 | block_boxes[threadIdx.x * 5 + 0] = 38 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; 39 | block_boxes[threadIdx.x * 5 + 1] = 40 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; 41 | block_boxes[threadIdx.x * 5 + 2] = 42 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; 43 | block_boxes[threadIdx.x * 5 + 3] = 44 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; 45 | block_boxes[threadIdx.x * 5 + 4] = 46 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; 47 | } 48 | __syncthreads(); 49 | 50 | if (threadIdx.x < row_size) { 51 | const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; 52 | const float *cur_box = dev_boxes + cur_box_idx * 5; 53 | int i = 0; 54 | unsigned long long t = 0; 55 | int start = 0; 56 | if (row_start == col_start) { 57 | start = threadIdx.x + 1; 58 | } 59 | for (i = start; i < col_size; i++) { 60 | if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) { 61 | t |= 1ULL << i; 62 | } 63 | } 64 | const int col_blocks = THCCeilDiv(n_boxes, threadsPerBlock); 65 | dev_mask[cur_box_idx * col_blocks + col_start] = t; 66 | } 67 | } 68 | 69 | // boxes is a N x 5 tensor 70 | at::Tensor nms_cuda(const at::Tensor boxes, float nms_overlap_thresh) { 71 | using scalar_t = float; 72 | AT_ASSERTM(boxes.type().is_cuda(), "boxes must be a CUDA tensor"); 73 | auto scores = boxes.select(1, 4); 74 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 75 | auto boxes_sorted = boxes.index_select(0, order_t); 76 | 77 | int boxes_num = boxes.size(0); 78 | 79 | const int col_blocks = THCCeilDiv(boxes_num, threadsPerBlock); 80 | 81 | scalar_t* boxes_dev = boxes_sorted.data(); 82 | 83 | THCState *state = at::globalContext().lazyInitCUDA(); // TODO replace with getTHCState 84 | 85 | unsigned long long* mask_dev = NULL; 86 | //THCudaCheck(THCudaMalloc(state, (void**) &mask_dev, 87 | // boxes_num * col_blocks * sizeof(unsigned long long))); 88 | 89 | mask_dev = (unsigned long long*) THCudaMalloc(state, boxes_num * col_blocks * sizeof(unsigned long long)); 90 | 91 | dim3 blocks(THCCeilDiv(boxes_num, threadsPerBlock), 92 | THCCeilDiv(boxes_num, threadsPerBlock)); 93 | dim3 threads(threadsPerBlock); 94 | nms_kernel<<>>(boxes_num, 95 | nms_overlap_thresh, 96 | boxes_dev, 97 | mask_dev); 98 | 99 | std::vector mask_host(boxes_num * col_blocks); 100 | THCudaCheck(cudaMemcpy(&mask_host[0], 101 | mask_dev, 102 | sizeof(unsigned long long) * boxes_num * col_blocks, 103 | cudaMemcpyDeviceToHost)); 104 | 105 | std::vector remv(col_blocks); 106 | memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); 107 | 108 | at::Tensor keep = at::empty({boxes_num}, boxes.options().dtype(at::kLong).device(at::kCPU)); 109 | int64_t* keep_out = keep.data(); 110 | 111 | int num_to_keep = 0; 112 | for (int i = 0; i < boxes_num; i++) { 113 | int nblock = i / threadsPerBlock; 114 | int inblock = i % threadsPerBlock; 115 | 116 | if (!(remv[nblock] & (1ULL << inblock))) { 117 | keep_out[num_to_keep++] = i; 118 | unsigned long long *p = &mask_host[0] + i * col_blocks; 119 | for (int j = nblock; j < col_blocks; j++) { 120 | remv[j] |= p[j]; 121 | } 122 | } 123 | } 124 | 125 | THCudaFree(state, mask_dev); 126 | // TODO improve this part 127 | return std::get<0>(order_t.index({ 128 | keep.narrow(/*dim=*/0, /*start=*/0, /*length=*/num_to_keep).to( 129 | order_t.device(), keep.scalar_type()) 130 | }).sort(0, false)); 131 | } -------------------------------------------------------------------------------- /ext/cuda/vision.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include 4 | 5 | at::Tensor nms_cuda(const at::Tensor boxes, float nms_overlap_thresh); 6 | 7 | 8 | -------------------------------------------------------------------------------- /ext/nms.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include "cpu/vision.h" 4 | 5 | #ifdef WITH_CUDA 6 | #include "cuda/vision.h" 7 | #endif 8 | 9 | 10 | at::Tensor nms(const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float threshold) { 13 | 14 | if (dets.type().is_cuda()) { 15 | #ifdef WITH_CUDA 16 | // TODO raise error if not compiled with CUDA 17 | if (dets.numel() == 0) 18 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 19 | auto b = at::cat({dets, scores.unsqueeze(1)}, 1); 20 | return nms_cuda(b, threshold); 21 | #else 22 | AT_ERROR("Not compiled with GPU support"); 23 | #endif 24 | } 25 | 26 | at::Tensor result = nms_cpu(dets, scores, threshold); 27 | return result; 28 | } 29 | -------------------------------------------------------------------------------- /ext/vision.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include "nms.h" 3 | 4 | 5 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 6 | m.def("nms", &nms, "non-maximum suppression"); 7 | } 8 | -------------------------------------------------------------------------------- /figures/004545.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/figures/004545.jpg -------------------------------------------------------------------------------- /figures/losses.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/figures/losses.png -------------------------------------------------------------------------------- /figures/lr.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/figures/lr.png -------------------------------------------------------------------------------- /figures/metrics.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/figures/metrics.png -------------------------------------------------------------------------------- /get_ratio.py: -------------------------------------------------------------------------------- 1 | import os 2 | import xml.etree.ElementTree as ET 3 | ann_dir='datasets/VOC2007/Annotations' 4 | img_dir='datasets/VOC2007/JPEGImages' 5 | 6 | annos=os.listdir(ann_dir) 7 | class_names = ('一次性快餐盒','书籍纸张','充电宝', 8 | '剩饭剩菜' ,'包','垃圾桶','塑料器皿','塑料玩具','塑料衣架', 9 | '大骨头','干电池','快递纸袋','插头电线','旧衣服','易拉罐', 10 | '枕头','果皮果肉','毛绒玩具','污损塑料','污损用纸','洗护用品', 11 | '烟蒂','牙签','玻璃器皿','砧板','筷子','纸盒纸箱','花盆', 12 | '茶叶渣','菜帮菜叶','蛋壳','调料瓶','软膏','过期药物', 13 | '酒瓶','金属厨具','金属器皿','金属食品罐','锅','陶瓷器皿', 14 | '鞋','食用油桶','饮料瓶','鱼骨') 15 | class_dict = {class_name: 0 for i, class_name in enumerate(class_names)} 16 | for ann in annos: 17 | annotation_file = os.path.join(ann_dir,ann) 18 | objects = ET.parse(annotation_file).findall("object") 19 | for obj in objects: 20 | class_name = obj.find('name').text.lower().strip() 21 | class_dict[class_name]+=1 22 | print(class_dict) 23 | print(sorted(class_dict.items(),key=lambda item:item[1],reverse=True)) 24 | ''' 25 | ratio={'一次性快餐盒': 329, '书籍纸张': 196, '充电宝': 410, '剩饭剩菜': 1416, '包': 442, '垃圾桶': 108, '塑料器皿': 666, '塑料玩具': 439, 26 | '塑料衣架': 325, '大骨头': 413, '干电池': 445, '快递纸袋': 248, '插头电线': 1277, '旧衣服': 394, '易拉罐': 414, '枕头': 382, '果皮果肉': 1052, 27 | '毛绒玩具': 655, '污损塑料': 1237, '污损用纸': 227, '洗护用品': 1030, '烟蒂': 433, '牙签': 156, '玻璃器皿': 878, '砧板': 466, '筷子': 815, 28 | '纸盒纸箱': 623, '花盆': 211, '茶叶渣': 409, '菜帮菜叶': 984, '蛋壳': 458, '调料瓶': 543, '软膏': 418, '过期药物': 710, '酒瓶': 346, 29 | '金属厨具': 284, '金属器皿': 341, '金属食品罐': 400, '锅': 531, '陶瓷器皿': 2662, '鞋': 463, '食用油桶': 446, '饮料瓶': 371, '鱼骨': 525 30 | } 31 | sort_ratio=[('陶瓷器皿', 2662), ('剩饭剩菜', 1416), ('插头电线', 1277), ('污损塑料', 1237), ('果皮果肉', 1052), 32 | ('洗护用品', 1030), ('菜帮菜叶', 984), ('玻璃器皿', 878), ('筷子', 815), ('过期药物', 710), 33 | ('塑料器皿', 666), ('毛绒玩具', 655), ('纸盒纸箱', 623), ('调料瓶', 543), ('锅', 531), 34 | ('鱼骨', 525), ('砧板', 466), ('鞋', 463), ('蛋壳', 458), ('食用油桶', 446), 35 | ('干电池', 445), ('包', 442), ('塑料玩具', 439), ('烟蒂', 433), ('软膏', 418), 36 | ('易拉罐', 414), ('大骨头', 413), ('充电宝', 410), ('茶叶渣', 409), ('金属食品罐', 400), 37 | ('旧衣服', 394), ('枕头', 382), ('饮料瓶', 371), ('酒瓶', 346), ('金属器皿', 341), 38 | ('一次性快餐盒', 329), ('塑料衣架', 325), ('金属厨具', 284), ('快递纸袋', 248), ('污损用纸', 227), 39 | ('花盆', 211), ('书籍纸张', 196), ('牙签', 156), ('垃圾桶', 108)] 40 | 41 | ratio=sorted(ratio.items(),key=lambda item:item[1],reverse=True) 42 | aug_need=[] 43 | for i in range(len(ratio)): 44 | if ratio[i][1]>1000: 45 | aug_need.append(ratio[i][0]) 46 | print(ratio) 47 | print(aug_need) 48 | ''' -------------------------------------------------------------------------------- /get_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | dir='./datasets/VOC2007/ImageSets/Main/test.txt' 4 | r=open(dir,'r') 5 | img_dir='datasets/VOC2007/JPEGImages' 6 | w_dir='demo/test_trash' 7 | if not os.path.exists(w_dir): 8 | os.makedirs(w_dir) 9 | for line in r.readlines(): 10 | line=line.strip('\n') 11 | shutil.copy(os.path.join(img_dir,line+'.jpg'),os.path.join(w_dir,line+'.jpg')) 12 | -------------------------------------------------------------------------------- /infer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import glob 3 | import os 4 | import time 5 | import cv2 6 | 7 | import random 8 | import torch 9 | from PIL import Image, ImageFont, ImageDraw 10 | from vizer.draw import draw_boxes 11 | 12 | from ssd.config import cfg 13 | from ssd.data.datasets import COCODataset, VOCDataset 14 | import argparse 15 | import numpy as np 16 | 17 | from ssd.data.transforms import build_transforms 18 | from ssd.modeling.detector import build_detection_model 19 | from ssd.utils import mkdir 20 | from ssd.utils.checkpoint import CheckPointer 21 | import pdb 22 | 23 | @torch.no_grad() 24 | def run_demo(cfg, ckpt, score_threshold, images_dir, output_dir, dataset_type): 25 | if dataset_type == "voc": 26 | class_names = VOCDataset.class_names 27 | elif dataset_type == 'coco': 28 | class_names = COCODataset.class_names 29 | else: 30 | raise NotImplementedError('Not implemented now.') 31 | device = torch.device(cfg.MODEL.DEVICE) 32 | smoke_name_dic = ('__background__','一次性快餐盒','书籍纸张','充电宝', 33 | '剩饭剩菜' ,'包','垃圾桶','塑料器皿','塑料玩具','塑料衣架', 34 | '大骨头','干电池','快递纸袋','插头电线','旧衣服','易拉罐', 35 | '枕头','果皮果肉','毛绒玩具','污损塑料','污损用纸','洗护用品', 36 | '烟蒂','牙签','玻璃器皿','砧板','筷子','纸盒纸箱','花盆', 37 | '茶叶渣','菜帮菜叶','蛋壳','调料瓶','软膏','过期药物', 38 | '酒瓶','金属厨具','金属器皿','金属食品罐','锅','陶瓷器皿', 39 | '鞋','食用油桶','饮料瓶','鱼骨') 40 | 41 | model = build_detection_model(cfg) 42 | cpu_device = torch.device("cpu") 43 | model = model.to(device) 44 | checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR) 45 | checkpointer.load(ckpt, use_latest=ckpt is None) 46 | weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file() 47 | print('Loaded weights from {}'.format(weight_file)) 48 | 49 | image_paths = glob.glob(os.path.join(images_dir, '*.jpg')) 50 | mkdir(output_dir) 51 | 52 | transforms = build_transforms(cfg, is_train=False) 53 | model.eval() 54 | miss = 0 55 | 56 | for i, image_path in enumerate(image_paths): 57 | start = time.time() 58 | image_name = os.path.basename(image_path) 59 | cv_image = cv2.imread(image_path) 60 | PIL_image = Image.open(image_path) 61 | 62 | image = np.array(Image.open(image_path).convert("RGB")) 63 | height, width = image.shape[:2] 64 | images = transforms(image)[0].unsqueeze(0) 65 | load_time = time.time() - start 66 | 67 | start = time.time() 68 | result = model(images.to(device))[0] 69 | inference_time = time.time() - start 70 | 71 | result = result.resize((width, height)).to(cpu_device).numpy() 72 | boxes, labels, scores = result['boxes'], result['labels'], result['scores'] 73 | 74 | indices = scores > score_threshold 75 | boxes = boxes[indices] 76 | labels = labels[indices] 77 | scores = scores[indices] 78 | 79 | miss = miss + (1 - len(boxes)) 80 | meters = ' | '.join( 81 | [ 82 | 'objects {:02d}'.format(len(boxes)), 83 | 'load {:03d}ms'.format(round(load_time * 1000)), 84 | 'inference {:03d}ms'.format(round(inference_time * 1000)), 85 | 'FPS {}'.format(round(1.0 / inference_time)) 86 | ] 87 | ) 88 | print('({:04d}/{:04d}) {}: {}'.format(i + 1, len(image_paths), image_name, meters)) 89 | 90 | draw_ = ImageDraw.Draw(PIL_image) 91 | for c in range(len(scores)): 92 | text = smoke_name_dic[labels[c]] 93 | font = ImageFont.truetype('/usr/share/fonts/truetype/arphic/uming.ttc', 40) 94 | draw_.text((int(boxes[c][0]) + 2, int(boxes[c][1]) - 2), text, (255, 0, 0), font=font) 95 | 96 | cv_image = cv2.cvtColor(np.asarray(PIL_image), cv2.COLOR_RGB2BGR) 97 | for c in range(len(scores)): 98 | cv2.rectangle(cv_image, (int(boxes[c][0]), int(boxes[c][1])), (int(boxes[c][2]), int(boxes[c][3])), 99 | (0, 0, 255), 4) 100 | cv2.imwrite(os.path.join(output_dir, image_name), cv_image) 101 | smoke_count = len(image_paths) 102 | print("出现：%d 漏掉： %d 漏检率：%.2f" % (smoke_count, miss, miss / smoke_count)) 103 | # print(len(label_list)) 104 | 105 | # print("漏检率: %.3f"%(miss/smoke_count)) 106 | 107 | 108 | def main(): 109 | parser = argparse.ArgumentParser(description='SSD demo.') 110 | parser.add_argument( 111 | "--config-file", 112 | default="configs/resnet50_224.yaml", # "configs/vgg_ssd300_voc0712.yaml", 113 | metavar="FILE", 114 | help="path to config file", 115 | type=str, 116 | ) 117 | parser.add_argument("--ckpt", type=str, default=None, help="Trained weights.") 118 | parser.add_argument("--score_threshold", type=float, default=0.5) 119 | parser.add_argument("--images_dir", default='demo/test_trash/', type=str, 120 | help='Specify a image dir to do prediction.') 121 | parser.add_argument("--output_dir", default='demo/result/5.18/', type=str, 122 | help='Specify a image dir to save predicted images.') 123 | parser.add_argument("--dataset_type", default="voc", type=str, 124 | help='Specify dataset type. Currently support voc and coco.') 125 | 126 | parser.add_argument( 127 | "opts", 128 | help="Modify config options using the command-line", 129 | default=None, 130 | nargs=argparse.REMAINDER, 131 | ) 132 | args = parser.parse_args() 133 | print(args) 134 | 135 | cfg.merge_from_file(args.config_file) 136 | cfg.merge_from_list(args.opts) 137 | cfg.freeze() 138 | 139 | print("Loaded configuration file {}".format(args.config_file)) 140 | with open(args.config_file, "r") as cf: 141 | config_str = "\n" + cf.read() 142 | print(config_str) 143 | print("Running with config:\n{}".format(cfg)) 144 | 145 | run_demo(cfg=cfg, 146 | ckpt=args.ckpt, 147 | score_threshold=args.score_threshold, 148 | images_dir=args.images_dir, 149 | output_dir=args.output_dir, 150 | dataset_type=args.dataset_type) 151 | 152 | 153 | if __name__ == '__main__': 154 | main() 155 | -------------------------------------------------------------------------------- /make_imagesets.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | 6 | root_path = 'datasets/VOC2007' 7 | 8 | xmlfilepath = root_path + '/Annotations' 9 | 10 | txtsavepath = root_path + '/ImageSets/Main' 11 | 12 | 13 | if not os.path.exists(txtsavepath): 14 | os.makedirs(txtsavepath) 15 | 16 | trainval_percent = 1 17 | train_percent = 1 18 | total_xml = os.listdir(xmlfilepath) 19 | num = len(total_xml) 20 | list = range(num) 21 | tv = int(num * trainval_percent) 22 | tr = int(tv * train_percent) 23 | trainval = random.sample(list, tv) 24 | train = random.sample(trainval, tr) 25 | 26 | print("train and val size:", tv) 27 | print("train size:", tr) 28 | 29 | ftrainval = open(txtsavepath + '/trainval.txt', 'w') 30 | ftest = open(txtsavepath + '/test.txt', 'w') 31 | ftrain = open(txtsavepath + '/train.txt', 'w') 32 | fval = open(txtsavepath + '/val.txt', 'w') 33 | 34 | for i in list: 35 | name = total_xml[i][:-4] + '\n' 36 | if i in trainval: 37 | ftrainval.write(name) 38 | if i in train: 39 | ftrain.write(name) 40 | else: 41 | fval.write(name) 42 | else: 43 | ftest.write(name) 44 | 45 | ftrainval.close() 46 | ftrain.close() 47 | fval.close() 48 | ftest.close() 49 | -------------------------------------------------------------------------------- /save_json.py: -------------------------------------------------------------------------------- 1 | # -*- encoding: utf-8 -*- 2 | """ 3 | @File : save_json.py 4 | @Time : 2020/5/4 18:08 5 | @Author : Alessia K 6 | @Email : ------ 7 | """ 8 | import json 9 | import numpy as np 10 | import os 11 | import csv 12 | 13 | 14 | def label_list(path): 15 | with open(path, 'r', encoding='utf-8') as f: 16 | class_csv = csv.reader(f, delimiter=',') 17 | result = {} 18 | 19 | for line, row in enumerate(class_csv): 20 | line += 1 21 | 22 | try: 23 | class_name, class_id = row 24 | except ValueError: 25 | # raise_from(ValueError('line {}: format should be \'class_name,class_id\''.format(line)), None) 26 | raise (ValueError('line {}: format should be \'class_name,class_id\''.format(line)), None) 27 | 28 | if class_name in result: 29 | raise ValueError('line {}: duplicate class name: \'{}\''.format(line, class_name)) 30 | result[int(class_id)] = class_name 31 | return result 32 | 33 | def create_class_dict(path): 34 | with open(os.path.join(path, 'classify_rule.json'), 'r', encoding='utf-8') as f: 35 | classify_rule = json.load(f) 36 | class_name = {} 37 | for keys, vals in classify_rule.items(): 38 | for i in range(len(vals)): 39 | class_name[str(vals[i])]=str(keys) 40 | 41 | return class_name 42 | 43 | 44 | def get_classes_name(class_name, class_name_i): 45 | return class_name[class_name_i] + '/' + class_name_i 46 | 47 | 48 | def save_result_as_json(img_name, classes, scores, bboxes, time): 49 | """ 50 | { 51 | "detection_classes": [] 52 | "detection_scores": [] (.4f) 53 | "detection_bboxes": [] (xmin、ymin、xmax、ymax) (.1f) 54 | "latency_time": "" (str(.1f)) 55 | } 56 | """ 57 | 58 | scores = np.around(scores.astype(np.float), decimals=4).tolist() 59 | bboxes = np.around(bboxes.astype(np.float), decimals=1).tolist() 60 | print(bboxes) 61 | 62 | 63 | save_file = {} 64 | save_file["detection_classes"] = classes 65 | save_file["detection_scores"] = scores 66 | save_file["detection_boxes"] = bboxes 67 | save_file["latency_time"] = "{:.1f} ms".format(time) 68 | 69 | return save_file 70 | 71 | # save_name = os.path.join('data', os.path.basename(img_name).replace('.jpg', '.json')) 72 | # with open(save_name, 'w+', encoding='utf-8')as f: 73 | # json.dump(save_file, f, indent=4, ensure_ascii=False) 74 | 75 | 76 | if __name__ == '__main__': 77 | class_name = create_class_dict(r'D:\Work\ohter\SodicData\train_val') 78 | print(class_name) 79 | labels = label_list('data/class_name.csv') 80 | print(labels) -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | with open("README.md", "r") as fh: 4 | long_description = fh.read() 5 | 6 | setup( 7 | name="torch-ssd", 8 | version="1.2.0", 9 | packages=find_packages(exclude=['ext']), 10 | install_requires=[ 11 | "torch~=1.0", 12 | "torchvision~=0.3", 13 | "opencv-python~=4.0", 14 | "yacs==0.1.6", 15 | "Vizer~=0.1.4", 16 | ], 17 | author="Congcong Li", 18 | author_email="luffy.lcc@gmail.com", 19 | description="High quality, fast, modular reference implementation of SSD in PyTorch", 20 | long_description=long_description, 21 | long_description_content_type="text/markdown", 22 | url="https://github.com/lufficc/SSD", 23 | classifiers=[ 24 | "Programming Language :: Python :: 3", 25 | "License :: OSI Approved :: MIT License", 26 | "Operating System :: OS Independent", 27 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 28 | ], 29 | license="MIT", 30 | python_requires=">=3.6", 31 | include_package_data=True, 32 | ) 33 | -------------------------------------------------------------------------------- /ssd/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/__init__.py -------------------------------------------------------------------------------- /ssd/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/config/__init__.py: -------------------------------------------------------------------------------- 1 | from .defaults import _C as cfg 2 | -------------------------------------------------------------------------------- /ssd/config/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/config/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/config/__pycache__/defaults.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/config/__pycache__/defaults.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/config/__pycache__/path_catlog.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/config/__pycache__/path_catlog.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/config/defaults.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode as CN 2 | 3 | _C = CN() 4 | 5 | _C.MODEL = CN() 6 | _C.MODEL.META_ARCHITECTURE = 'SSDDetector' 7 | _C.MODEL.DEVICE = "cuda" 8 | # match default boxes to any ground truth with jaccard overlap higher than a threshold (0.5) 9 | _C.MODEL.THRESHOLD = 0.5 10 | _C.MODEL.NUM_CLASSES = 45 11 | # Hard negative mining 12 | _C.MODEL.NEG_POS_RATIO = 3 13 | _C.MODEL.CENTER_VARIANCE = 0.1 14 | _C.MODEL.SIZE_VARIANCE = 0.2 15 | 16 | # ---------------------------------------------------------------------------- # 17 | # Backbone 18 | # ---------------------------------------------------------------------------- # 19 | _C.MODEL.BACKBONE = CN() 20 | _C.MODEL.BACKBONE.NAME = 'efficient_net-b3'#'efficient_net-b3' 21 | _C.MODEL.BACKBONE.OUT_CHANNELS = (512, 1024, 512, 256, 256, 256) 22 | _C.MODEL.BACKBONE.PRETRAINED = True 23 | 24 | # ----------------------------------------------------------------------------- 25 | # PRIORS 26 | # ----------------------------------------------------------------------------- 27 | _C.MODEL.PRIORS = CN() 28 | _C.MODEL.PRIORS.FEATURE_MAPS = [38, 19, 10, 5, 3, 1] 29 | _C.MODEL.PRIORS.STRIDES = [8, 16, 32, 64, 100, 300] 30 | _C.MODEL.PRIORS.MIN_SIZES = [30, 60, 111, 162, 213, 264] 31 | _C.MODEL.PRIORS.MAX_SIZES = [60, 111, 162, 213, 264, 315] 32 | _C.MODEL.PRIORS.ASPECT_RATIOS = [[2], [2, 3], [2, 3], [2, 3], [2], [2]] 33 | # When has 1 aspect ratio, every location has 4 boxes, 2 ratio 6 boxes. 34 | # #boxes = 2 + #ratio * 2 35 | _C.MODEL.PRIORS.BOXES_PER_LOCATION = [4, 6, 6, 6, 4, 4] # number of boxes per feature map location 36 | _C.MODEL.PRIORS.CLIP = True 37 | 38 | # ----------------------------------------------------------------------------- 39 | # Box Head 40 | # ----------------------------------------------------------------------------- 41 | _C.MODEL.BOX_HEAD = CN() 42 | _C.MODEL.BOX_HEAD.NAME = 'SSDBoxHead' 43 | _C.MODEL.BOX_HEAD.PREDICTOR = 'SSDBoxPredictor' 44 | 45 | # ----------------------------------------------------------------------------- 46 | # INPUT 47 | # ----------------------------------------------------------------------------- 48 | _C.INPUT = CN() 49 | # Image size 50 | _C.INPUT.IMAGE_SIZE = 300 51 | # Values to be used for image normalization, RGB layout 52 | _C.INPUT.PIXEL_MEAN = [123, 117, 104] 53 | 54 | # ----------------------------------------------------------------------------- 55 | # Dataset 56 | # ----------------------------------------------------------------------------- 57 | _C.DATASETS = CN() 58 | # List of the dataset names for training, as present in paths_catalog.py 59 | _C.DATASETS.TRAIN = () 60 | # List of the dataset names for testing, as present in paths_catalog.py 61 | _C.DATASETS.TEST = () 62 | 63 | # ----------------------------------------------------------------------------- 64 | # DataLoader 65 | # ----------------------------------------------------------------------------- 66 | _C.DATA_LOADER = CN() 67 | # Number of data loading threads 68 | _C.DATA_LOADER.NUM_WORKERS = 8 69 | _C.DATA_LOADER.PIN_MEMORY = True 70 | 71 | # ---------------------------------------------------------------------------- # 72 | # Solver 73 | # ---------------------------------------------------------------------------- # 74 | _C.SOLVER = CN() 75 | # train configs 76 | _C.SOLVER.MAX_ITER = 120000 77 | _C.SOLVER.LR_STEPS = [80000, 100000] 78 | _C.SOLVER.GAMMA = 0.1 79 | _C.SOLVER.BATCH_SIZE = 32 80 | _C.SOLVER.LR = 1e-3 81 | _C.SOLVER.MOMENTUM = 0.9 82 | _C.SOLVER.WEIGHT_DECAY = 5e-4 83 | _C.SOLVER.WARMUP_FACTOR = 1.0 / 3 84 | _C.SOLVER.WARMUP_ITERS = 500 85 | 86 | # ---------------------------------------------------------------------------- # 87 | # Specific test options 88 | # ---------------------------------------------------------------------------- # 89 | _C.TEST = CN() 90 | _C.TEST.NMS_THRESHOLD = 0.45 91 | _C.TEST.CONFIDENCE_THRESHOLD = 0.01 92 | _C.TEST.MAX_PER_CLASS = -1 93 | _C.TEST.MAX_PER_IMAGE = 100 94 | _C.TEST.BATCH_SIZE = 10 95 | 96 | _C.OUTPUT_DIR = 'outputs' 97 | -------------------------------------------------------------------------------- /ssd/config/path_catlog.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | class DatasetCatalog: 5 | DATA_DIR = 'datasets' 6 | DATASETS = { 7 | 'voc_2007_train': { 8 | "data_dir": "VOC2007", 9 | "split": "train" 10 | }, 11 | 'voc_2007_val': { 12 | "data_dir": "VOC2007", 13 | "split": "val" 14 | }, 15 | 'voc_2007_trainval': { 16 | "data_dir": "VOC2007", 17 | "split": "trainval" 18 | }, 19 | 'voc_2007_test': { 20 | "data_dir": "VOC2007", 21 | "split": "test" 22 | }, 23 | 'voc_2012_train': { 24 | "data_dir": "VOC2012", 25 | "split": "train" 26 | }, 27 | 'voc_2012_val': { 28 | "data_dir": "VOC2012", 29 | "split": "val" 30 | }, 31 | 'voc_2012_trainval': { 32 | "data_dir": "VOC2012", 33 | "split": "trainval" 34 | }, 35 | 'voc_2012_test': { 36 | "data_dir": "VOC2012", 37 | "split": "test" 38 | }, 39 | 'coco_2014_valminusminival': { 40 | "data_dir": "val2014", 41 | "ann_file": "annotations/instances_valminusminival2014.json" 42 | }, 43 | 'coco_2014_minival': { 44 | "data_dir": "val2014", 45 | "ann_file": "annotations/instances_minival2014.json" 46 | }, 47 | 'coco_2014_train': { 48 | "data_dir": "train2014", 49 | "ann_file": "annotations/instances_train2014.json" 50 | }, 51 | 'coco_2014_val': { 52 | "data_dir": "val2014", 53 | "ann_file": "annotations/instances_val2014.json" 54 | }, 55 | } 56 | 57 | @staticmethod 58 | def get(name): 59 | if "voc" in name: 60 | voc_root = DatasetCatalog.DATA_DIR 61 | if 'VOC_ROOT' in os.environ: 62 | voc_root = os.environ['VOC_ROOT'] 63 | 64 | attrs = DatasetCatalog.DATASETS[name] 65 | args = dict( 66 | data_dir=os.path.join(voc_root, attrs["data_dir"]), 67 | split=attrs["split"], 68 | ) 69 | return dict(factory="VOCDataset", args=args) 70 | elif "coco" in name: 71 | coco_root = DatasetCatalog.DATA_DIR 72 | if 'COCO_ROOT' in os.environ: 73 | coco_root = os.environ['COCO_ROOT'] 74 | 75 | attrs = DatasetCatalog.DATASETS[name] 76 | args = dict( 77 | data_dir=os.path.join(coco_root, attrs["data_dir"]), 78 | ann_file=os.path.join(coco_root, attrs["ann_file"]), 79 | ) 80 | return dict(factory="COCODataset", args=args) 81 | 82 | raise RuntimeError("Dataset not available: {}".format(name)) 83 | -------------------------------------------------------------------------------- /ssd/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/__init__.py -------------------------------------------------------------------------------- /ssd/data/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/__pycache__/build.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/__pycache__/build.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/build.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | from torch.utils.data.dataloader import default_collate 4 | 5 | from ssd.data import samplers 6 | from ssd.data.datasets import build_dataset 7 | from ssd.data.transforms import build_transforms, build_target_transform 8 | from ssd.structures.container import Container 9 | 10 | 11 | class BatchCollator: 12 | def __init__(self, is_train=True): 13 | self.is_train = is_train 14 | 15 | def __call__(self, batch): 16 | transposed_batch = list(zip(*batch)) 17 | images = default_collate(transposed_batch[0]) 18 | img_ids = default_collate(transposed_batch[2]) 19 | 20 | if self.is_train: 21 | list_targets = transposed_batch[1] 22 | targets = Container( 23 | {key: default_collate([d[key] for d in list_targets]) for key in list_targets[0]} 24 | ) 25 | else: 26 | targets = None 27 | return images, targets, img_ids 28 | 29 | 30 | def make_data_loader(cfg, is_train=True, distributed=False, max_iter=None, start_iter=0): 31 | train_transform = build_transforms(cfg, is_train=is_train) 32 | target_transform = build_target_transform(cfg) if is_train else None 33 | dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST 34 | datasets = build_dataset(dataset_list, transform=train_transform, target_transform=target_transform, is_train=is_train) 35 | 36 | shuffle = is_train or distributed 37 | 38 | data_loaders = [] 39 | 40 | for dataset in datasets: 41 | if distributed: 42 | sampler = samplers.DistributedSampler(dataset, shuffle=shuffle) 43 | elif shuffle: 44 | sampler = torch.utils.data.RandomSampler(dataset) 45 | else: 46 | sampler = torch.utils.data.sampler.SequentialSampler(dataset) 47 | 48 | batch_size = cfg.SOLVER.BATCH_SIZE if is_train else cfg.TEST.BATCH_SIZE 49 | batch_sampler = torch.utils.data.sampler.BatchSampler(sampler=sampler, batch_size=batch_size, drop_last=False) 50 | if max_iter is not None: 51 | batch_sampler = samplers.IterationBasedBatchSampler(batch_sampler, num_iterations=max_iter, start_iter=start_iter) 52 | 53 | #data_loader = DataLoader(dataset, num_workers=cfg.DATA_LOADER.NUM_WORKERS, batch_sampler=batch_sampler, 54 | #pin_memory=cfg.DATA_LOADER.PIN_MEMORY, collate_fn=BatchCollator(is_train)) 55 | data_loader = DataLoader(dataset, batch_sampler=batch_sampler, 56 | pin_memory=cfg.DATA_LOADER.PIN_MEMORY, collate_fn=BatchCollator(is_train)) 57 | data_loaders.append(data_loader) 58 | 59 | if is_train: 60 | # during training, a single (possibly concatenated) data_loader is returned 61 | assert len(data_loaders) == 1 62 | return data_loaders[0] 63 | return data_loaders 64 | -------------------------------------------------------------------------------- /ssd/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import ConcatDataset 2 | 3 | from ssd.config.path_catlog import DatasetCatalog 4 | from .voc import VOCDataset 5 | from .coco import COCODataset 6 | 7 | _DATASETS = { 8 | 'VOCDataset': VOCDataset, 9 | 'COCODataset': COCODataset, 10 | } 11 | 12 | 13 | def build_dataset(dataset_list, transform=None, target_transform=None, is_train=True): 14 | assert len(dataset_list) > 0 15 | datasets = [] 16 | for dataset_name in dataset_list: 17 | data = DatasetCatalog.get(dataset_name) 18 | args = data['args'] 19 | factory = _DATASETS[data['factory']] 20 | args['transform'] = transform 21 | args['target_transform'] = target_transform 22 | if factory == VOCDataset: 23 | args['keep_difficult'] = not is_train 24 | elif factory == COCODataset: 25 | args['remove_empty'] = is_train 26 | dataset = factory(**args) 27 | datasets.append(dataset) 28 | # for testing, return a list of datasets 29 | if not is_train: 30 | return datasets 31 | dataset = datasets[0] 32 | if len(datasets) > 1: 33 | dataset = ConcatDataset(datasets) 34 | 35 | return [dataset] 36 | -------------------------------------------------------------------------------- /ssd/data/datasets/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/__pycache__/coco.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/__pycache__/coco.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/__pycache__/voc.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/__pycache__/voc.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch.utils.data 3 | import numpy as np 4 | from PIL import Image 5 | 6 | from ssd.structures.container import Container 7 | 8 | 9 | class COCODataset(torch.utils.data.Dataset): 10 | class_names = ('__background__', 11 | 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 12 | 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 13 | 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 14 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 15 | 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 16 | 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 17 | 'kite', 'baseball bat', 'baseball glove', 'skateboard', 18 | 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 19 | 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 20 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 21 | 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 22 | 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 23 | 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 24 | 'refrigerator', 'book', 'clock', 'vase', 'scissors', 25 | 'teddy bear', 'hair drier', 'toothbrush') 26 | 27 | def __init__(self, data_dir, ann_file, transform=None, target_transform=None, remove_empty=False): 28 | from pycocotools.coco import COCO 29 | self.coco = COCO(ann_file) 30 | self.data_dir = data_dir 31 | self.transform = transform 32 | self.target_transform = target_transform 33 | self.remove_empty = remove_empty 34 | if self.remove_empty: 35 | # when training, images without annotations are removed. 36 | self.ids = list(self.coco.imgToAnns.keys()) 37 | else: 38 | # when testing, all images used. 39 | self.ids = list(self.coco.imgs.keys()) 40 | coco_categories = sorted(self.coco.getCatIds()) 41 | self.coco_id_to_contiguous_id = {coco_id: i + 1 for i, coco_id in enumerate(coco_categories)} 42 | self.contiguous_id_to_coco_id = {v: k for k, v in self.coco_id_to_contiguous_id.items()} 43 | 44 | def __getitem__(self, index): 45 | image_id = self.ids[index] 46 | boxes, labels = self._get_annotation(image_id) 47 | image = self._read_image(image_id) 48 | if self.transform: 49 | image, boxes, labels = self.transform(image, boxes, labels) 50 | if self.target_transform: 51 | boxes, labels = self.target_transform(boxes, labels) 52 | targets = Container( 53 | boxes=boxes, 54 | labels=labels, 55 | ) 56 | return image, targets, index 57 | 58 | def get_annotation(self, index): 59 | image_id = self.ids[index] 60 | return image_id, self._get_annotation(image_id) 61 | 62 | def __len__(self): 63 | return len(self.ids) 64 | 65 | def _get_annotation(self, image_id): 66 | ann_ids = self.coco.getAnnIds(imgIds=image_id) 67 | ann = self.coco.loadAnns(ann_ids) 68 | # filter crowd annotations 69 | ann = [obj for obj in ann if obj["iscrowd"] == 0] 70 | boxes = np.array([self._xywh2xyxy(obj["bbox"]) for obj in ann], np.float32).reshape((-1, 4)) 71 | labels = np.array([self.coco_id_to_contiguous_id[obj["category_id"]] for obj in ann], np.int64).reshape((-1,)) 72 | # remove invalid boxes 73 | keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) 74 | boxes = boxes[keep] 75 | labels = labels[keep] 76 | return boxes, labels 77 | 78 | def _xywh2xyxy(self, box): 79 | x1, y1, w, h = box 80 | return [x1, y1, x1 + w, y1 + h] 81 | 82 | def get_img_info(self, index): 83 | image_id = self.ids[index] 84 | img_data = self.coco.imgs[image_id] 85 | return img_data 86 | 87 | def _read_image(self, image_id): 88 | file_name = self.coco.loadImgs(image_id)[0]['file_name'] 89 | image_file = os.path.join(self.data_dir, file_name) 90 | image = Image.open(image_file).convert("RGB") 91 | image = np.array(image) 92 | return image 93 | -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from ssd.data.datasets import VOCDataset, COCODataset 2 | from .coco import coco_evaluation 3 | from .voc import voc_evaluation 4 | 5 | 6 | def evaluate(dataset, predictions, output_dir, **kwargs): 7 | """evaluate dataset using different methods based on dataset type. 8 | Args: 9 | dataset: Dataset object 10 | predictions(list[(boxes, labels, scores)]): Each item in the list represents the 11 | prediction results for one image. And the index should match the dataset index. 12 | output_dir: output folder, to save evaluation files or results. 13 | Returns: 14 | evaluation result 15 | """ 16 | args = dict( 17 | dataset=dataset, predictions=predictions, output_dir=output_dir, **kwargs, 18 | ) 19 | if isinstance(dataset, VOCDataset): 20 | return voc_evaluation(**args) 21 | elif isinstance(dataset, COCODataset): 22 | return coco_evaluation(**args) 23 | else: 24 | raise NotImplementedError 25 | -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/evaluation/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/coco/__init__.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | from datetime import datetime 5 | 6 | 7 | def coco_evaluation(dataset, predictions, output_dir, iteration=None): 8 | coco_results = [] 9 | for i, prediction in enumerate(predictions): 10 | img_info = dataset.get_img_info(i) 11 | prediction = prediction.resize((img_info['width'], img_info['height'])).numpy() 12 | boxes, labels, scores = prediction['boxes'], prediction['labels'], prediction['scores'] 13 | 14 | image_id, annotation = dataset.get_annotation(i) 15 | class_mapper = dataset.contiguous_id_to_coco_id 16 | if labels.shape[0] == 0: 17 | continue 18 | 19 | boxes = boxes.tolist() 20 | labels = labels.tolist() 21 | scores = scores.tolist() 22 | coco_results.extend( 23 | [ 24 | { 25 | "image_id": image_id, 26 | "category_id": class_mapper[labels[k]], 27 | "bbox": [box[0], box[1], box[2] - box[0], box[3] - box[1]], # to xywh format 28 | "score": scores[k], 29 | } 30 | for k, box in enumerate(boxes) 31 | ] 32 | ) 33 | iou_type = 'bbox' 34 | json_result_file = os.path.join(output_dir, iou_type + ".json") 35 | logger = logging.getLogger("SSD.inference") 36 | logger.info('Writing results to {}...'.format(json_result_file)) 37 | with open(json_result_file, "w") as f: 38 | json.dump(coco_results, f) 39 | from pycocotools.cocoeval import COCOeval 40 | coco_gt = dataset.coco 41 | coco_dt = coco_gt.loadRes(json_result_file) 42 | coco_eval = COCOeval(coco_gt, coco_dt, iou_type) 43 | coco_eval.evaluate() 44 | coco_eval.accumulate() 45 | coco_eval.summarize() 46 | 47 | result_strings = [] 48 | keys = ["AP", "AP50", "AP75", "APs", "APm", "APl"] 49 | metrics = {} 50 | for i, key in enumerate(keys): 51 | metrics[key] = coco_eval.stats[i] 52 | logger.info('{:<10}: {}'.format(key, round(coco_eval.stats[i], 3))) 53 | result_strings.append('{:<10}: {}'.format(key, round(coco_eval.stats[i], 3))) 54 | 55 | if iteration is not None: 56 | result_path = os.path.join(output_dir, 'result_{:07d}.txt'.format(iteration)) 57 | else: 58 | result_path = os.path.join(output_dir, 'result_{}.txt'.format(datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))) 59 | with open(result_path, "w") as f: 60 | f.write('\n'.join(result_strings)) 61 | 62 | return dict(metrics=metrics) 63 | -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/coco/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/evaluation/coco/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/voc/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from datetime import datetime 4 | 5 | import numpy as np 6 | 7 | from .eval_detection_voc import eval_detection_voc 8 | 9 | 10 | def voc_evaluation(dataset, predictions, output_dir, iteration=None): 11 | class_names = dataset.class_names 12 | 13 | pred_boxes_list = [] 14 | pred_labels_list = [] 15 | pred_scores_list = [] 16 | gt_boxes_list = [] 17 | gt_labels_list = [] 18 | gt_difficults = [] 19 | 20 | for i in range(len(dataset)): 21 | image_id, annotation = dataset.get_annotation(i) 22 | gt_boxes, gt_labels, is_difficult = annotation 23 | gt_boxes_list.append(gt_boxes) 24 | gt_labels_list.append(gt_labels) 25 | gt_difficults.append(is_difficult.astype(np.bool)) 26 | 27 | img_info = dataset.get_img_info(i) 28 | prediction = predictions[i] 29 | prediction = prediction.resize((img_info['width'], img_info['height'])).numpy() 30 | boxes, labels, scores = prediction['boxes'], prediction['labels'], prediction['scores'] 31 | 32 | pred_boxes_list.append(boxes) 33 | pred_labels_list.append(labels) 34 | pred_scores_list.append(scores) 35 | result = eval_detection_voc(pred_bboxes=pred_boxes_list, 36 | pred_labels=pred_labels_list, 37 | pred_scores=pred_scores_list, 38 | gt_bboxes=gt_boxes_list, 39 | gt_labels=gt_labels_list, 40 | gt_difficults=gt_difficults, 41 | iou_thresh=0.5, 42 | use_07_metric=True) 43 | logger = logging.getLogger("SSD.inference") 44 | result_str = "mAP: {:.4f}\n".format(result["map"]) 45 | metrics = {'mAP': result["map"]} 46 | for i, ap in enumerate(result["ap"]): 47 | if i == 0: # skip background 48 | continue 49 | metrics[class_names[i]] = ap 50 | result_str += "{:<16}: {:.4f}\n".format(class_names[i], ap) 51 | logger.info(result_str) 52 | 53 | if iteration is not None: 54 | result_path = os.path.join(output_dir, 'result_{:07d}.txt'.format(iteration)) 55 | else: 56 | result_path = os.path.join(output_dir, 'result_{}.txt'.format(datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))) 57 | with open(result_path, "w") as f: 58 | f.write(result_str) 59 | 60 | return dict(metrics=metrics) 61 | -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/voc/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/evaluation/voc/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/evaluation/voc/__pycache__/eval_detection_voc.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/datasets/evaluation/voc/__pycache__/eval_detection_voc.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/datasets/voc.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch.utils.data 3 | import numpy as np 4 | import xml.etree.ElementTree as ET 5 | from PIL import Image 6 | 7 | from ssd.structures.container import Container 8 | 9 | 10 | class VOCDataset(torch.utils.data.Dataset): 11 | class_names = ('__background__','一次性快餐盒','书籍纸张','充电宝', 12 | '剩饭剩菜' ,'包','垃圾桶','塑料器皿','塑料玩具','塑料衣架', 13 | '大骨头','干电池','快递纸袋','插头电线','旧衣服','易拉罐', 14 | '枕头','果皮果肉','毛绒玩具','污损塑料','污损用纸','洗护用品', 15 | '烟蒂','牙签','玻璃器皿','砧板','筷子','纸盒纸箱','花盆', 16 | '茶叶渣','菜帮菜叶','蛋壳','调料瓶','软膏','过期药物', 17 | '酒瓶','金属厨具','金属器皿','金属食品罐','锅','陶瓷器皿', 18 | '鞋','食用油桶','饮料瓶','鱼骨') 19 | 20 | def __init__(self, data_dir, split, transform=None, target_transform=None, keep_difficult=False): 21 | """Dataset for VOC data. 22 | Args: 23 | data_dir: the root of the VOC2007 or VOC2012 dataset, the directory contains the following sub-directories: 24 | Annotations, ImageSets, JPEGImages, SegmentationClass, SegmentationObject. 25 | """ 26 | self.data_dir = data_dir 27 | self.split = split 28 | self.transform = transform 29 | self.target_transform = target_transform 30 | image_sets_file = os.path.join(self.data_dir, "ImageSets", "Main", "%s.txt" % self.split) 31 | self.ids = VOCDataset._read_image_ids(image_sets_file) 32 | self.keep_difficult = keep_difficult 33 | 34 | self.class_dict = {class_name: i for i, class_name in enumerate(self.class_names)} 35 | 36 | def __getitem__(self, index): 37 | image_id = self.ids[index] 38 | boxes, labels, is_difficult = self._get_annotation(image_id) 39 | if not self.keep_difficult: 40 | boxes = boxes[is_difficult == 0] 41 | labels = labels[is_difficult == 0] 42 | image = self._read_image(image_id) 43 | if self.transform: 44 | image, boxes, labels = self.transform(image, boxes, labels) 45 | if self.target_transform: 46 | boxes, labels = self.target_transform(boxes, labels) 47 | targets = Container( 48 | boxes=boxes, 49 | labels=labels, 50 | ) 51 | return image, targets, index 52 | 53 | def get_annotation(self, index): 54 | image_id = self.ids[index] 55 | return image_id, self._get_annotation(image_id) 56 | 57 | def __len__(self): 58 | return len(self.ids) 59 | 60 | @staticmethod 61 | def _read_image_ids(image_sets_file): 62 | ids = [] 63 | with open(image_sets_file) as f: 64 | for line in f: 65 | ids.append(line.rstrip()) 66 | return ids 67 | 68 | def _get_annotation(self, image_id): 69 | annotation_file = os.path.join(self.data_dir, "Annotations", "%s.xml" % image_id) 70 | objects = ET.parse(annotation_file).findall("object") 71 | boxes = [] 72 | labels = [] 73 | is_difficult = [] 74 | for obj in objects: 75 | class_name = obj.find('name').text.lower().strip() 76 | bbox = obj.find('bndbox') 77 | # VOC dataset format follows Matlab, in which indexes start from 0 78 | x1 = float(bbox.find('xmin').text) - 1 79 | y1 = float(bbox.find('ymin').text) - 1 80 | x2 = float(bbox.find('xmax').text) - 1 81 | y2 = float(bbox.find('ymax').text) - 1 82 | boxes.append([x1, y1, x2, y2]) 83 | labels.append(self.class_dict[class_name]) 84 | is_difficult_str = obj.find('difficult').text 85 | is_difficult.append(int(is_difficult_str) if is_difficult_str else 0) 86 | 87 | return (np.array(boxes, dtype=np.float32), 88 | np.array(labels, dtype=np.int64), 89 | np.array(is_difficult, dtype=np.uint8)) 90 | 91 | def get_img_info(self, index): 92 | img_id = self.ids[index] 93 | annotation_file = os.path.join(self.data_dir, "Annotations", "%s.xml" % img_id) 94 | anno = ET.parse(annotation_file).getroot() 95 | size = anno.find("size") 96 | im_info = tuple(map(int, (size.find("height").text, size.find("width").text))) 97 | return {"height": im_info[0], "width": im_info[1]} 98 | 99 | def _read_image(self, image_id): 100 | image_file = os.path.join(self.data_dir, "JPEGImages", "%s.jpg" % image_id) 101 | image = Image.open(image_file).convert("RGB") 102 | image = np.array(image) 103 | return image 104 | -------------------------------------------------------------------------------- /ssd/data/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | from .iteration_based_batch_sampler import IterationBasedBatchSampler 2 | from .distributed import DistributedSampler 3 | 4 | __all__ = ['IterationBasedBatchSampler', 'DistributedSampler'] 5 | -------------------------------------------------------------------------------- /ssd/data/samplers/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/samplers/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/samplers/__pycache__/distributed.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/samplers/__pycache__/distributed.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/samplers/__pycache__/iteration_based_batch_sampler.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/samplers/__pycache__/iteration_based_batch_sampler.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/samplers/distributed.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # Code is copy-pasted exactly as in torch.utils.data.distributed. 3 | # FIXME remove this once c10d fixes the bug it has 4 | import math 5 | import torch 6 | import torch.distributed as dist 7 | from torch.utils.data.sampler import Sampler 8 | 9 | 10 | class DistributedSampler(Sampler): 11 | """Sampler that restricts data loading to a subset of the dataset. 12 | It is especially useful in conjunction with 13 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 14 | process can pass a DistributedSampler instance as a DataLoader sampler, 15 | and load a subset of the original dataset that is exclusive to it. 16 | .. note:: 17 | Dataset is assumed to be of constant size. 18 | Arguments: 19 | dataset: Dataset used for sampling. 20 | num_replicas (optional): Number of processes participating in 21 | distributed training. 22 | rank (optional): Rank of the current process within num_replicas. 23 | """ 24 | 25 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 26 | if num_replicas is None: 27 | if not dist.is_available(): 28 | raise RuntimeError("Requires distributed package to be available") 29 | num_replicas = dist.get_world_size() 30 | if rank is None: 31 | if not dist.is_available(): 32 | raise RuntimeError("Requires distributed package to be available") 33 | rank = dist.get_rank() 34 | self.dataset = dataset 35 | self.num_replicas = num_replicas 36 | self.rank = rank 37 | self.epoch = 0 38 | self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) 39 | self.total_size = self.num_samples * self.num_replicas 40 | self.shuffle = shuffle 41 | 42 | def __iter__(self): 43 | if self.shuffle: 44 | # deterministically shuffle based on epoch 45 | g = torch.Generator() 46 | g.manual_seed(self.epoch) 47 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 48 | else: 49 | indices = torch.arange(len(self.dataset)).tolist() 50 | 51 | # add extra samples to make it evenly divisible 52 | indices += indices[: (self.total_size - len(indices))] 53 | assert len(indices) == self.total_size 54 | 55 | # subsample 56 | offset = self.num_samples * self.rank 57 | indices = indices[offset: offset + self.num_samples] 58 | assert len(indices) == self.num_samples 59 | 60 | return iter(indices) 61 | 62 | def __len__(self): 63 | return self.num_samples 64 | 65 | def set_epoch(self, epoch): 66 | self.epoch = epoch 67 | -------------------------------------------------------------------------------- /ssd/data/samplers/iteration_based_batch_sampler.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data.sampler import BatchSampler 2 | 3 | 4 | class IterationBasedBatchSampler(BatchSampler): 5 | """ 6 | Wraps a BatchSampler, re-sampling from it until 7 | a specified number of iterations have been sampled 8 | """ 9 | 10 | def __init__(self, batch_sampler, num_iterations, start_iter=0): 11 | self.batch_sampler = batch_sampler 12 | self.num_iterations = num_iterations 13 | self.start_iter = start_iter 14 | 15 | def __iter__(self): 16 | iteration = self.start_iter 17 | while iteration <= self.num_iterations: 18 | # if the underlying sampler has a set_epoch method, like 19 | # DistributedSampler, used for making each process see 20 | # a different split of the dataset, then set it 21 | if hasattr(self.batch_sampler.sampler, "set_epoch"): 22 | self.batch_sampler.sampler.set_epoch(iteration) 23 | for batch in self.batch_sampler: 24 | iteration += 1 25 | if iteration > self.num_iterations: 26 | break 27 | yield batch 28 | 29 | def __len__(self): 30 | return self.num_iterations 31 | -------------------------------------------------------------------------------- /ssd/data/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from ssd.modeling.anchors.prior_box import PriorBox 2 | from .target_transform import SSDTargetTransform 3 | from .transforms import * 4 | 5 | 6 | def build_transforms(cfg, is_train=True): 7 | if is_train: 8 | transform = [ 9 | ConvertFromInts(), 10 | PhotometricDistort(), 11 | Expand(cfg.INPUT.PIXEL_MEAN), 12 | #PaddingSqure(cfg.INPUT.PIXEL_MEAN), 13 | #RandomSampleCrop(), 14 | RandomMirror(), 15 | ToPercentCoords(), 16 | Resize(cfg.INPUT.IMAGE_SIZE), 17 | SubtractMeans(cfg.INPUT.PIXEL_MEAN), 18 | ToTensor(), 19 | ] 20 | else: 21 | transform = [ 22 | Resize(cfg.INPUT.IMAGE_SIZE), 23 | SubtractMeans(cfg.INPUT.PIXEL_MEAN), 24 | ToTensor() 25 | ] 26 | transform = Compose(transform) 27 | return transform 28 | 29 | 30 | def build_target_transform(cfg): 31 | transform = SSDTargetTransform(PriorBox(cfg)(), 32 | cfg.MODEL.CENTER_VARIANCE, 33 | cfg.MODEL.SIZE_VARIANCE, 34 | cfg.MODEL.THRESHOLD) 35 | return transform 36 | -------------------------------------------------------------------------------- /ssd/data/transforms/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/transforms/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/transforms/__pycache__/target_transform.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/transforms/__pycache__/target_transform.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/transforms/__pycache__/transforms.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/data/transforms/__pycache__/transforms.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/data/transforms/target_transform.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from ssd.utils import box_utils 5 | 6 | 7 | class SSDTargetTransform: 8 | def __init__(self, center_form_priors, center_variance, size_variance, iou_threshold): 9 | self.center_form_priors = center_form_priors 10 | self.corner_form_priors = box_utils.center_form_to_corner_form(center_form_priors) 11 | self.center_variance = center_variance 12 | self.size_variance = size_variance 13 | self.iou_threshold = iou_threshold 14 | 15 | def __call__(self, gt_boxes, gt_labels): 16 | if type(gt_boxes) is np.ndarray: 17 | gt_boxes = torch.from_numpy(gt_boxes) 18 | if type(gt_labels) is np.ndarray: 19 | gt_labels = torch.from_numpy(gt_labels) 20 | boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels, 21 | self.corner_form_priors, self.iou_threshold) 22 | boxes = box_utils.corner_form_to_center_form(boxes) 23 | locations = box_utils.convert_boxes_to_locations(boxes, self.center_form_priors, self.center_variance, self.size_variance) 24 | 25 | return locations, labels 26 | -------------------------------------------------------------------------------- /ssd/engine/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/engine/__init__.py -------------------------------------------------------------------------------- /ssd/engine/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/engine/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/engine/__pycache__/inference.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/engine/__pycache__/inference.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/engine/__pycache__/trainer.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/engine/__pycache__/trainer.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/engine/inference.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | 4 | import torch 5 | import torch.utils.data 6 | from tqdm import tqdm 7 | 8 | from ssd.data.build import make_data_loader 9 | from ssd.data.datasets.evaluation import evaluate 10 | 11 | from ssd.utils import dist_util, mkdir 12 | from ssd.utils.dist_util import synchronize, is_main_process 13 | 14 | 15 | def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu): 16 | all_predictions = dist_util.all_gather(predictions_per_gpu) 17 | if not dist_util.is_main_process(): 18 | return 19 | # merge the list of dicts 20 | predictions = {} 21 | for p in all_predictions: 22 | predictions.update(p) 23 | # convert a dict where the key is the index in a list 24 | image_ids = list(sorted(predictions.keys())) 25 | if len(image_ids) != image_ids[-1] + 1: 26 | logger = logging.getLogger("SSD.inference") 27 | logger.warning( 28 | "Number of images that were gathered from multiple processes is not " 29 | "a contiguous set. Some images might be missing from the evaluation" 30 | ) 31 | 32 | # convert to a list 33 | predictions = [predictions[i] for i in image_ids] 34 | return predictions 35 | 36 | 37 | def compute_on_dataset(model, data_loader, device): 38 | results_dict = {} 39 | for batch in tqdm(data_loader): 40 | images, targets, image_ids = batch 41 | cpu_device = torch.device("cpu") 42 | with torch.no_grad(): 43 | outputs = model(images.to(device)) 44 | 45 | outputs = [o.to(cpu_device) for o in outputs] 46 | results_dict.update( 47 | {img_id: result for img_id, result in zip(image_ids, outputs)} 48 | ) 49 | return results_dict 50 | 51 | 52 | def inference(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs): 53 | dataset = data_loader.dataset 54 | logger = logging.getLogger("SSD.inference") 55 | logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset))) 56 | predictions_path = os.path.join(output_folder, 'predictions.pth') 57 | if use_cached and os.path.exists(predictions_path): 58 | predictions = torch.load(predictions_path, map_location='cpu') 59 | else: 60 | predictions = compute_on_dataset(model, data_loader, device) 61 | synchronize() 62 | predictions = _accumulate_predictions_from_multiple_gpus(predictions) 63 | if not is_main_process(): 64 | return 65 | if output_folder: 66 | torch.save(predictions, predictions_path) 67 | return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs) 68 | 69 | 70 | @torch.no_grad() 71 | def do_evaluation(cfg, model, distributed, **kwargs): 72 | if isinstance(model, torch.nn.parallel.DistributedDataParallel): 73 | model = model.module 74 | model.eval() 75 | device = torch.device(cfg.MODEL.DEVICE) 76 | data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed) 77 | eval_results = [] 78 | for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val): 79 | output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name) 80 | if not os.path.exists(output_folder): 81 | mkdir(output_folder) 82 | eval_result = inference(model, data_loader, dataset_name, device, output_folder, **kwargs) 83 | eval_results.append(eval_result) 84 | return eval_results 85 | -------------------------------------------------------------------------------- /ssd/engine/trainer.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import datetime 3 | import logging 4 | import os 5 | import time 6 | import torch 7 | import torch.distributed as dist 8 | 9 | from ssd.engine.inference import do_evaluation 10 | from ssd.utils import dist_util 11 | from ssd.utils.metric_logger import MetricLogger 12 | 13 | 14 | def write_metric(eval_result, prefix, summary_writer, global_step): 15 | for key in eval_result: 16 | value = eval_result[key] 17 | tag = '{}/{}'.format(prefix, key) 18 | if isinstance(value, collections.Mapping): 19 | write_metric(value, tag, summary_writer, global_step) 20 | else: 21 | summary_writer.add_scalar(tag, value, global_step=global_step) 22 | 23 | 24 | def reduce_loss_dict(loss_dict): 25 | """ 26 | Reduce the loss dictionary from all processes so that process with rank 27 | 0 has the averaged results. Returns a dict with the same fields as 28 | loss_dict, after reduction. 29 | """ 30 | world_size = dist_util.get_world_size() 31 | if world_size < 2: 32 | return loss_dict 33 | with torch.no_grad(): 34 | loss_names = [] 35 | all_losses = [] 36 | for k in sorted(loss_dict.keys()): 37 | loss_names.append(k) 38 | all_losses.append(loss_dict[k]) 39 | all_losses = torch.stack(all_losses, dim=0) 40 | dist.reduce(all_losses, dst=0) 41 | if dist.get_rank() == 0: 42 | # only main process gets accumulated, so only divide by 43 | # world_size in this case 44 | all_losses /= world_size 45 | reduced_losses = {k: v for k, v in zip(loss_names, all_losses)} 46 | return reduced_losses 47 | 48 | 49 | def do_train(cfg, model, 50 | data_loader, 51 | optimizer, 52 | scheduler, 53 | checkpointer, 54 | device, 55 | arguments, 56 | args): 57 | logger = logging.getLogger("SSD.trainer") 58 | logger.info("Start training ...") 59 | meters = MetricLogger() 60 | 61 | model.train() 62 | save_to_disk = dist_util.get_rank() == 0 63 | if args.use_tensorboard and save_to_disk: 64 | import tensorboardX 65 | 66 | summary_writer = tensorboardX.SummaryWriter(log_dir=os.path.join(cfg.OUTPUT_DIR, 'tf_logs')) 67 | else: 68 | summary_writer = None 69 | 70 | max_iter = len(data_loader) 71 | start_iter = arguments["iteration"] 72 | start_training_time = time.time() 73 | end = time.time() 74 | for iteration, (images, targets, _) in enumerate(data_loader, start_iter): 75 | iteration = iteration + 1 76 | arguments["iteration"] = iteration 77 | 78 | images = images.to(device) 79 | targets = targets.to(device) 80 | loss_dict = model(images, targets=targets) 81 | loss = sum(loss for loss in loss_dict.values()) 82 | 83 | # reduce losses over all GPUs for logging purposes 84 | loss_dict_reduced = reduce_loss_dict(loss_dict) 85 | losses_reduced = sum(loss for loss in loss_dict_reduced.values()) 86 | meters.update(total_loss=losses_reduced, **loss_dict_reduced) 87 | 88 | optimizer.zero_grad() 89 | loss.backward() 90 | optimizer.step() 91 | scheduler.step() 92 | 93 | batch_time = time.time() - end 94 | end = time.time() 95 | meters.update(time=batch_time) 96 | if iteration % args.log_step == 0: 97 | eta_seconds = meters.time.global_avg * (max_iter - iteration) 98 | eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) 99 | logger.info( 100 | meters.delimiter.join([ 101 | "iter: {iter:06d}", 102 | "lr: {lr:.5f}", 103 | '{meters}', 104 | "eta: {eta}", 105 | 'mem: {mem}M', 106 | ]).format( 107 | iter=iteration, 108 | lr=optimizer.param_groups[0]['lr'], 109 | meters=str(meters), 110 | eta=eta_string, 111 | mem=round(torch.cuda.max_memory_allocated() / 1024.0 / 1024.0), 112 | ) 113 | ) 114 | if summary_writer: 115 | global_step = iteration 116 | summary_writer.add_scalar('losses/total_loss', losses_reduced, global_step=global_step) 117 | for loss_name, loss_item in loss_dict_reduced.items(): 118 | summary_writer.add_scalar('losses/{}'.format(loss_name), loss_item, global_step=global_step) 119 | summary_writer.add_scalar('lr', optimizer.param_groups[0]['lr'], global_step=global_step) 120 | 121 | if iteration % args.save_step == 0: 122 | checkpointer.save("model_{:06d}".format(iteration), **arguments) 123 | 124 | if args.eval_step > 0 and iteration % args.eval_step == 0 and not iteration == max_iter: 125 | eval_results = do_evaluation(cfg, model, distributed=args.distributed, iteration=iteration) 126 | if dist_util.get_rank() == 0 and summary_writer: 127 | for eval_result, dataset in zip(eval_results, cfg.DATASETS.TEST): 128 | write_metric(eval_result['metrics'], 'metrics/' + dataset, summary_writer, iteration) 129 | model.train() # *IMPORTANT*: change to train mode after eval. 130 | 131 | checkpointer.save("model_final", **arguments) 132 | # compute training time 133 | total_training_time = int(time.time() - start_training_time) 134 | total_time_str = str(datetime.timedelta(seconds=total_training_time)) 135 | logger.info("Total training time: {} ({:.4f} s / it)".format(total_time_str, total_training_time / max_iter)) 136 | return model 137 | -------------------------------------------------------------------------------- /ssd/layers/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.init as init 4 | from .separable_conv import SeparableConv2d 5 | 6 | __all__ = ['L2Norm', 'SeparableConv2d'] 7 | 8 | 9 | class L2Norm(nn.Module): 10 | def __init__(self, n_channels, scale): 11 | super(L2Norm, self).__init__() 12 | self.n_channels = n_channels 13 | self.gamma = scale or None 14 | self.eps = 1e-10 15 | self.weight = nn.Parameter(torch.Tensor(self.n_channels)) 16 | self.reset_parameters() 17 | 18 | def reset_parameters(self): 19 | init.constant_(self.weight, self.gamma) 20 | 21 | def forward(self, x): 22 | norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps 23 | x = torch.div(x, norm) 24 | out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x 25 | return out 26 | -------------------------------------------------------------------------------- /ssd/layers/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/layers/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/layers/__pycache__/separable_conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/layers/__pycache__/separable_conv.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/layers/separable_conv.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | 4 | class SeparableConv2d(nn.Module): 5 | def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False): 6 | super().__init__() 7 | ReLU = nn.ReLU if onnx_compatible else nn.ReLU6 8 | self.conv = nn.Sequential( 9 | nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, 10 | groups=in_channels, stride=stride, padding=padding), 11 | nn.BatchNorm2d(in_channels), 12 | ReLU(), 13 | nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1), 14 | ) 15 | 16 | def forward(self, x): 17 | return self.conv(x) 18 | -------------------------------------------------------------------------------- /ssd/modeling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/__init__.py -------------------------------------------------------------------------------- /ssd/modeling/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/__pycache__/registry.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/__pycache__/registry.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/anchors/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/anchors/__init__.py -------------------------------------------------------------------------------- /ssd/modeling/anchors/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/anchors/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/anchors/__pycache__/prior_box.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/anchors/__pycache__/prior_box.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/anchors/prior_box.py: -------------------------------------------------------------------------------- 1 | from itertools import product 2 | 3 | import torch 4 | from math import sqrt 5 | 6 | 7 | class PriorBox: 8 | def __init__(self, cfg): 9 | self.image_size = cfg.INPUT.IMAGE_SIZE 10 | prior_config = cfg.MODEL.PRIORS 11 | self.feature_maps = prior_config.FEATURE_MAPS 12 | self.min_sizes = prior_config.MIN_SIZES 13 | self.max_sizes = prior_config.MAX_SIZES 14 | self.strides = prior_config.STRIDES 15 | self.aspect_ratios = prior_config.ASPECT_RATIOS 16 | self.clip = prior_config.CLIP 17 | 18 | def __call__(self): 19 | """Generate SSD Prior Boxes. 20 | It returns the center, height and width of the priors. The values are relative to the image size 21 | Returns: 22 | priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values 23 | are relative to the image size. 24 | """ 25 | priors = [] 26 | for k, f in enumerate(self.feature_maps): 27 | scale = self.image_size / self.strides[k] 28 | for i, j in product(range(f), repeat=2): 29 | # unit center x,y 30 | cx = (j + 0.5) / scale 31 | cy = (i + 0.5) / scale 32 | 33 | # small sized square box 34 | size = self.min_sizes[k] 35 | h = w = size / self.image_size 36 | priors.append([cx, cy, w, h]) 37 | 38 | # big sized square box 39 | size = sqrt(self.min_sizes[k] * self.max_sizes[k]) 40 | h = w = size / self.image_size 41 | priors.append([cx, cy, w, h]) 42 | 43 | # change h/w ratio of the small sized box 44 | size = self.min_sizes[k] 45 | h = w = size / self.image_size 46 | for ratio in self.aspect_ratios[k]: 47 | ratio = sqrt(ratio) 48 | priors.append([cx, cy, w * ratio, h / ratio]) 49 | priors.append([cx, cy, w / ratio, h * ratio]) 50 | 51 | priors = torch.tensor(priors) 52 | if self.clip: 53 | priors.clamp_(max=1, min=0) 54 | return priors 55 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | from ssd.modeling import registry 2 | from .vgg import VGG 3 | from .mobilenet import MobileNetV2 4 | from .efficient_net import EfficientNet 5 | from .resnet import res_backbone 6 | __all__ = ['build_backbone', 'VGG', 'MobileNetV2', 'EfficientNet','res_backbone'] 7 | 8 | 9 | def build_backbone(cfg): 10 | return registry.BACKBONES[cfg.MODEL.BACKBONE.NAME](cfg, cfg.MODEL.BACKBONE.PRETRAINED) 11 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/__pycache__/mobilenet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/__pycache__/mobilenet.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/__pycache__/resnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/__pycache__/resnet.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/__pycache__/vgg.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/__pycache__/vgg.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/__init__.py: -------------------------------------------------------------------------------- 1 | from ssd.modeling import registry 2 | from .efficient_net import EfficientNet 3 | 4 | __all__ = ['efficient_net_b3', 'EfficientNet'] 5 | 6 | 7 | @registry.BACKBONES.register('efficient_net-b3') 8 | def efficient_net_b3(cfg, pretrained=True): 9 | if pretrained: 10 | model = EfficientNet.from_pretrained('efficientnet-b3') 11 | else: 12 | model = EfficientNet.from_name('efficientnet-b3') 13 | 14 | return model 15 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/efficient_net/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/__pycache__/efficient_net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/efficient_net/__pycache__/efficient_net.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/backbone/efficient_net/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/efficient_net.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | from .utils import ( 5 | relu_fn, 6 | round_filters, 7 | round_repeats, 8 | drop_connect, 9 | Conv2dSamePadding, 10 | get_model_params, 11 | efficientnet_params, 12 | load_pretrained_weights, 13 | ) 14 | 15 | INDICES = { 16 | 'efficientnet-b3': [7, 17, 25] 17 | } 18 | 19 | EXTRAS = { 20 | 'efficientnet-b3': [ 21 | # in, out, k, s, p 22 | [(384, 128, 1, 1, 0), (128, 256, 3, 2, 1)], # 5 x 5 23 | [(256, 128, 1, 1, 0), (128, 256, 3, 1, 0)], # 3 x 3 24 | [(256, 128, 1, 1, 0), (128, 256, 3, 1, 0)], # 1 x 1 25 | 26 | ] 27 | } 28 | 29 | 30 | def add_extras(cfgs): 31 | extras = nn.ModuleList() 32 | for cfg in cfgs: 33 | extra = [] 34 | for params in cfg: 35 | in_channels, out_channels, kernel_size, stride, padding = params 36 | extra.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)) 37 | extra.append(nn.ReLU()) 38 | extras.append(nn.Sequential(*extra)) 39 | return extras 40 | 41 | 42 | class MBConvBlock(nn.Module): 43 | """ 44 | Mobile Inverted Residual Bottleneck Block 45 | 46 | Args: 47 | block_args (namedtuple): BlockArgs, see above 48 | global_params (namedtuple): GlobalParam, see above 49 | 50 | Attributes: 51 | has_se (bool): Whether the block contains a Squeeze and Excitation layer. 52 | """ 53 | 54 | def __init__(self, block_args, global_params): 55 | super().__init__() 56 | self._block_args = block_args 57 | self._bn_mom = 1 - global_params.batch_norm_momentum 58 | self._bn_eps = global_params.batch_norm_epsilon 59 | self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) 60 | self.id_skip = block_args.id_skip # skip connection and drop connect 61 | 62 | # Expansion phase 63 | inp = self._block_args.input_filters # number of input channels 64 | oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels 65 | if self._block_args.expand_ratio != 1: 66 | self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) 67 | self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 68 | 69 | # Depthwise convolution phase 70 | k = self._block_args.kernel_size 71 | s = self._block_args.stride 72 | self._depthwise_conv = Conv2dSamePadding( 73 | in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise 74 | kernel_size=k, stride=s, bias=False) 75 | self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 76 | 77 | # Squeeze and Excitation layer, if desired 78 | if self.has_se: 79 | num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) 80 | self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) 81 | self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) 82 | 83 | # Output phase 84 | final_oup = self._block_args.output_filters 85 | self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) 86 | self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) 87 | 88 | def forward(self, inputs, drop_connect_rate=None): 89 | """ 90 | :param inputs: input tensor 91 | :param drop_connect_rate: drop connect rate (float, between 0 and 1) 92 | :return: output of block 93 | """ 94 | 95 | # Expansion and Depthwise Convolution 96 | x = inputs 97 | if self._block_args.expand_ratio != 1: 98 | x = relu_fn(self._bn0(self._expand_conv(inputs))) 99 | x = relu_fn(self._bn1(self._depthwise_conv(x))) 100 | 101 | # Squeeze and Excitation 102 | if self.has_se: 103 | x_squeezed = F.adaptive_avg_pool2d(x, 1) 104 | x_squeezed = self._se_expand(relu_fn(self._se_reduce(x_squeezed))) 105 | x = torch.sigmoid(x_squeezed) * x 106 | 107 | x = self._bn2(self._project_conv(x)) 108 | 109 | # Skip connection and drop connect 110 | input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters 111 | if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters: 112 | if drop_connect_rate: 113 | x = drop_connect(x, p=drop_connect_rate, training=self.training) 114 | x = x + inputs # skip connection 115 | return x 116 | 117 | 118 | class EfficientNet(nn.Module): 119 | """ 120 | An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods 121 | 122 | Args: 123 | blocks_args (list): A list of BlockArgs to construct blocks 124 | global_params (namedtuple): A set of GlobalParams shared between blocks 125 | 126 | Example: 127 | model = EfficientNet.from_pretrained('efficientnet-b0') 128 | 129 | """ 130 | 131 | def __init__(self, model_name, blocks_args=None, global_params=None): 132 | super().__init__() 133 | self.indices = INDICES[model_name] 134 | self.extras = add_extras(EXTRAS[model_name]) 135 | assert isinstance(blocks_args, list), 'blocks_args should be a list' 136 | assert len(blocks_args) > 0, 'block args must be greater than 0' 137 | self._global_params = global_params 138 | self._blocks_args = blocks_args 139 | 140 | # Batch norm parameters 141 | bn_mom = 1 - self._global_params.batch_norm_momentum 142 | bn_eps = self._global_params.batch_norm_epsilon 143 | 144 | # Stem 145 | in_channels = 3 # rgb 146 | out_channels = round_filters(32, self._global_params) # number of output channels 147 | self._conv_stem = Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False) 148 | self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) 149 | 150 | # Build blocks 151 | self._blocks = nn.ModuleList([]) 152 | for block_args in self._blocks_args: 153 | 154 | # Update block input and output filters based on depth multiplier. 155 | block_args = block_args._replace( 156 | input_filters=round_filters(block_args.input_filters, self._global_params), 157 | output_filters=round_filters(block_args.output_filters, self._global_params), 158 | num_repeat=round_repeats(block_args.num_repeat, self._global_params) 159 | ) 160 | 161 | # The first block needs to take care of stride and filter size increase. 162 | self._blocks.append(MBConvBlock(block_args, self._global_params)) 163 | if block_args.num_repeat > 1: 164 | block_args = block_args._replace(input_filters=block_args.output_filters, stride=1) 165 | for _ in range(block_args.num_repeat - 1): 166 | self._blocks.append(MBConvBlock(block_args, self._global_params)) 167 | self.reset_parameters() 168 | 169 | def reset_parameters(self): 170 | for m in self.extras.modules(): 171 | if isinstance(m, nn.Conv2d): 172 | nn.init.xavier_uniform_(m.weight) 173 | nn.init.zeros_(m.bias) 174 | 175 | def extract_features(self, inputs): 176 | """ Returns output of the final convolution layer """ 177 | 178 | # Stem 179 | x = relu_fn(self._bn0(self._conv_stem(inputs))) 180 | 181 | features = [] 182 | 183 | # Blocks 184 | for idx, block in enumerate(self._blocks): 185 | drop_connect_rate = self._global_params.drop_connect_rate 186 | if drop_connect_rate: 187 | drop_connect_rate *= float(idx) / len(self._blocks) 188 | x = block(x, drop_connect_rate) 189 | if idx in self.indices: 190 | features.append(x) 191 | 192 | return x, features 193 | 194 | def forward(self, inputs): 195 | """ Calls extract_features to extract features, applies final linear layer, and returns logits. """ 196 | 197 | # Convolution layers 198 | x, features = self.extract_features(inputs) 199 | 200 | for layer in self.extras: 201 | x = layer(x) 202 | features.append(x) 203 | 204 | return tuple(features) 205 | 206 | @classmethod 207 | def from_name(cls, model_name, override_params=None): 208 | cls._check_model_name_is_valid(model_name) 209 | blocks_args, global_params = get_model_params(model_name, override_params) 210 | return EfficientNet(model_name, blocks_args, global_params) 211 | 212 | @classmethod 213 | def from_pretrained(cls, model_name): 214 | model = EfficientNet.from_name(model_name) 215 | load_pretrained_weights(model, model_name) 216 | return model 217 | 218 | @classmethod 219 | def get_image_size(cls, model_name): 220 | cls._check_model_name_is_valid(model_name) 221 | _, _, res, _ = efficientnet_params(model_name) 222 | return res 223 | 224 | @classmethod 225 | def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False): 226 | """ Validates model name. None that pretrained weights are only available for 227 | the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """ 228 | num_models = 4 if also_need_pretrained_weights else 8 229 | valid_models = ['efficientnet_b' + str(i) for i in range(num_models)] 230 | if model_name.replace('-', '_') not in valid_models: 231 | raise ValueError('model_name should be one of: ' + ', '.join(valid_models)) 232 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/efficient_net/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains helper functions for building the model and for loading model parameters. 3 | These helper functions are built to mirror those in the official TensorFlow implementation. 4 | """ 5 | 6 | import re 7 | import math 8 | import collections 9 | import torch 10 | from torch import nn 11 | from torch.nn import functional as F 12 | from ssd.utils.model_zoo import load_state_dict_from_url 13 | 14 | ######################################################################## 15 | ############### HELPERS FUNCTIONS FOR MODEL ARCHITECTURE ############### 16 | ######################################################################## 17 | 18 | 19 | # Parameters for the entire model (stem, all blocks, and head) 20 | 21 | GlobalParams = collections.namedtuple('GlobalParams', [ 22 | 'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate', 23 | 'num_classes', 'width_coefficient', 'depth_coefficient', 24 | 'depth_divisor', 'min_depth', 'drop_connect_rate', ]) 25 | 26 | # Parameters for an individual model block 27 | BlockArgs = collections.namedtuple('BlockArgs', [ 28 | 'kernel_size', 'num_repeat', 'input_filters', 'output_filters', 29 | 'expand_ratio', 'id_skip', 'stride', 'se_ratio']) 30 | 31 | # Change namedtuple defaults 32 | GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields) 33 | BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields) 34 | 35 | 36 | def relu_fn(x): 37 | """ Swish activation function """ 38 | return x * torch.sigmoid(x) 39 | 40 | 41 | def round_filters(filters, global_params): 42 | """ Calculate and round number of filters based on depth multiplier. """ 43 | multiplier = global_params.width_coefficient 44 | if not multiplier: 45 | return filters 46 | divisor = global_params.depth_divisor 47 | min_depth = global_params.min_depth 48 | filters *= multiplier 49 | min_depth = min_depth or divisor 50 | new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor) 51 | if new_filters < 0.9 * filters: # prevent rounding by more than 10% 52 | new_filters += divisor 53 | return int(new_filters) 54 | 55 | 56 | def round_repeats(repeats, global_params): 57 | """ Round number of filters based on depth multiplier. """ 58 | multiplier = global_params.depth_coefficient 59 | if not multiplier: 60 | return repeats 61 | return int(math.ceil(multiplier * repeats)) 62 | 63 | 64 | def drop_connect(inputs, p, training): 65 | """ Drop connect. """ 66 | if not training: return inputs 67 | batch_size = inputs.shape[0] 68 | keep_prob = 1 - p 69 | random_tensor = keep_prob 70 | random_tensor += torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device) 71 | binary_tensor = torch.floor(random_tensor) 72 | output = inputs / keep_prob * binary_tensor 73 | return output 74 | 75 | 76 | class Conv2dSamePadding(nn.Conv2d): 77 | """ 2D Convolutions like TensorFlow """ 78 | 79 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True): 80 | super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias) 81 | self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2 82 | 83 | def forward(self, x): 84 | ih, iw = x.size()[-2:] 85 | kh, kw = self.weight.size()[-2:] 86 | sh, sw = self.stride 87 | oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) 88 | pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) 89 | pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) 90 | if pad_h > 0 or pad_w > 0: 91 | x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) 92 | return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) 93 | 94 | 95 | ######################################################################## 96 | ############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ############## 97 | ######################################################################## 98 | 99 | 100 | def efficientnet_params(model_name): 101 | """ Map EfficientNet model name to parameter coefficients. """ 102 | params_dict = { 103 | # Coefficients: width,depth,res,dropout 104 | 'efficientnet-b0': (1.0, 1.0, 224, 0.2), 105 | 'efficientnet-b1': (1.0, 1.1, 240, 0.2), 106 | 'efficientnet-b2': (1.1, 1.2, 260, 0.3), 107 | 'efficientnet-b3': (1.2, 1.4, 300, 0.3), 108 | 'efficientnet-b4': (1.4, 1.8, 380, 0.4), 109 | 'efficientnet-b5': (1.6, 2.2, 456, 0.4), 110 | 'efficientnet-b6': (1.8, 2.6, 528, 0.5), 111 | 'efficientnet-b7': (2.0, 3.1, 600, 0.5), 112 | } 113 | return params_dict[model_name] 114 | 115 | 116 | class BlockDecoder(object): 117 | """ Block Decoder for readability, straight from the official TensorFlow repository """ 118 | 119 | @staticmethod 120 | def _decode_block_string(block_string): 121 | """ Gets a block through a string notation of arguments. """ 122 | assert isinstance(block_string, str) 123 | 124 | ops = block_string.split('_') 125 | options = {} 126 | for op in ops: 127 | splits = re.split(r'(\d.*)', op) 128 | if len(splits) >= 2: 129 | key, value = splits[:2] 130 | options[key] = value 131 | 132 | # Check stride 133 | assert (('s' in options and len(options['s']) == 1) or 134 | (len(options['s']) == 2 and options['s'][0] == options['s'][1])) 135 | 136 | return BlockArgs( 137 | kernel_size=int(options['k']), 138 | num_repeat=int(options['r']), 139 | input_filters=int(options['i']), 140 | output_filters=int(options['o']), 141 | expand_ratio=int(options['e']), 142 | id_skip=('noskip' not in block_string), 143 | se_ratio=float(options['se']) if 'se' in options else None, 144 | stride=[int(options['s'][0])]) 145 | 146 | @staticmethod 147 | def _encode_block_string(block): 148 | """Encodes a block to a string.""" 149 | args = [ 150 | 'r%d' % block.num_repeat, 151 | 'k%d' % block.kernel_size, 152 | 's%d%d' % (block.strides[0], block.strides[1]), 153 | 'e%s' % block.expand_ratio, 154 | 'i%d' % block.input_filters, 155 | 'o%d' % block.output_filters 156 | ] 157 | if 0 < block.se_ratio <= 1: 158 | args.append('se%s' % block.se_ratio) 159 | if block.id_skip is False: 160 | args.append('noskip') 161 | return '_'.join(args) 162 | 163 | @staticmethod 164 | def decode(string_list): 165 | """ 166 | Decodes a list of string notations to specify blocks inside the network. 167 | 168 | :param string_list: a list of strings, each string is a notation of block 169 | :return: a list of BlockArgs namedtuples of block args 170 | """ 171 | assert isinstance(string_list, list) 172 | blocks_args = [] 173 | for block_string in string_list: 174 | blocks_args.append(BlockDecoder._decode_block_string(block_string)) 175 | return blocks_args 176 | 177 | @staticmethod 178 | def encode(blocks_args): 179 | """ 180 | Encodes a list of BlockArgs to a list of strings. 181 | 182 | :param blocks_args: a list of BlockArgs namedtuples of block args 183 | :return: a list of strings, each string is a notation of block 184 | """ 185 | block_strings = [] 186 | for block in blocks_args: 187 | block_strings.append(BlockDecoder._encode_block_string(block)) 188 | return block_strings 189 | 190 | 191 | def efficientnet(width_coefficient=None, depth_coefficient=None, 192 | dropout_rate=0.2, drop_connect_rate=0.2): 193 | """ Creates a efficientnet model. """ 194 | 195 | blocks_args = [ 196 | 'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25', 197 | 'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25', 198 | 'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25', 199 | 'r1_k3_s11_e6_i192_o320_se0.25', 200 | ] 201 | blocks_args = BlockDecoder.decode(blocks_args) 202 | 203 | global_params = GlobalParams( 204 | batch_norm_momentum=0.99, 205 | batch_norm_epsilon=1e-3, 206 | dropout_rate=dropout_rate, 207 | drop_connect_rate=drop_connect_rate, 208 | # data_format='channels_last', # removed, this is always true in PyTorch 209 | num_classes=1000, 210 | width_coefficient=width_coefficient, 211 | depth_coefficient=depth_coefficient, 212 | depth_divisor=8, 213 | min_depth=None 214 | ) 215 | 216 | return blocks_args, global_params 217 | 218 | 219 | def get_model_params(model_name, override_params): 220 | """ Get the block args and global params for a given model """ 221 | if model_name.startswith('efficientnet'): 222 | w, d, _, p = efficientnet_params(model_name) 223 | # note: all models have drop connect rate = 0.2 224 | blocks_args, global_params = efficientnet(width_coefficient=w, depth_coefficient=d, dropout_rate=p) 225 | else: 226 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 227 | if override_params: 228 | # ValueError will be raised here if override_params has fields not included in global_params. 229 | global_params = global_params._replace(**override_params) 230 | return blocks_args, global_params 231 | 232 | 233 | url_map = { 234 | 'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet-b0-08094119.pth', 235 | 'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet-b1-dbc7070a.pth', 236 | 'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet-b2-27687264.pth', 237 | 'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet-b3-c8376fa2.pth', 238 | 'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet-b4-e116e8b3.pth', 239 | 'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet-b5-586e6cc6.pth', 240 | } 241 | 242 | 243 | def load_pretrained_weights(model, model_name): 244 | """ Loads pretrained weights, and downloads if loading for the first time. """ 245 | state_dict = load_state_dict_from_url(url_map[model_name]) 246 | model.load_state_dict(state_dict, strict=False) 247 | print('Loaded pretrained weights for {}'.format(model_name)) 248 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/mobilenet.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from ssd.modeling import registry 4 | from ssd.utils.model_zoo import load_state_dict_from_url 5 | 6 | model_urls = { 7 | 'mobilenet_v2': 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth', 8 | } 9 | 10 | 11 | class ConvBNReLU(nn.Sequential): 12 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1): 13 | padding = (kernel_size - 1) // 2 14 | super(ConvBNReLU, self).__init__( 15 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False), 16 | nn.BatchNorm2d(out_planes), 17 | nn.ReLU6(inplace=True) 18 | ) 19 | 20 | 21 | class InvertedResidual(nn.Module): 22 | def __init__(self, inp, oup, stride, expand_ratio): 23 | super(InvertedResidual, self).__init__() 24 | self.stride = stride 25 | assert stride in [1, 2] 26 | 27 | hidden_dim = int(round(inp * expand_ratio)) 28 | self.use_res_connect = self.stride == 1 and inp == oup 29 | 30 | layers = [] 31 | if expand_ratio != 1: 32 | # pw 33 | layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1)) 34 | layers.extend([ 35 | # dw 36 | ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim), 37 | # pw-linear 38 | nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 39 | nn.BatchNorm2d(oup), 40 | ]) 41 | self.conv = nn.Sequential(*layers) 42 | 43 | def forward(self, x): 44 | if self.use_res_connect: 45 | return x + self.conv(x) 46 | else: 47 | return self.conv(x) 48 | 49 | 50 | class MobileNetV2(nn.Module): 51 | def __init__(self, width_mult=1.0, inverted_residual_setting=None): 52 | super(MobileNetV2, self).__init__() 53 | block = InvertedResidual 54 | input_channel = 32 55 | last_channel = 1280 56 | 57 | if inverted_residual_setting is None: 58 | inverted_residual_setting = [ 59 | # t, c, n, s 60 | [1, 16, 1, 1], 61 | [6, 24, 2, 2], 62 | [6, 32, 3, 2], 63 | [6, 64, 4, 2], 64 | [6, 96, 3, 1], 65 | [6, 160, 3, 2], 66 | [6, 320, 1, 1], 67 | ] 68 | 69 | # only check the first element, assuming user knows t,c,n,s are required 70 | if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4: 71 | raise ValueError("inverted_residual_setting should be non-empty " 72 | "or a 4-element list, got {}".format(inverted_residual_setting)) 73 | 74 | # building first layer 75 | input_channel = int(input_channel * width_mult) 76 | self.last_channel = int(last_channel * max(1.0, width_mult)) 77 | features = [ConvBNReLU(3, input_channel, stride=2)] 78 | # building inverted residual blocks 79 | for t, c, n, s in inverted_residual_setting: 80 | output_channel = int(c * width_mult) 81 | for i in range(n): 82 | stride = s if i == 0 else 1 83 | features.append(block(input_channel, output_channel, stride, expand_ratio=t)) 84 | input_channel = output_channel 85 | # building last several layers 86 | features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1)) 87 | # make it nn.Sequential 88 | self.features = nn.Sequential(*features) 89 | self.extras = nn.ModuleList([ 90 | InvertedResidual(1280, 512, 2, 0.2), 91 | InvertedResidual(512, 256, 2, 0.25), 92 | InvertedResidual(256, 256, 2, 0.5), 93 | InvertedResidual(256, 64, 2, 0.25) 94 | ]) 95 | 96 | self.reset_parameters() 97 | 98 | def reset_parameters(self): 99 | # weight initialization 100 | for m in self.modules(): 101 | if isinstance(m, nn.Conv2d): 102 | nn.init.kaiming_normal_(m.weight, mode='fan_out') 103 | if m.bias is not None: 104 | nn.init.zeros_(m.bias) 105 | elif isinstance(m, nn.BatchNorm2d): 106 | nn.init.ones_(m.weight) 107 | nn.init.zeros_(m.bias) 108 | elif isinstance(m, nn.Linear): 109 | nn.init.normal_(m.weight, 0, 0.01) 110 | nn.init.zeros_(m.bias) 111 | 112 | def forward(self, x): 113 | features = [] 114 | for i in range(14): 115 | x = self.features[i](x) 116 | features.append(x) 117 | 118 | for i in range(14, len(self.features)): 119 | x = self.features[i](x) 120 | features.append(x) 121 | 122 | for i in range(len(self.extras)): 123 | x = self.extras[i](x) 124 | features.append(x) 125 | 126 | return tuple(features) 127 | 128 | 129 | @registry.BACKBONES.register('mobilenet_v2') 130 | def mobilenet_v2(cfg, pretrained=True): 131 | model = MobileNetV2() 132 | if pretrained: 133 | model.load_state_dict(load_state_dict_from_url(model_urls['mobilenet_v2']), strict=False) 134 | return model 135 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/resnet.py: -------------------------------------------------------------------------------- 1 | import torchvision.models as models 2 | from collections import OrderedDict 3 | import torch.nn as nn 4 | import torch 5 | from ssd.modeling import registry 6 | class IntermediateLayerGetter(nn.ModuleDict): 7 | def __init__(self, model, return_layers): 8 | if not set(return_layers).issubset([name for name, _ in model.named_children()]): 9 | raise ValueError("return_layers are not present in model") 10 | 11 | orig_return_layers = return_layers 12 | return_layers = {k: v for k, v in return_layers.items()} 13 | layers = OrderedDict() 14 | for name, module in model.named_children(): 15 | layers[name] = module 16 | if name in return_layers: 17 | del return_layers[name] 18 | if not return_layers: 19 | break 20 | 21 | super(IntermediateLayerGetter, self).__init__(layers) 22 | self.return_layers = orig_return_layers 23 | 24 | def forward(self, x): 25 | out = OrderedDict() 26 | for name, module in self.named_children(): 27 | x = module(x) 28 | if name in self.return_layers: 29 | out_name = self.return_layers[name] 30 | out[out_name] = x 31 | return out 32 | 33 | return_layers = {'layer1': 1, 'layer2': 2, 'layer3': 3, 'layer4': 4} 34 | backbone = models.resnet101(pretrained=True) 35 | body = IntermediateLayerGetter(backbone, return_layers) 36 | class resnet(nn.Module): 37 | def __init__(self): 38 | super(resnet, self).__init__() 39 | self.body=body 40 | def forward(self, x): 41 | features = self.body(x) 42 | # return them as a tuple 43 | return tuple(features.values()) 44 | #model=resnet() 45 | #x=torch.randn(1,3,512,512) 46 | #out=model(x) 47 | #print(out[0].shape,out[1].shape,out[2].shape,out[3].shape) 48 | @registry.BACKBONES.register('res_backbone') 49 | def res_backbone(cfg, pretrained=True): 50 | model = resnet() 51 | return model 52 | -------------------------------------------------------------------------------- /ssd/modeling/backbone/utils.py: -------------------------------------------------------------------------------- 1 | try: 2 | from torch.hub import load_state_dict_from_url 3 | except ImportError: 4 | from torch.utils.model_zoo import load_url as load_state_dict_from_url -------------------------------------------------------------------------------- /ssd/modeling/backbone/vgg.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | from ssd.layers import L2Norm 5 | from ssd.modeling import registry 6 | from ssd.utils.model_zoo import load_state_dict_from_url 7 | 8 | model_urls = { 9 | 'vgg': 'https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth', 10 | } 11 | 12 | 13 | # borrowed from https://github.com/amdegroot/ssd.pytorch/blob/master/ssd.py 14 | def add_vgg(cfg, batch_norm=False): 15 | layers = [] 16 | in_channels = 3 17 | for v in cfg: 18 | if v == 'M': 19 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 20 | elif v == 'C': 21 | layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)] 22 | else: 23 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 24 | if batch_norm: 25 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 26 | else: 27 | layers += [conv2d, nn.ReLU(inplace=True)] 28 | in_channels = v 29 | pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1) 30 | conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6) 31 | conv7 = nn.Conv2d(1024, 1024, kernel_size=1) 32 | layers += [pool5, conv6, 33 | nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)] 34 | return layers 35 | 36 | 37 | def add_extras(cfg, i, size=300): 38 | # Extra layers added to VGG for feature scaling 39 | layers = [] 40 | in_channels = i 41 | flag = False 42 | for k, v in enumerate(cfg): 43 | if in_channels != 'S': 44 | if v == 'S': 45 | layers += [nn.Conv2d(in_channels, cfg[k + 1], kernel_size=(1, 3)[flag], stride=2, padding=1)] 46 | else: 47 | layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])] 48 | flag = not flag 49 | in_channels = v 50 | if size == 512: 51 | layers.append(nn.Conv2d(in_channels, 128, kernel_size=1, stride=1)) 52 | layers.append(nn.Conv2d(128, 256, kernel_size=4, stride=1, padding=1)) 53 | return layers 54 | 55 | 56 | vgg_base = { 57 | '300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 58 | 512, 512, 512], 59 | '512': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 60 | 512, 512, 512], 61 | } 62 | extras_base = { 63 | '300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256], 64 | '512': [256, 'S', 512, 128, 'S', 256, 128, 'S', 256, 128, 'S', 256], 65 | } 66 | 67 | 68 | class VGG(nn.Module): 69 | def __init__(self, cfg): 70 | super().__init__() 71 | size = cfg.INPUT.IMAGE_SIZE 72 | vgg_config = vgg_base[str(size)] 73 | extras_config = extras_base[str(size)] 74 | 75 | self.vgg = nn.ModuleList(add_vgg(vgg_config)) 76 | self.extras = nn.ModuleList(add_extras(extras_config, i=1024, size=size)) 77 | self.l2_norm = L2Norm(512, scale=20) 78 | self.reset_parameters() 79 | 80 | def reset_parameters(self): 81 | for m in self.extras.modules(): 82 | if isinstance(m, nn.Conv2d): 83 | nn.init.xavier_uniform_(m.weight) 84 | nn.init.zeros_(m.bias) 85 | 86 | def init_from_pretrain(self, state_dict): 87 | self.vgg.load_state_dict(state_dict) 88 | 89 | def forward(self, x): 90 | features = [] 91 | for i in range(23): 92 | x = self.vgg[i](x) 93 | s = self.l2_norm(x) # Conv4_3 L2 normalization 94 | features.append(s) 95 | 96 | # apply vgg up to fc7 97 | for i in range(23, len(self.vgg)): 98 | x = self.vgg[i](x) 99 | features.append(x) 100 | 101 | for k, v in enumerate(self.extras): 102 | x = F.relu(v(x), inplace=True) 103 | if k % 2 == 1: 104 | features.append(x) 105 | 106 | return tuple(features) 107 | 108 | 109 | @registry.BACKBONES.register('vgg') 110 | def vgg(cfg, pretrained=True): 111 | model = VGG(cfg) 112 | if pretrained: 113 | model.init_from_pretrain(load_state_dict_from_url(model_urls['vgg'])) 114 | return model 115 | -------------------------------------------------------------------------------- /ssd/modeling/box/__init__.py: -------------------------------------------------------------------------------- 1 | from .prior_box import PriorBox 2 | from .box_utils import decode,nms, diounms 3 | from .box_utils import match, log_sum_exp,match_ious,bbox_overlaps_iou, bbox_overlaps_giou, bbox_overlaps_diou, bbox_overlaps_ciou 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/__init__.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/__init__.cpython-35.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/box_utils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/box_utils.cpython-35.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/box_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/box_utils.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/prior_box.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/prior_box.cpython-35.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/__pycache__/prior_box.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box/__pycache__/prior_box.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box/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 | s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size)) 45 | mean += [cx, cy, s_k_prime, s_k_prime] 46 | 47 | # rest of aspect ratios 48 | #print(self.aspect_ratios[k]) 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 | #print(output.shape) 57 | return output 58 | -------------------------------------------------------------------------------- /ssd/modeling/box_head/__init__.py: -------------------------------------------------------------------------------- 1 | from ssd.modeling import registry 2 | from .box_head import SSDBoxHead 3 | 4 | __all__ = ['build_box_head', 'SSDBoxHead'] 5 | 6 | 7 | def build_box_head(cfg): 8 | return registry.BOX_HEADS[cfg.MODEL.BOX_HEAD.NAME](cfg) 9 | -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/box_head.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/box_head.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/box_predictor.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/box_predictor.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/inference.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/inference.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/loss.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/__pycache__/loss_Iou.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/box_head/__pycache__/loss_Iou.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/box_head/box_head.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | import torch.nn.functional as F 3 | 4 | from ssd.modeling import registry 5 | from ssd.modeling.anchors.prior_box import PriorBox 6 | from ssd.modeling.box_head.box_predictor import make_box_predictor 7 | from ssd.utils import box_utils 8 | from .inference import PostProcessor 9 | from .loss import MultiBoxLoss 10 | 11 | 12 | @registry.BOX_HEADS.register('SSDBoxHead') 13 | class SSDBoxHead(nn.Module): 14 | def __init__(self, cfg): 15 | super().__init__() 16 | self.cfg = cfg 17 | self.predictor = make_box_predictor(cfg) 18 | self.loss_evaluator = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO) 19 | self.post_processor = PostProcessor(cfg) 20 | self.priors = None 21 | 22 | def forward(self, features, targets=None): 23 | cls_logits, bbox_pred = self.predictor(features) 24 | if self.training: 25 | return self._forward_train(cls_logits, bbox_pred, targets) 26 | else: 27 | return self._forward_test(cls_logits, bbox_pred) 28 | 29 | def _forward_train(self, cls_logits, bbox_pred, targets): 30 | gt_boxes, gt_labels = targets['boxes'], targets['labels'] 31 | #print(gt_boxes.shape)#torch.Size([2, 45180, 4]) 32 | reg_loss, cls_loss = self.loss_evaluator(cls_logits, bbox_pred, gt_labels, gt_boxes) 33 | loss_dict = dict( 34 | reg_loss=reg_loss, 35 | cls_loss=cls_loss, 36 | ) 37 | detections = (cls_logits, bbox_pred) 38 | return detections, loss_dict 39 | 40 | def _forward_test(self, cls_logits, bbox_pred): 41 | if self.priors is None: 42 | self.priors = PriorBox(self.cfg)().to(bbox_pred.device) 43 | scores = F.softmax(cls_logits, dim=2) 44 | boxes = box_utils.convert_locations_to_boxes( 45 | bbox_pred, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE 46 | ) 47 | boxes = box_utils.center_form_to_corner_form(boxes) 48 | detections = (scores, boxes) 49 | detections = self.post_processor(detections) 50 | return detections, {} 51 | -------------------------------------------------------------------------------- /ssd/modeling/box_head/box_predictor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from ssd.layers import SeparableConv2d 5 | from ssd.modeling import registry 6 | 7 | 8 | class BoxPredictor(nn.Module): 9 | def __init__(self, cfg): 10 | super().__init__() 11 | self.cfg = cfg 12 | self.cls_headers = nn.ModuleList() 13 | self.reg_headers = nn.ModuleList() 14 | for level, (boxes_per_location, out_channels) in enumerate(zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)): 15 | self.cls_headers.append(self.cls_block(level, out_channels, boxes_per_location)) 16 | self.reg_headers.append(self.reg_block(level, out_channels, boxes_per_location)) 17 | self.reset_parameters() 18 | 19 | def cls_block(self, level, out_channels, boxes_per_location): 20 | raise NotImplementedError 21 | 22 | def reg_block(self, level, out_channels, boxes_per_location): 23 | raise NotImplementedError 24 | 25 | def reset_parameters(self): 26 | for m in self.modules(): 27 | if isinstance(m, nn.Conv2d): 28 | nn.init.xavier_uniform_(m.weight) 29 | nn.init.zeros_(m.bias) 30 | 31 | def forward(self, features): 32 | cls_logits = [] 33 | bbox_pred = [] 34 | for feature, cls_header, reg_header in zip(features, self.cls_headers, self.reg_headers): 35 | cls_logits.append(cls_header(feature).permute(0, 2, 3, 1).contiguous()) 36 | bbox_pred.append(reg_header(feature).permute(0, 2, 3, 1).contiguous()) 37 | 38 | batch_size = features[0].shape[0] 39 | cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1, self.cfg.MODEL.NUM_CLASSES) 40 | bbox_pred = torch.cat([l.view(l.shape[0], -1) for l in bbox_pred], dim=1).view(batch_size, -1, 4) 41 | 42 | return cls_logits, bbox_pred 43 | 44 | 45 | @registry.BOX_PREDICTORS.register('SSDBoxPredictor') 46 | class SSDBoxPredictor(BoxPredictor): 47 | def cls_block(self, level, out_channels, boxes_per_location): 48 | return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1) 49 | 50 | def reg_block(self, level, out_channels, boxes_per_location): 51 | return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1) 52 | 53 | 54 | @registry.BOX_PREDICTORS.register('SSDLiteBoxPredictor') 55 | class SSDLiteBoxPredictor(BoxPredictor): 56 | def cls_block(self, level, out_channels, boxes_per_location): 57 | num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS) 58 | if level == num_levels - 1: 59 | return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=1) 60 | return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1) 61 | 62 | def reg_block(self, level, out_channels, boxes_per_location): 63 | num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS) 64 | if level == num_levels - 1: 65 | return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=1) 66 | return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1) 67 | 68 | 69 | def make_box_predictor(cfg): 70 | return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg) 71 | -------------------------------------------------------------------------------- /ssd/modeling/box_head/inference.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ssd.structures.container import Container 4 | from ssd.utils.nms import batched_nms 5 | 6 | 7 | class PostProcessor: 8 | def __init__(self, cfg): 9 | super().__init__() 10 | self.cfg = cfg 11 | self.width = cfg.INPUT.IMAGE_SIZE 12 | self.height = cfg.INPUT.IMAGE_SIZE 13 | 14 | def __call__(self, detections): 15 | batches_scores, batches_boxes = detections 16 | device = batches_scores.device 17 | batch_size = batches_scores.size(0) 18 | results = [] 19 | for batch_id in range(batch_size): 20 | scores, boxes = batches_scores[batch_id], batches_boxes[batch_id] # (N, #CLS) (N, 4) 21 | num_boxes = scores.shape[0] 22 | num_classes = scores.shape[1] 23 | 24 | boxes = boxes.view(num_boxes, 1, 4).expand(num_boxes, num_classes, 4) 25 | labels = torch.arange(num_classes, device=device) 26 | labels = labels.view(1, num_classes).expand_as(scores) 27 | 28 | # remove predictions with the background label 29 | boxes = boxes[:, 1:] 30 | scores = scores[:, 1:] 31 | labels = labels[:, 1:] 32 | 33 | # batch everything, by making every class prediction be a separate instance 34 | boxes = boxes.reshape(-1, 4) 35 | scores = scores.reshape(-1) 36 | labels = labels.reshape(-1) 37 | 38 | # remove low scoring boxes 39 | indices = torch.nonzero(scores > self.cfg.TEST.CONFIDENCE_THRESHOLD).squeeze(1) 40 | boxes, scores, labels = boxes[indices], scores[indices], labels[indices] 41 | 42 | boxes[:, 0::2] *= self.width 43 | boxes[:, 1::2] *= self.height 44 | 45 | keep = batched_nms(boxes, scores, labels, self.cfg.TEST.NMS_THRESHOLD) 46 | # keep only topk scoring predictions 47 | keep = keep[:self.cfg.TEST.MAX_PER_IMAGE] 48 | boxes, scores, labels = boxes[keep], scores[keep], labels[keep] 49 | 50 | container = Container(boxes=boxes, labels=labels, scores=scores) 51 | container.img_width = self.width 52 | container.img_height = self.height 53 | results.append(container) 54 | return results 55 | -------------------------------------------------------------------------------- /ssd/modeling/box_head/loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | import torch 4 | import argparse 5 | from ssd.utils import box_utils 6 | from .loss_Iou import IouLoss 7 | from itertools import product 8 | 9 | from math import sqrt 10 | 11 | 12 | class PriorBox: 13 | def __init__(self): 14 | self.image_size = 300 15 | self.feature_maps = [75, 38, 19, 10] 16 | self.min_sizes = [36, 100, 159, 253] 17 | self.max_sizes = [100, 159, 253, 300] 18 | self.strides = [4, 8, 16, 30] 19 | self.aspect_ratios = [[2, 3], [4, 3], [3, 2], [1, 1]] 20 | self.clip = True 21 | 22 | def __call__(self): 23 | """Generate SSD Prior Boxes. 24 | It returns the center, height and width of the priors. The values are relative to the image size 25 | Returns: 26 | priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values 27 | are relative to the image size. 28 | """ 29 | priors = [] 30 | for k, f in enumerate(self.feature_maps): 31 | scale = self.image_size / self.strides[k] 32 | for i, j in product(range(f), repeat=2): 33 | # unit center x,y 34 | cx = (j + 0.5) / scale 35 | cy = (i + 0.5) / scale 36 | 37 | # small sized square box 38 | size = self.min_sizes[k] 39 | h = w = size / self.image_size 40 | priors.append([cx, cy, w, h]) 41 | 42 | # big sized square box 43 | size = sqrt(self.min_sizes[k] * self.max_sizes[k]) 44 | h = w = size / self.image_size 45 | priors.append([cx, cy, w, h]) 46 | 47 | # change h/w ratio of the small sized box 48 | size = self.min_sizes[k] 49 | h = w = size / self.image_size 50 | for ratio in self.aspect_ratios[k]: 51 | ratio = sqrt(ratio) 52 | priors.append([cx, cy, w * ratio, h / ratio]) 53 | priors.append([cx, cy, w / ratio, h * ratio]) 54 | 55 | priors = torch.tensor(priors) 56 | if self.clip: 57 | priors.clamp_(max=1, min=0) 58 | return priors 59 | 60 | 61 | class MultiBoxLoss(nn.Module): 62 | def __init__(self, neg_pos_ratio): 63 | """Implement SSD MultiBox Loss. 64 | 65 | Basically, MultiBox loss combines classification loss 66 | and Smooth L1 regression loss. 67 | """ 68 | super(MultiBoxLoss, self).__init__() 69 | self.neg_pos_ratio = neg_pos_ratio 70 | self.iouloss=IouLoss(size_sum=True, losstype='CIou') 71 | 72 | def forward(self, confidence, predicted_locations, labels, gt_locations): 73 | """Compute classification loss and smooth l1 loss. 74 | 75 | Args: 76 | confidence (batch_size, num_priors, num_classes): class predictions. 77 | predicted_locations (batch_size, num_priors, 4): predicted locations. 78 | labels (batch_size, num_priors): real labels of all the priors. 79 | gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors. 80 | """ 81 | priors = PriorBox()().to(predicted_locations.device) 82 | priors=priors.unsqueeze(0).expand(predicted_locations.shape[0],-1,-1) 83 | num_classes = confidence.size(2) 84 | with torch.no_grad(): 85 | # derived from cross_entropy=sum(log(p)) 86 | loss = -F.log_softmax(confidence, dim=2)[:, :, 0] 87 | mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio) 88 | 89 | confidence = confidence[mask, :] 90 | classification_loss = F.cross_entropy(confidence.view(-1, num_classes), labels[mask], reduction='sum') 91 | 92 | pos_mask = labels > 0 93 | predicted_locations = predicted_locations[pos_mask, :].view(-1, 4) 94 | gt_locations = gt_locations[pos_mask, :].view(-1, 4) 95 | priors=priors[pos_mask, :].view(-1, 4) 96 | smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum') 97 | #smooth_l1_loss = self.Giou_loss(predicted_locations, gt_locations, reduction='mean') 98 | iou_loss=self.iouloss(predicted_locations, gt_locations,priors) 99 | #print(iou_loss) 100 | num_pos = gt_locations.size(0) 101 | return iou_loss/ num_pos, classification_loss / num_pos -------------------------------------------------------------------------------- /ssd/modeling/detector/__init__.py: -------------------------------------------------------------------------------- 1 | from .ssd_detector import SSDDetector 2 | 3 | _DETECTION_META_ARCHITECTURES = { 4 | "SSDDetector": SSDDetector 5 | } 6 | 7 | 8 | def build_detection_model(cfg): 9 | meta_arch = _DETECTION_META_ARCHITECTURES[cfg.MODEL.META_ARCHITECTURE] 10 | return meta_arch(cfg) 11 | -------------------------------------------------------------------------------- /ssd/modeling/detector/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/detector/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/detector/__pycache__/ssd_detector.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/modeling/detector/__pycache__/ssd_detector.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/modeling/detector/ssd_detector.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from ssd.modeling.backbone import build_backbone 4 | from ssd.modeling.box_head import build_box_head 5 | 6 | 7 | class SSDDetector(nn.Module): 8 | def __init__(self, cfg): 9 | super().__init__() 10 | self.cfg = cfg 11 | self.backbone = build_backbone(cfg) 12 | self.box_head = build_box_head(cfg) 13 | 14 | def forward(self, images, targets=None): 15 | features = self.backbone(images) 16 | detections, detector_losses = self.box_head(features, targets) 17 | if self.training: 18 | return detector_losses 19 | return detections 20 | -------------------------------------------------------------------------------- /ssd/modeling/registry.py: -------------------------------------------------------------------------------- 1 | from ssd.utils.registry import Registry 2 | 3 | BACKBONES = Registry() 4 | BOX_HEADS = Registry() 5 | BOX_PREDICTORS = Registry() 6 | -------------------------------------------------------------------------------- /ssd/solver/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/solver/__init__.py -------------------------------------------------------------------------------- /ssd/solver/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/solver/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/solver/__pycache__/build.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/solver/__pycache__/build.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/solver/__pycache__/lr_scheduler.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/solver/__pycache__/lr_scheduler.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/solver/build.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .lr_scheduler import WarmupMultiStepLR 4 | 5 | 6 | def make_optimizer(cfg, model, lr=None): 7 | lr = cfg.SOLVER.BASE_LR if lr is None else lr 8 | return torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY) 9 | 10 | 11 | def make_lr_scheduler(cfg, optimizer, milestones=None): 12 | return WarmupMultiStepLR(optimizer=optimizer, 13 | milestones=cfg.SOLVER.LR_STEPS if milestones is None else milestones, 14 | gamma=cfg.SOLVER.GAMMA, 15 | warmup_factor=cfg.SOLVER.WARMUP_FACTOR, 16 | warmup_iters=cfg.SOLVER.WARMUP_ITERS) 17 | -------------------------------------------------------------------------------- /ssd/solver/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | from bisect import bisect_right 2 | 3 | from torch.optim.lr_scheduler import _LRScheduler 4 | 5 | 6 | class WarmupMultiStepLR(_LRScheduler): 7 | def __init__(self, optimizer, milestones, gamma=0.1, warmup_factor=1.0 / 3, 8 | warmup_iters=500, last_epoch=-1): 9 | if not list(milestones) == sorted(milestones): 10 | raise ValueError( 11 | "Milestones should be a list of" " increasing integers. Got {}", 12 | milestones, 13 | ) 14 | 15 | self.milestones = milestones 16 | self.gamma = gamma 17 | self.warmup_factor = warmup_factor 18 | self.warmup_iters = warmup_iters 19 | super().__init__(optimizer, last_epoch) 20 | 21 | def get_lr(self): 22 | warmup_factor = 1 23 | if self.last_epoch < self.warmup_iters: 24 | alpha = float(self.last_epoch) / self.warmup_iters 25 | warmup_factor = self.warmup_factor * (1 - alpha) + alpha 26 | return [ 27 | base_lr 28 | * warmup_factor 29 | * self.gamma ** bisect_right(self.milestones, self.last_epoch) 30 | for base_lr in self.base_lrs 31 | ] 32 | -------------------------------------------------------------------------------- /ssd/structures/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/structures/__init__.py -------------------------------------------------------------------------------- /ssd/structures/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/structures/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/structures/__pycache__/container.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/structures/__pycache__/container.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/structures/container.py: -------------------------------------------------------------------------------- 1 | class Container: 2 | """ 3 | Help class for manage boxes, labels, etc... 4 | Not inherit dict due to `default_collate` will change dict's subclass to dict. 5 | """ 6 | 7 | def __init__(self, *args, **kwargs): 8 | self._data_dict = dict(*args, **kwargs) 9 | 10 | def __setattr__(self, key, value): 11 | object.__setattr__(self, key, value) 12 | 13 | def __getitem__(self, key): 14 | return self._data_dict[key] 15 | 16 | def __iter__(self): 17 | return self._data_dict.__iter__() 18 | 19 | def __setitem__(self, key, value): 20 | self._data_dict[key] = value 21 | 22 | def _call(self, name, *args, **kwargs): 23 | keys = list(self._data_dict.keys()) 24 | for key in keys: 25 | value = self._data_dict[key] 26 | if hasattr(value, name): 27 | self._data_dict[key] = getattr(value, name)(*args, **kwargs) 28 | return self 29 | 30 | def to(self, *args, **kwargs): 31 | return self._call('to', *args, **kwargs) 32 | 33 | def numpy(self): 34 | return self._call('numpy') 35 | 36 | def resize(self, size): 37 | """resize boxes 38 | Args: 39 | size: (width, height) 40 | Returns: 41 | self 42 | """ 43 | img_width = getattr(self, 'img_width', -1) 44 | img_height = getattr(self, 'img_height', -1) 45 | assert img_width > 0 and img_height > 0 46 | assert 'boxes' in self._data_dict 47 | boxes = self._data_dict['boxes'] 48 | new_width, new_height = size 49 | boxes[:, 0::2] *= (new_width / img_width) 50 | boxes[:, 1::2] *= (new_height / img_height) 51 | return self 52 | 53 | def __repr__(self): 54 | return self._data_dict.__repr__() 55 | -------------------------------------------------------------------------------- /ssd/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .misc import * 2 | -------------------------------------------------------------------------------- /ssd/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/box_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/box_utils.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/checkpoint.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/checkpoint.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/dist_util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/dist_util.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/logger.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/logger.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/metric_logger.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/metric_logger.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/misc.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/misc.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/model_zoo.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/model_zoo.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/nms.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/nms.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/__pycache__/registry.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DLLXW/GarbageDetection/03aef87dd3d5b658207e3fd32c0046260831d40c/ssd/utils/__pycache__/registry.cpython-36.pyc -------------------------------------------------------------------------------- /ssd/utils/box_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | 4 | 5 | def convert_locations_to_boxes(locations, priors, center_variance, 6 | size_variance): 7 | """Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w). 8 | 9 | The conversion: 10 | $$predicted\_center * center_variance = \frac {real\_center - prior\_center} {prior\_hw}$$ 11 | $$exp(predicted\_hw * size_variance) = \frac {real\_hw} {prior\_hw}$$ 12 | We do it in the inverse direction here. 13 | Args: 14 | locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well. 15 | priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes. 16 | center_variance: a float used to change the scale of center. 17 | size_variance: a float used to change of scale of size. 18 | Returns: 19 | boxes: priors: [[center_x, center_y, w, h]]. All the values 20 | are relative to the image size. 21 | """ 22 | # priors can have one dimension less. 23 | if priors.dim() + 1 == locations.dim(): 24 | priors = priors.unsqueeze(0) 25 | return torch.cat([ 26 | locations[..., :2] * center_variance * priors[..., 2:] + priors[..., :2], 27 | torch.exp(locations[..., 2:] * size_variance) * priors[..., 2:] 28 | ], dim=locations.dim() - 1) 29 | 30 | 31 | def convert_boxes_to_locations(center_form_boxes, center_form_priors, center_variance, size_variance): 32 | # priors can have one dimension less 33 | if center_form_priors.dim() + 1 == center_form_boxes.dim(): 34 | center_form_priors = center_form_priors.unsqueeze(0) 35 | return torch.cat([ 36 | (center_form_boxes[..., :2] - center_form_priors[..., :2]) / center_form_priors[..., 2:] / center_variance, 37 | torch.log(center_form_boxes[..., 2:] / center_form_priors[..., 2:]) / size_variance 38 | ], dim=center_form_boxes.dim() - 1) 39 | 40 | 41 | def area_of(left_top, right_bottom) -> torch.Tensor: 42 | """Compute the areas of rectangles given two corners. 43 | 44 | Args: 45 | left_top (N, 2): left top corner. 46 | right_bottom (N, 2): right bottom corner. 47 | 48 | Returns: 49 | area (N): return the area. 50 | """ 51 | hw = torch.clamp(right_bottom - left_top, min=0.0) 52 | return hw[..., 0] * hw[..., 1] 53 | 54 | 55 | def iou_of(boxes0, boxes1, eps=1e-5): 56 | """Return intersection-over-union (Jaccard index) of boxes. 57 | 58 | Args: 59 | boxes0 (N, 4): ground truth boxes. 60 | boxes1 (N or 1, 4): predicted boxes. 61 | eps: a small number to avoid 0 as denominator. 62 | Returns: 63 | iou (N): IoU values. 64 | """ 65 | overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2]) 66 | overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:]) 67 | 68 | overlap_area = area_of(overlap_left_top, overlap_right_bottom) 69 | area0 = area_of(boxes0[..., :2], boxes0[..., 2:]) 70 | area1 = area_of(boxes1[..., :2], boxes1[..., 2:]) 71 | return overlap_area / (area0 + area1 - overlap_area + eps) 72 | 73 | 74 | def assign_priors(gt_boxes, gt_labels, corner_form_priors, 75 | iou_threshold): 76 | """Assign ground truth boxes and targets to priors. 77 | 78 | Args: 79 | gt_boxes (num_targets, 4): ground truth boxes. 80 | gt_labels (num_targets): labels of targets. 81 | priors (num_priors, 4): corner form priors 82 | Returns: 83 | boxes (num_priors, 4): real values for priors. 84 | labels (num_priros): labels for priors. 85 | """ 86 | # size: num_priors x num_targets 87 | ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1)) 88 | # size: num_priors 89 | best_target_per_prior, best_target_per_prior_index = ious.max(1) 90 | # size: num_targets 91 | best_prior_per_target, best_prior_per_target_index = ious.max(0) 92 | 93 | for target_index, prior_index in enumerate(best_prior_per_target_index): 94 | best_target_per_prior_index[prior_index] = target_index 95 | # 2.0 is used to make sure every target has a prior assigned 96 | best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2) 97 | # size: num_priors 98 | labels = gt_labels[best_target_per_prior_index] 99 | labels[best_target_per_prior < iou_threshold] = 0 # the backgournd id 100 | boxes = gt_boxes[best_target_per_prior_index] 101 | return boxes, labels 102 | 103 | 104 | def hard_negative_mining(loss, labels, neg_pos_ratio): 105 | """ 106 | It used to suppress the presence of a large number of negative prediction. 107 | It works on image level not batch level. 108 | For any example/image, it keeps all the positive predictions and 109 | cut the number of negative predictions to make sure the ratio 110 | between the negative examples and positive examples is no more 111 | the given ratio for an image. 112 | 113 | Args: 114 | loss (N, num_priors): the loss for each example. 115 | labels (N, num_priors): the labels. 116 | neg_pos_ratio: the ratio between the negative examples and positive examples. 117 | """ 118 | pos_mask = labels > 0 119 | num_pos = pos_mask.long().sum(dim=1, keepdim=True) 120 | num_neg = num_pos * neg_pos_ratio 121 | 122 | loss[pos_mask] = -math.inf 123 | _, indexes = loss.sort(dim=1, descending=True) 124 | _, orders = indexes.sort(dim=1) 125 | neg_mask = orders < num_neg 126 | return pos_mask | neg_mask 127 | 128 | 129 | def center_form_to_corner_form(locations): 130 | return torch.cat([locations[..., :2] - locations[..., 2:] / 2, 131 | locations[..., :2] + locations[..., 2:] / 2], locations.dim() - 1) 132 | 133 | 134 | def corner_form_to_center_form(boxes): 135 | return torch.cat([ 136 | (boxes[..., :2] + boxes[..., 2:]) / 2, 137 | boxes[..., 2:] - boxes[..., :2] 138 | ], boxes.dim() - 1) 139 | -------------------------------------------------------------------------------- /ssd/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | 4 | import torch 5 | from torch.nn.parallel import DistributedDataParallel 6 | 7 | from ssd.utils.model_zoo import cache_url 8 | 9 | 10 | class CheckPointer: 11 | _last_checkpoint_name = 'last_checkpoint.txt' 12 | 13 | def __init__(self, 14 | model, 15 | optimizer=None, 16 | scheduler=None, 17 | save_dir="", 18 | save_to_disk=None, 19 | logger=None): 20 | self.model = model 21 | self.optimizer = optimizer 22 | self.scheduler = scheduler 23 | self.save_dir = save_dir 24 | self.save_to_disk = save_to_disk 25 | if logger is None: 26 | logger = logging.getLogger(__name__) 27 | self.logger = logger 28 | 29 | def save(self, name, **kwargs): 30 | if not self.save_dir: 31 | return 32 | 33 | if not self.save_to_disk: 34 | return 35 | 36 | data = {} 37 | if isinstance(self.model, DistributedDataParallel): 38 | data['model'] = self.model.module.state_dict() 39 | else: 40 | data['model'] = self.model.state_dict() 41 | if self.optimizer is not None: 42 | data["optimizer"] = self.optimizer.state_dict() 43 | if self.scheduler is not None: 44 | data["scheduler"] = self.scheduler.state_dict() 45 | data.update(kwargs) 46 | 47 | save_file = os.path.join(self.save_dir, "{}.pth".format(name)) 48 | self.logger.info("Saving checkpoint to {}".format(save_file)) 49 | torch.save(data, save_file) 50 | 51 | self.tag_last_checkpoint(save_file) 52 | 53 | def load(self, f=None, use_latest=True): 54 | if self.has_checkpoint() and use_latest: 55 | # override argument with existing checkpoint 56 | f = self.get_checkpoint_file() 57 | if not f: 58 | # no checkpoint could be found 59 | self.logger.info("No checkpoint found.") 60 | return {} 61 | 62 | self.logger.info("Loading checkpoint from {}".format(f)) 63 | checkpoint = self._load_file(f) 64 | model = self.model 65 | if isinstance(model, DistributedDataParallel): 66 | model = self.model.module 67 | 68 | model.load_state_dict(checkpoint.pop("model")) 69 | if "optimizer" in checkpoint and self.optimizer: 70 | self.logger.info("Loading optimizer from {}".format(f)) 71 | self.optimizer.load_state_dict(checkpoint.pop("optimizer")) 72 | if "scheduler" in checkpoint and self.scheduler: 73 | self.logger.info("Loading scheduler from {}".format(f)) 74 | self.scheduler.load_state_dict(checkpoint.pop("scheduler")) 75 | 76 | # return any further checkpoint data 77 | return checkpoint 78 | 79 | def get_checkpoint_file(self): 80 | save_file = os.path.join(self.save_dir, self._last_checkpoint_name) 81 | try: 82 | with open(save_file, "r") as f: 83 | last_saved = f.read() 84 | last_saved = last_saved.strip() 85 | except IOError: 86 | # if file doesn't exist, maybe because it has just been 87 | # deleted by a separate process 88 | last_saved = "" 89 | return last_saved 90 | 91 | def has_checkpoint(self): 92 | save_file = os.path.join(self.save_dir, self._last_checkpoint_name) 93 | return os.path.exists(save_file) 94 | 95 | def tag_last_checkpoint(self, last_filename): 96 | save_file = os.path.join(self.save_dir, self._last_checkpoint_name) 97 | with open(save_file, "w") as f: 98 | f.write(last_filename) 99 | 100 | def _load_file(self, f): 101 | # download url files 102 | if f.startswith("http"): 103 | # if the file is a url path, download it and cache it 104 | cached_f = cache_url(f) 105 | self.logger.info("url {} cached in {}".format(f, cached_f)) 106 | f = cached_f 107 | return torch.load(f, map_location=torch.device("cpu")) 108 | -------------------------------------------------------------------------------- /ssd/utils/dist_util.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import torch 4 | import torch.distributed as dist 5 | 6 | 7 | def get_world_size(): 8 | if not dist.is_available(): 9 | return 1 10 | if not dist.is_initialized(): 11 | return 1 12 | return dist.get_world_size() 13 | 14 | 15 | def get_rank(): 16 | if not dist.is_available(): 17 | return 0 18 | if not dist.is_initialized(): 19 | return 0 20 | return dist.get_rank() 21 | 22 | 23 | def is_main_process(): 24 | return get_rank() == 0 25 | 26 | 27 | def synchronize(): 28 | """ 29 | Helper function to synchronize (barrier) among all processes when 30 | using distributed training 31 | """ 32 | if not dist.is_available(): 33 | return 34 | if not dist.is_initialized(): 35 | return 36 | world_size = dist.get_world_size() 37 | if world_size == 1: 38 | return 39 | dist.barrier() 40 | 41 | 42 | def _encode(encoded_data, data): 43 | # gets a byte representation for the data 44 | encoded_bytes = pickle.dumps(data) 45 | # convert this byte string into a byte tensor 46 | storage = torch.ByteStorage.from_buffer(encoded_bytes) 47 | tensor = torch.ByteTensor(storage).to("cuda") 48 | # encoding: first byte is the size and then rest is the data 49 | s = tensor.numel() 50 | assert s <= 255, "Can't encode data greater than 255 bytes" 51 | # put the encoded data in encoded_data 52 | encoded_data[0] = s 53 | encoded_data[1: (s + 1)] = tensor 54 | 55 | 56 | def all_gather(data): 57 | """ 58 | Run all_gather on arbitrary picklable data (not necessarily tensors) 59 | Args: 60 | data: any picklable object 61 | Returns: 62 | list[data]: list of data gathered from each rank 63 | """ 64 | world_size = get_world_size() 65 | if world_size == 1: 66 | return [data] 67 | 68 | # serialized to a Tensor 69 | buffer = pickle.dumps(data) 70 | storage = torch.ByteStorage.from_buffer(buffer) 71 | tensor = torch.ByteTensor(storage).to("cuda") 72 | 73 | # obtain Tensor size of each rank 74 | local_size = torch.LongTensor([tensor.numel()]).to("cuda") 75 | size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)] 76 | dist.all_gather(size_list, local_size) 77 | size_list = [int(size.item()) for size in size_list] 78 | max_size = max(size_list) 79 | 80 | # receiving Tensor from all ranks 81 | # we pad the tensor because torch all_gather does not support 82 | # gathering tensors of different shapes 83 | tensor_list = [] 84 | for _ in size_list: 85 | tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda")) 86 | if local_size != max_size: 87 | padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda") 88 | tensor = torch.cat((tensor, padding), dim=0) 89 | dist.all_gather(tensor_list, tensor) 90 | 91 | data_list = [] 92 | for size, tensor in zip(size_list, tensor_list): 93 | buffer = tensor.cpu().numpy().tobytes()[:size] 94 | data_list.append(pickle.loads(buffer)) 95 | 96 | return data_list 97 | -------------------------------------------------------------------------------- /ssd/utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import sys 4 | 5 | 6 | def setup_logger(name, distributed_rank, save_dir=None): 7 | logger = logging.getLogger(name) 8 | logger.setLevel(logging.DEBUG) 9 | # don't log results for the non-master process 10 | if distributed_rank > 0: 11 | return logger 12 | stream_handler = logging.StreamHandler(stream=sys.stdout) 13 | stream_handler.setLevel(logging.DEBUG) 14 | formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s") 15 | stream_handler.setFormatter(formatter) 16 | logger.addHandler(stream_handler) 17 | if save_dir: 18 | fh = logging.FileHandler(os.path.join(save_dir, 'log.txt')) 19 | fh.setLevel(logging.DEBUG) 20 | fh.setFormatter(formatter) 21 | logger.addHandler(fh) 22 | return logger 23 | -------------------------------------------------------------------------------- /ssd/utils/metric_logger.py: -------------------------------------------------------------------------------- 1 | from collections import deque, defaultdict 2 | import numpy as np 3 | import torch 4 | 5 | 6 | class SmoothedValue: 7 | """Track a series of values and provide access to smoothed values over a 8 | window or the global series average. 9 | """ 10 | 11 | def __init__(self, window_size=10): 12 | self.deque = deque(maxlen=window_size) 13 | self.value = np.nan 14 | self.series = [] 15 | self.total = 0.0 16 | self.count = 0 17 | 18 | def update(self, value): 19 | self.deque.append(value) 20 | self.series.append(value) 21 | self.count += 1 22 | self.total += value 23 | self.value = value 24 | 25 | @property 26 | def median(self): 27 | values = np.array(self.deque) 28 | return np.median(values) 29 | 30 | @property 31 | def avg(self): 32 | values = np.array(self.deque) 33 | return np.mean(values) 34 | 35 | @property 36 | def global_avg(self): 37 | return self.total / self.count 38 | 39 | 40 | class MetricLogger: 41 | def __init__(self, delimiter=", "): 42 | self.meters = defaultdict(SmoothedValue) 43 | self.delimiter = delimiter 44 | 45 | def update(self, **kwargs): 46 | for k, v in kwargs.items(): 47 | if isinstance(v, torch.Tensor): 48 | v = v.item() 49 | assert isinstance(v, (float, int)) 50 | self.meters[k].update(v) 51 | 52 | def __getattr__(self, attr): 53 | if attr in self.meters: 54 | return self.meters[attr] 55 | if attr in self.__dict__: 56 | return self.__dict__[attr] 57 | raise AttributeError("'{}' object has no attribute '{}'".format( 58 | type(self).__name__, attr)) 59 | 60 | def __str__(self): 61 | loss_str = [] 62 | for name, meter in self.meters.items(): 63 | loss_str.append( 64 | "{}: {:.3f} ({:.3f})".format(name, meter.avg, meter.global_avg) 65 | ) 66 | return self.delimiter.join(loss_str) 67 | -------------------------------------------------------------------------------- /ssd/utils/misc.py: -------------------------------------------------------------------------------- 1 | import errno 2 | import os 3 | 4 | 5 | def str2bool(s): 6 | return s.lower() in ('true', '1') 7 | 8 | 9 | def mkdir(path): 10 | try: 11 | os.makedirs(path) 12 | except OSError as e: 13 | if e.errno != errno.EEXIST: 14 | raise 15 | -------------------------------------------------------------------------------- /ssd/utils/model_zoo.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import os 3 | import sys 4 | 5 | import torch 6 | 7 | from ssd.utils.dist_util import is_main_process, synchronize 8 | 9 | try: 10 | from torch.hub import _download_url_to_file 11 | from torch.hub import urlparse 12 | from torch.hub import HASH_REGEX 13 | except ImportError: 14 | from torch.utils.model_zoo import _download_url_to_file 15 | from torch.utils.model_zoo import urlparse 16 | from torch.utils.model_zoo import HASH_REGEX 17 | 18 | 19 | # very similar to https://github.com/pytorch/pytorch/blob/master/torch/utils/model_zoo.py 20 | # but with a few improvements and modifications 21 | def cache_url(url, model_dir=None, progress=True): 22 | r"""Loads the Torch serialized object at the given URL. 23 | If the object is already present in `model_dir`, it's deserialized and 24 | returned. The filename part of the URL should follow the naming convention 25 | ``filename-.ext`` where ```` is the first eight or more 26 | digits of the SHA256 hash of the contents of the file. The hash is used to 27 | ensure unique names and to verify the contents of the file. 28 | The default value of `model_dir` is ``$TORCH_HOME/models`` where 29 | ``$TORCH_HOME`` defaults to ``~/.torch``. The default directory can be 30 | overridden with the ``$TORCH_MODEL_ZOO`` environment variable. 31 | Args: 32 | url (string): URL of the object to download 33 | model_dir (string, optional): directory in which to save the object 34 | progress (bool, optional): whether or not to display a progress bar to stderr 35 | Example: 36 | >>> cached_file = maskrcnn_benchmark.utils.model_zoo.cache_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth') 37 | """ 38 | if model_dir is None: 39 | torch_home = os.path.expanduser(os.getenv("TORCH_HOME", "~/.torch")) 40 | model_dir = os.getenv("TORCH_MODEL_ZOO", os.path.join(torch_home, "models")) 41 | if not os.path.exists(model_dir): 42 | os.makedirs(model_dir) 43 | parts = urlparse(url) 44 | filename = os.path.basename(parts.path) 45 | if filename == "model_final.pkl": 46 | # workaround as pre-trained Caffe2 models from Detectron have all the same filename 47 | # so make the full path the filename by replacing / with _ 48 | filename = parts.path.replace("/", "_") 49 | cached_file = os.path.join(model_dir, filename) 50 | if not os.path.exists(cached_file) and is_main_process(): 51 | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) 52 | hash_prefix = HASH_REGEX.search(filename) 53 | if hash_prefix is not None: 54 | hash_prefix = hash_prefix.group(1) 55 | # workaround: Caffe2 models don't have a hash, but follow the R-50 convention, 56 | # which matches the hash PyTorch uses. So we skip the hash matching 57 | # if the hash_prefix is less than 6 characters 58 | if len(hash_prefix) < 6: 59 | hash_prefix = None 60 | _download_url_to_file(url, cached_file, hash_prefix, progress=progress) 61 | synchronize() 62 | return cached_file 63 | 64 | 65 | def load_state_dict_from_url(url, map_location='cpu'): 66 | cached_file = cache_url(url) 67 | return torch.load(cached_file, map_location=map_location) 68 | -------------------------------------------------------------------------------- /ssd/utils/nms.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import warnings 3 | import torchvision 4 | 5 | try: 6 | import torch 7 | import torch_extension 8 | 9 | _nms = torch_extension.nms 10 | except ImportError: 11 | if torchvision.__version__ >= '0.3.0': 12 | _nms = torchvision.ops.nms 13 | else: 14 | warnings.warn('No NMS is available. Please upgrade torchvision to 0.3.0+ or compile c++ NMS ' 15 | 'using `cd ext & python build.py build_ext develop`') 16 | sys.exit(-1) 17 | 18 | 19 | def nms(boxes, scores, nms_thresh): 20 | """ Performs non-maximum suppression, run on GPU or CPU according to 21 | boxes's device. 22 | Args: 23 | boxes(Tensor[N, 4]): boxes in (x1, y1, x2, y2) format, use absolute coordinates(or relative coordinates) 24 | scores(Tensor[N]): scores 25 | nms_thresh(float): thresh 26 | Returns: 27 | indices kept. 28 | """ 29 | keep = _nms(boxes, scores, nms_thresh) 30 | return keep 31 | 32 | 33 | def batched_nms(boxes, scores, idxs, iou_threshold): 34 | """ 35 | Performs non-maximum suppression in a batched fashion. 36 | 37 | Each index value correspond to a category, and NMS 38 | will not be applied between elements of different categories. 39 | 40 | Parameters 41 | ---------- 42 | boxes : Tensor[N, 4] 43 | boxes where NMS will be performed. They 44 | are expected to be in (x1, y1, x2, y2) format 45 | scores : Tensor[N] 46 | scores for each one of the boxes 47 | idxs : Tensor[N] 48 | indices of the categories for each one of the boxes. 49 | iou_threshold : float 50 | discards all overlapping boxes 51 | with IoU < iou_threshold 52 | 53 | Returns 54 | ------- 55 | keep : Tensor 56 | int64 tensor with the indices of 57 | the elements that have been kept by NMS, sorted 58 | in decreasing order of scores 59 | """ 60 | if boxes.numel() == 0: 61 | return torch.empty((0,), dtype=torch.int64, device=boxes.device) 62 | # strategy: in order to perform NMS independently per class. 63 | # we add an offset to all the boxes. The offset is dependent 64 | # only on the class idx, and is large enough so that boxes 65 | # from different classes do not overlap 66 | max_coordinate = boxes.max() 67 | offsets = idxs.to(boxes) * (max_coordinate + 1) 68 | boxes_for_nms = boxes + offsets[:, None] 69 | keep = nms(boxes_for_nms, scores, iou_threshold) 70 | return keep 71 | -------------------------------------------------------------------------------- /ssd/utils/registry.py: -------------------------------------------------------------------------------- 1 | def _register_generic(module_dict, module_name, module): 2 | assert module_name not in module_dict 3 | module_dict[module_name] = module 4 | 5 | 6 | class Registry(dict): 7 | """ 8 | A helper class for managing registering modules, it extends a dictionary 9 | and provides a register functions. 10 | Eg. creating a registry: 11 | some_registry = Registry({"default": default_module}) 12 | There're two ways of registering new modules: 13 | 1): normal way is just calling register function: 14 | def foo(): 15 | ... 16 | some_registry.register("foo_module", foo) 17 | 2): used as decorator when declaring the module: 18 | @some_registry.register("foo_module") 19 | @some_registry.register("foo_module_nickname") 20 | def foo(): 21 | ... 22 | Access of module is just like using a dictionary, eg: 23 | f = some_registry["foo_module"] 24 | """ 25 | 26 | def __init__(self, *args, **kwargs): 27 | super(Registry, self).__init__(*args, **kwargs) 28 | 29 | def register(self, module_name, module=None): 30 | # used as function call 31 | if module is not None: 32 | _register_generic(self, module_name, module) 33 | return 34 | 35 | # used as decorator 36 | def register_fn(fn): 37 | _register_generic(self, module_name, fn) 38 | return fn 39 | 40 | return register_fn 41 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import os 4 | 5 | import torch 6 | import torch.utils.data 7 | 8 | from ssd.config import cfg 9 | from ssd.engine.inference import do_evaluation 10 | from ssd.modeling.detector import build_detection_model 11 | from ssd.utils import dist_util 12 | from ssd.utils.checkpoint import CheckPointer 13 | from ssd.utils.dist_util import synchronize 14 | from ssd.utils.logger import setup_logger 15 | 16 | 17 | def evaluation(cfg, ckpt, distributed): 18 | logger = logging.getLogger("SSD.inference") 19 | 20 | model = build_detection_model(cfg) 21 | checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR, logger=logger) 22 | device = torch.device(cfg.MODEL.DEVICE) 23 | model.to(device) 24 | checkpointer.load(ckpt, use_latest=ckpt is None) 25 | do_evaluation(cfg, model, distributed) 26 | 27 | 28 | def main(): 29 | parser = argparse.ArgumentParser(description='SSD Evaluation on VOC and COCO dataset.') 30 | parser.add_argument( 31 | "--config-file", 32 | default="", 33 | metavar="FILE", 34 | help="path to config file", 35 | type=str, 36 | ) 37 | parser.add_argument("--local_rank", type=int, default=0) 38 | parser.add_argument( 39 | "--ckpt", 40 | help="The path to the checkpoint for test, default is the latest checkpoint.", 41 | default=None, 42 | type=str, 43 | ) 44 | 45 | parser.add_argument("--output_dir", default="eval_results", type=str, help="The directory to store evaluation results.") 46 | 47 | parser.add_argument( 48 | "opts", 49 | help="Modify config options using the command-line", 50 | default=None, 51 | nargs=argparse.REMAINDER, 52 | ) 53 | args = parser.parse_args() 54 | 55 | num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 56 | distributed = num_gpus > 1 57 | 58 | if torch.cuda.is_available(): 59 | # This flag allows you to enable the inbuilt cudnn auto-tuner to 60 | # find the best algorithm to use for your hardware. 61 | torch.backends.cudnn.benchmark = True 62 | if distributed: 63 | torch.cuda.set_device(args.local_rank) 64 | torch.distributed.init_process_group(backend="nccl", init_method="env://") 65 | synchronize() 66 | 67 | cfg.merge_from_file(args.config_file) 68 | cfg.merge_from_list(args.opts) 69 | cfg.freeze() 70 | 71 | logger = setup_logger("SSD", dist_util.get_rank(), cfg.OUTPUT_DIR) 72 | logger.info("Using {} GPUs".format(num_gpus)) 73 | logger.info(args) 74 | 75 | logger.info("Loaded configuration file {}".format(args.config_file)) 76 | with open(args.config_file, "r") as cf: 77 | config_str = "\n" + cf.read() 78 | logger.info(config_str) 79 | logger.info("Running with config:\n{}".format(cfg)) 80 | evaluation(cfg, ckpt=args.ckpt, distributed=distributed) 81 | 82 | 83 | if __name__ == '__main__': 84 | main() 85 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import os 4 | 5 | import torch 6 | import torch.distributed as dist 7 | 8 | from ssd.engine.inference import do_evaluation 9 | from ssd.config import cfg 10 | from ssd.data.build import make_data_loader 11 | from ssd.engine.trainer import do_train 12 | from ssd.modeling.detector import build_detection_model 13 | from ssd.solver.build import make_optimizer, make_lr_scheduler 14 | from ssd.utils import dist_util, mkdir 15 | from ssd.utils.checkpoint import CheckPointer 16 | from ssd.utils.dist_util import synchronize 17 | from ssd.utils.logger import setup_logger 18 | from ssd.utils.misc import str2bool 19 | #export CUDA_VISIBLE_DEVICES=1 20 | 21 | def train(cfg, args): 22 | logger = logging.getLogger('SSD.trainer') 23 | model = build_detection_model(cfg) 24 | device = torch.device(cfg.MODEL.DEVICE) 25 | model.to(device) 26 | if args.distributed: 27 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) 28 | 29 | lr = cfg.SOLVER.LR * args.num_gpus # scale by num gpus 30 | optimizer = make_optimizer(cfg, model, lr) 31 | 32 | milestones = [step // args.num_gpus for step in cfg.SOLVER.LR_STEPS] 33 | scheduler = make_lr_scheduler(cfg, optimizer, milestones) 34 | 35 | arguments = {"iteration": 0} 36 | save_to_disk = dist_util.get_rank() == 0 37 | checkpointer = CheckPointer(model, optimizer, scheduler, cfg.OUTPUT_DIR, save_to_disk, logger) 38 | extra_checkpoint_data = checkpointer.load() 39 | arguments.update(extra_checkpoint_data) 40 | 41 | max_iter = cfg.SOLVER.MAX_ITER // args.num_gpus 42 | train_loader = make_data_loader(cfg, is_train=True, distributed=args.distributed, max_iter=max_iter, start_iter=arguments['iteration']) 43 | 44 | model = do_train(cfg, model, train_loader, optimizer, scheduler, checkpointer, device, arguments, args) 45 | return model 46 | 47 | 48 | def main(): 49 | parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training With PyTorch') 50 | parser.add_argument( 51 | "--config-file", 52 | default='configs/resnet50_300.yaml',#"configs/vgg_ssd512_voc0712.yaml",efficient_net_b3_ssd300_voc0712.yaml 53 | metavar="FILE", 54 | help="path to config file", 55 | type=str, 56 | ) 57 | parser.add_argument("--local_rank", type=int, default=0) 58 | parser.add_argument('--log_step', default=20, type=int, help='Print logs every log_step') 59 | parser.add_argument('--save_step', default=5000, type=int, help='Save checkpoint every save_step') 60 | parser.add_argument('--eval_step', default=-1, type=int, help='Evaluate dataset every eval_step, disabled when eval_step < 0') 61 | parser.add_argument('--use_tensorboard', default=True, type=str2bool) 62 | parser.add_argument( 63 | "--skip-test", 64 | dest="skip_test", 65 | help="Do not test the final model", 66 | action="store_true", 67 | ) 68 | parser.add_argument( 69 | "opts", 70 | help="Modify config options using the command-line", 71 | default=None, 72 | nargs=argparse.REMAINDER, 73 | ) 74 | args = parser.parse_args() 75 | num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 76 | args.distributed = num_gpus > 1 77 | args.num_gpus = num_gpus 78 | 79 | if torch.cuda.is_available(): 80 | # This flag allows you to enable the inbuilt cudnn auto-tuner to 81 | # find the best algorithm to use for your hardware. 82 | torch.backends.cudnn.benchmark = True 83 | if args.distributed: 84 | torch.cuda.set_device(args.local_rank) 85 | torch.distributed.init_process_group(backend="nccl", init_method="env://") 86 | synchronize() 87 | 88 | cfg.merge_from_file(args.config_file) 89 | cfg.merge_from_list(args.opts) 90 | cfg.freeze() 91 | 92 | if cfg.OUTPUT_DIR: 93 | mkdir(cfg.OUTPUT_DIR) 94 | 95 | logger = setup_logger("SSD", dist_util.get_rank(), cfg.OUTPUT_DIR) 96 | logger.info("Using {} GPUs".format(num_gpus)) 97 | logger.info(args) 98 | 99 | logger.info("Loaded configuration file {}".format(args.config_file)) 100 | with open(args.config_file, "r") as cf: 101 | config_str = "\n" + cf.read() 102 | logger.info(config_str) 103 | logger.info("Running with config:\n{}".format(cfg)) 104 | 105 | model = train(cfg, args) 106 | 107 | if not args.skip_test: 108 | logger.info('Start evaluating...') 109 | torch.cuda.empty_cache() # speed up evaluating after training finished 110 | do_evaluation(cfg, model, distributed=args.distributed) 111 | 112 | 113 | if __name__ == '__main__': 114 | main() 115 | --------------------------------------------------------------------------------