├── LICENSE ├── README.md ├── dataset └── dataset.py ├── demo ├── demo.gif ├── images │ ├── 12826985603_075fcfd119_h.jpg │ ├── 6958051094_0ea6dac630_w.jpg │ ├── ILSVRC2012_test_00002969.jpg │ ├── ILSVRC2012_test_00007857.jpg │ └── ILSVRC2012_test_00043529.jpg ├── img.lst └── predictions │ ├── ILSVRC2012_test_00002969_edge.png │ ├── ILSVRC2012_test_00002969_sal.png │ ├── ILSVRC2012_test_00002969_skel.png │ ├── ILSVRC2012_test_00007857_edge.png │ ├── ILSVRC2012_test_00007857_sal.png │ ├── ILSVRC2012_test_00007857_skel.png │ ├── ILSVRC2012_test_00043529_edge.png │ ├── ILSVRC2012_test_00043529_sal.png │ └── ILSVRC2012_test_00043529_skel.png ├── main.py ├── networks ├── dfi.py └── resnet.py └── solver.py /LICENSE: -------------------------------------------------------------------------------- 1 | Attribution-NonCommercial-ShareAlike 4.0 International 2 | 3 | ======================================================================= 4 | 5 | Creative Commons Corporation ("Creative Commons") is not a law firm and 6 | does not provide legal services or legal advice. Distribution of 7 | Creative Commons public licenses does not create a lawyer-client or 8 | other relationship. Creative Commons makes its licenses and related 9 | information available on an "as-is" basis. Creative Commons gives no 10 | warranties regarding its licenses, any material licensed under their 11 | terms and conditions, or any related information. Creative Commons 12 | disclaims all liability for damages resulting from their use to the 13 | fullest extent possible. 14 | 15 | Using Creative Commons Public Licenses 16 | 17 | Creative Commons public licenses provide a standard set of terms and 18 | conditions that creators and other rights holders may use to share 19 | original works of authorship and other material subject to copyright 20 | and certain other rights specified in the public license below. The 21 | following considerations are for informational purposes only, are not 22 | exhaustive, and do not form part of our licenses. 23 | 24 | Considerations for licensors: Our public licenses are 25 | intended for use by those authorized to give the public 26 | permission to use material in ways otherwise restricted by 27 | copyright and certain other rights. Our licenses are 28 | irrevocable. Licensors should read and understand the terms 29 | and conditions of the license they choose before applying it. 30 | Licensors should also secure all rights necessary before 31 | applying our licenses so that the public can reuse the 32 | material as expected. Licensors should clearly mark any 33 | material not subject to the license. This includes other CC- 34 | licensed material, or material used under an exception or 35 | limitation to copyright. More considerations for licensors: 36 | wiki.creativecommons.org/Considerations_for_licensors 37 | 38 | Considerations for the public: By using one of our public 39 | licenses, a licensor grants the public permission to use the 40 | licensed material under specified terms and conditions. If 41 | the licensor's permission is not necessary for any reason--for 42 | example, because of any applicable exception or limitation to 43 | copyright--then that use is not regulated by the license. Our 44 | licenses grant only permissions under copyright and certain 45 | other rights that a licensor has authority to grant. Use of 46 | the licensed material may still be restricted for other 47 | reasons, including because others have copyright or other 48 | rights in the material. A licensor may make special requests, 49 | such as asking that all changes be marked or described. 50 | Although not required by our licenses, you are encouraged to 51 | respect those requests where reasonable. More considerations 52 | for the public: 53 | wiki.creativecommons.org/Considerations_for_licensees 54 | 55 | ======================================================================= 56 | 57 | Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 58 | Public License 59 | 60 | By exercising the Licensed Rights (defined below), You accept and agree 61 | to be bound by the terms and conditions of this Creative Commons 62 | Attribution-NonCommercial-ShareAlike 4.0 International Public License 63 | ("Public License"). To the extent this Public License may be 64 | interpreted as a contract, You are granted the Licensed Rights in 65 | consideration of Your acceptance of these terms and conditions, and the 66 | Licensor grants You such rights in consideration of benefits the 67 | Licensor receives from making the Licensed Material available under 68 | these terms and conditions. 69 | 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. BY-NC-SA Compatible License means a license listed at 88 | creativecommons.org/compatiblelicenses, approved by Creative 89 | Commons as essentially the equivalent of this Public License. 90 | 91 | d. Copyright and Similar Rights means copyright and/or similar rights 92 | closely related to copyright including, without limitation, 93 | performance, broadcast, sound recording, and Sui Generis Database 94 | Rights, without regard to how the rights are labeled or 95 | categorized. For purposes of this Public License, the rights 96 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 97 | Rights. 98 | 99 | e. Effective Technological Measures means those measures that, in the 100 | absence of proper authority, may not be circumvented under laws 101 | fulfilling obligations under Article 11 of the WIPO Copyright 102 | Treaty adopted on December 20, 1996, and/or similar international 103 | agreements. 104 | 105 | f. Exceptions and Limitations means fair use, fair dealing, and/or 106 | any other exception or limitation to Copyright and Similar Rights 107 | that applies to Your use of the Licensed Material. 108 | 109 | g. License Elements means the license attributes listed in the name 110 | of a Creative Commons Public License. The License Elements of this 111 | Public License are Attribution, NonCommercial, and ShareAlike. 112 | 113 | h. Licensed Material means the artistic or literary work, database, 114 | or other material to which the Licensor applied this Public 115 | License. 116 | 117 | i. Licensed Rights means the rights granted to You subject to the 118 | terms and conditions of this Public License, which are limited to 119 | all Copyright and Similar Rights that apply to Your use of the 120 | Licensed Material and that the Licensor has authority to license. 121 | 122 | j. Licensor means the individual(s) or entity(ies) granting rights 123 | under this Public License. 124 | 125 | k. NonCommercial means not primarily intended for or directed towards 126 | commercial advantage or monetary compensation. For purposes of 127 | this Public License, the exchange of the Licensed Material for 128 | other material subject to Copyright and Similar Rights by digital 129 | file-sharing or similar means is NonCommercial provided there is 130 | no payment of monetary compensation in connection with the 131 | exchange. 132 | 133 | l. Share means to provide material to the public by any means or 134 | process that requires permission under the Licensed Rights, such 135 | as reproduction, public display, public performance, distribution, 136 | dissemination, communication, or importation, and to make material 137 | available to the public including in ways that members of the 138 | public may access the material from a place and at a time 139 | individually chosen by them. 140 | 141 | m. Sui Generis Database Rights means rights other than copyright 142 | resulting from Directive 96/9/EC of the European Parliament and of 143 | the Council of 11 March 1996 on the legal protection of databases, 144 | as amended and/or succeeded, as well as other essentially 145 | equivalent rights anywhere in the world. 146 | 147 | n. You means the individual or entity exercising the Licensed Rights 148 | under this Public License. Your has a corresponding meaning. 149 | 150 | 151 | Section 2 -- Scope. 152 | 153 | a. License grant. 154 | 155 | 1. Subject to the terms and conditions of this Public License, 156 | the Licensor hereby grants You a worldwide, royalty-free, 157 | non-sublicensable, non-exclusive, irrevocable license to 158 | exercise the Licensed Rights in the Licensed Material to: 159 | 160 | a. reproduce and Share the Licensed Material, in whole or 161 | in part, for NonCommercial purposes only; and 162 | 163 | b. produce, reproduce, and Share Adapted Material for 164 | NonCommercial purposes only. 165 | 166 | 2. Exceptions and Limitations. For the avoidance of doubt, where 167 | Exceptions and Limitations apply to Your use, this Public 168 | License does not apply, and You do not need to comply with 169 | its terms and conditions. 170 | 171 | 3. Term. The term of this Public License is specified in Section 172 | 6(a). 173 | 174 | 4. Media and formats; technical modifications allowed. The 175 | Licensor authorizes You to exercise the Licensed Rights in 176 | all media and formats whether now known or hereafter created, 177 | and to make technical modifications necessary to do so. The 178 | Licensor waives and/or agrees not to assert any right or 179 | authority to forbid You from making technical modifications 180 | necessary to exercise the Licensed Rights, including 181 | technical modifications necessary to circumvent Effective 182 | Technological Measures. For purposes of this Public License, 183 | simply making modifications authorized by this Section 2(a) 184 | (4) never produces Adapted Material. 185 | 186 | 5. Downstream recipients. 187 | 188 | a. Offer from the Licensor -- Licensed Material. Every 189 | recipient of the Licensed Material automatically 190 | receives an offer from the Licensor to exercise the 191 | Licensed Rights under the terms and conditions of this 192 | Public License. 193 | 194 | b. Additional offer from the Licensor -- Adapted Material. 195 | Every recipient of Adapted Material from You 196 | automatically receives an offer from the Licensor to 197 | exercise the Licensed Rights in the Adapted Material 198 | under the conditions of the Adapter's License You apply. 199 | 200 | c. No downstream restrictions. You may not offer or impose 201 | any additional or different terms or conditions on, or 202 | apply any Effective Technological Measures to, the 203 | Licensed Material if doing so restricts exercise of the 204 | Licensed Rights by any recipient of the Licensed 205 | Material. 206 | 207 | 6. No endorsement. Nothing in this Public License constitutes or 208 | may be construed as permission to assert or imply that You 209 | are, or that Your use of the Licensed Material is, connected 210 | with, or sponsored, endorsed, or granted official status by, 211 | the Licensor or others designated to receive attribution as 212 | provided in Section 3(a)(1)(A)(i). 213 | 214 | b. Other rights. 215 | 216 | 1. Moral rights, such as the right of integrity, are not 217 | licensed under this Public License, nor are publicity, 218 | privacy, and/or other similar personality rights; however, to 219 | the extent possible, the Licensor waives and/or agrees not to 220 | assert any such rights held by the Licensor to the limited 221 | extent necessary to allow You to exercise the Licensed 222 | Rights, but not otherwise. 223 | 224 | 2. Patent and trademark rights are not licensed under this 225 | Public License. 226 | 227 | 3. To the extent possible, the Licensor waives any right to 228 | collect royalties from You for the exercise of the Licensed 229 | Rights, whether directly or through a collecting society 230 | under any voluntary or waivable statutory or compulsory 231 | licensing scheme. In all other cases the Licensor expressly 232 | reserves any right to collect such royalties, including when 233 | the Licensed Material is used other than for NonCommercial 234 | purposes. 235 | 236 | 237 | Section 3 -- License Conditions. 238 | 239 | Your exercise of the Licensed Rights is expressly made subject to the 240 | following conditions. 241 | 242 | a. Attribution. 243 | 244 | 1. If You Share the Licensed Material (including in modified 245 | form), You must: 246 | 247 | a. retain the following if it is supplied by the Licensor 248 | with the Licensed Material: 249 | 250 | i. identification of the creator(s) of the Licensed 251 | Material and any others designated to receive 252 | attribution, in any reasonable manner requested by 253 | the Licensor (including by pseudonym if 254 | designated); 255 | 256 | ii. a copyright notice; 257 | 258 | iii. a notice that refers to this Public License; 259 | 260 | iv. a notice that refers to the disclaimer of 261 | warranties; 262 | 263 | v. a URI or hyperlink to the Licensed Material to the 264 | extent reasonably practicable; 265 | 266 | b. indicate if You modified the Licensed Material and 267 | retain an indication of any previous modifications; and 268 | 269 | c. indicate the Licensed Material is licensed under this 270 | Public License, and include the text of, or the URI or 271 | hyperlink to, this Public License. 272 | 273 | 2. You may satisfy the conditions in Section 3(a)(1) in any 274 | reasonable manner based on the medium, means, and context in 275 | which You Share the Licensed Material. For example, it may be 276 | reasonable to satisfy the conditions by providing a URI or 277 | hyperlink to a resource that includes the required 278 | information. 279 | 3. If requested by the Licensor, You must remove any of the 280 | information required by Section 3(a)(1)(A) to the extent 281 | reasonably practicable. 282 | 283 | b. ShareAlike. 284 | 285 | In addition to the conditions in Section 3(a), if You Share 286 | Adapted Material You produce, the following conditions also apply. 287 | 288 | 1. The Adapter's License You apply must be a Creative Commons 289 | license with the same License Elements, this version or 290 | later, or a BY-NC-SA Compatible License. 291 | 292 | 2. You must include the text of, or the URI or hyperlink to, the 293 | Adapter's License You apply. You may satisfy this condition 294 | in any reasonable manner based on the medium, means, and 295 | context in which You Share Adapted Material. 296 | 297 | 3. You may not offer or impose any additional or different terms 298 | or conditions on, or apply any Effective Technological 299 | Measures to, Adapted Material that restrict exercise of the 300 | rights granted under the Adapter's License You apply. 301 | 302 | 303 | Section 4 -- Sui Generis Database Rights. 304 | 305 | Where the Licensed Rights include Sui Generis Database Rights that 306 | apply to Your use of the Licensed Material: 307 | 308 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 309 | to extract, reuse, reproduce, and Share all or a substantial 310 | portion of the contents of the database for NonCommercial purposes 311 | only; 312 | 313 | b. if You include all or a substantial portion of the database 314 | contents in a database in which You have Sui Generis Database 315 | Rights, then the database in which You have Sui Generis Database 316 | Rights (but not its individual contents) is Adapted Material, 317 | including for purposes of Section 3(b); and 318 | 319 | c. You must comply with the conditions in Section 3(a) if You Share 320 | all or a substantial portion of the contents of the database. 321 | 322 | For the avoidance of doubt, this Section 4 supplements and does not 323 | replace Your obligations under this Public License where the Licensed 324 | Rights include other Copyright and Similar Rights. 325 | 326 | 327 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 328 | 329 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 330 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 331 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 332 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 333 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 334 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 335 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 336 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 337 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 338 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 339 | 340 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 341 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 342 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 343 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 344 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 345 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 346 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 347 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 348 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 349 | 350 | c. The disclaimer of warranties and limitation of liability provided 351 | above shall be interpreted in a manner that, to the extent 352 | possible, most closely approximates an absolute disclaimer and 353 | waiver of all liability. 354 | 355 | 356 | Section 6 -- Term and Termination. 357 | 358 | a. This Public License applies for the term of the Copyright and 359 | Similar Rights licensed here. However, if You fail to comply with 360 | this Public License, then Your rights under this Public License 361 | terminate automatically. 362 | 363 | b. Where Your right to use the Licensed Material has terminated under 364 | Section 6(a), it reinstates: 365 | 366 | 1. automatically as of the date the violation is cured, provided 367 | it is cured within 30 days of Your discovery of the 368 | violation; or 369 | 370 | 2. upon express reinstatement by the Licensor. 371 | 372 | For the avoidance of doubt, this Section 6(b) does not affect any 373 | right the Licensor may have to seek remedies for Your violations 374 | of this Public License. 375 | 376 | c. For the avoidance of doubt, the Licensor may also offer the 377 | Licensed Material under separate terms or conditions or stop 378 | distributing the Licensed Material at any time; however, doing so 379 | will not terminate this Public License. 380 | 381 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 382 | License. 383 | 384 | 385 | Section 7 -- Other Terms and Conditions. 386 | 387 | a. The Licensor shall not be bound by any additional or different 388 | terms or conditions communicated by You unless expressly agreed. 389 | 390 | b. Any arrangements, understandings, or agreements regarding the 391 | Licensed Material not stated herein are separate from and 392 | independent of the terms and conditions of this Public License. 393 | 394 | 395 | Section 8 -- Interpretation. 396 | 397 | a. For the avoidance of doubt, this Public License does not, and 398 | shall not be interpreted to, reduce, limit, restrict, or impose 399 | conditions on any use of the Licensed Material that could lawfully 400 | be made without permission under this Public License. 401 | 402 | b. To the extent possible, if any provision of this Public License is 403 | deemed unenforceable, it shall be automatically reformed to the 404 | minimum extent necessary to make it enforceable. If the provision 405 | cannot be reformed, it shall be severed from this Public License 406 | without affecting the enforceability of the remaining terms and 407 | conditions. 408 | 409 | c. No term or condition of this Public License will be waived and no 410 | failure to comply consented to unless expressly agreed to by the 411 | Licensor. 412 | 413 | d. Nothing in this Public License constitutes or may be interpreted 414 | as a limitation upon, or waiver of, any privileges and immunities 415 | that apply to the Licensor or You, including from the legal 416 | processes of any jurisdiction or authority. 417 | 418 | ======================================================================= 419 | 420 | Creative Commons is not a party to its public 421 | licenses. Notwithstanding, Creative Commons may elect to apply one of 422 | its public licenses to material it publishes and in those instances 423 | will be considered the “Licensor.” The text of the Creative Commons 424 | public licenses is dedicated to the public domain under the CC0 Public 425 | Domain Dedication. Except for the limited purpose of indicating that 426 | material is shared under a Creative Commons public license or as 427 | otherwise permitted by the Creative Commons policies published at 428 | creativecommons.org/policies, Creative Commons does not authorize the 429 | use of the trademark "Creative Commons" or any other trademark or logo 430 | of Creative Commons without its prior written consent including, 431 | without limitation, in connection with any unauthorized modifications 432 | to any of its public licenses or any other arrangements, 433 | understandings, or agreements concerning use of licensed material. For 434 | the avoidance of doubt, this paragraph does not form part of the 435 | public licenses. 436 | 437 | Creative Commons may be contacted at creativecommons.org. 438 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Dynamic Feature Integration for Simultaneous Detection of Salient Object, Edge and Skeleton 2 | 3 | ### This is a demo PyTorch implementation of our IEEE TIP 2020 [paper](http://mftp.mmcheng.net/Papers/20TIP-DFI.pdf). 4 | ### We also provide an [Online Demo](http://mc.nankai.edu.cn/dfi). 5 | 6 |

7 | animated 8 |

9 | 10 | ## Prerequisites 11 | 12 | - [Pytorch 0.4.1+](http://pytorch.org/) 13 | - [torchvision](http://pytorch.org/) 14 | - [opencv](https://opencv.org/) 15 | 16 | 17 | ## Demo usage 18 | ### 1. Clone the repository 19 | ```shell 20 | git clone https://github.com/backseason/DFI.git 21 | cd DFI/ 22 | ``` 23 | 24 | ### 2. Download the pretrained model 25 | `dfi.pth` [GoogleDrive](https://drive.google.com/file/d/1N29cJghKKJOHbKgpwR2_Ui64umCE-XG3/view?usp=sharing) | [BaiduYun](https://pan.baidu.com/s/1WPQiUPo7t8REK3LtmG9_KA) (pwd: **wkeb**) 26 | and move it to the `pretrained` folder. 27 | 28 | 29 | ### 3. Test (demo) 30 | The source images are in the `demo/images` folder. 31 | By running 32 | ```shell 33 | python main.py 34 | ``` 35 | you'll get the predictions under 36 | the `demo/predictions` folder. The predictions of all the three tasks are performed simultaneously. 37 | 38 | 39 | ### 4. Pre-computed results and evaluation results 40 | 41 | You can find the pre-computed predictions maps of all the three tasks and 42 | their corresponding evaluation scores with 43 | the following link: 44 | `Results reported in the paper` [GoogleDrive](https://drive.google.com/file/d/17SBs3v3h_FnImbHOZk0zy4JzDUHSK1zv/view?usp=sharing) | [BaiduYun](https://pan.baidu.com/s/1WP3WP5oaNWRuaUcKH4oZ7g) (pwd: **7eg3**) 45 | 46 | ### 5. Contact 47 | If you have any questions, feel free to contact me via: `j04.liu(at)gmail.com`. 48 | 49 | ### If you think this work is helpful, please cite 50 | ```latex 51 | @article{liu2020dynamic, 52 | title={Dynamic Feature Integration for Simultaneous Detection of Salient Object, Edge and Skeleton}, 53 | author={Jiang-Jiang Liu and Qibin Hou and Ming-Ming Cheng}, 54 | journal={IEEE Transactions on Image Processing}, 55 | year={2020}, 56 | volume={}, 57 | number={}, 58 | pages={1-15}, 59 | doi={10.1109/TIP.2020.3017352}, 60 | } 61 | ``` 62 | -------------------------------------------------------------------------------- /dataset/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.utils import data 4 | from torchvision.transforms import functional as F 5 | import cv2 6 | import numpy as np 7 | from PIL import Image 8 | 9 | class ImageDataTest(data.Dataset): 10 | def __init__(self, test_mode=1, sal_mode='e'): 11 | if test_mode == 0: 12 | self.image_root = './data/HED-BSDS_PASCAL/HED-BSDS/test/' 13 | self.image_source = './data/HED-BSDS_PASCAL/HED-BSDS/test.lst' 14 | elif test_mode == 1: 15 | if sal_mode == 'e': 16 | self.image_root = './data/ECSSD/Imgs/' 17 | self.image_source = './data/ECSSD/test.lst' 18 | elif sal_mode == 'p': 19 | self.image_root = './data/PASCALS/Imgs/' 20 | self.image_source = './data/PASCALS/test.lst' 21 | elif sal_mode == 'd': 22 | self.image_root = './data/DUTOMRON/Imgs/' 23 | self.image_source = './data/DUTOMRON/test.lst' 24 | elif sal_mode == 'h': 25 | self.image_root = './data/HKU-IS/Imgs/' 26 | self.image_source = './data/HKU-IS/test.lst' 27 | elif sal_mode == 's': 28 | self.image_root = './data/SOD/Imgs/' 29 | self.image_source = './data/SOD/test.lst' 30 | elif sal_mode == 't': 31 | self.image_root = './data/DUTS-TE/Imgs/' 32 | self.image_source = './data/DUTS-TE/test.lst' 33 | elif test_mode == 2: 34 | self.image_root = './data/SK-LARGE/images/test/' 35 | self.image_source = './data/SK-LARGE/test.lst' 36 | elif test_mode == 3: 37 | self.image_root = './demo/images/' 38 | self.image_source = './demo/img.lst' 39 | 40 | with open(self.image_source, 'r') as f: 41 | self.image_list = [x.strip() for x in f.readlines()] 42 | 43 | self.image_num = len(self.image_list) 44 | 45 | def __getitem__(self, item): 46 | image, im_size = load_image_test(os.path.join(self.image_root, self.image_list[item])) 47 | image = torch.Tensor(image) 48 | 49 | return {'image': image, 'name': self.image_list[item%self.image_num], 'size': im_size} 50 | 51 | def __len__(self): 52 | return self.image_num 53 | 54 | def get_loader(test_mode=0, sal_mode='e', pin=False): 55 | dataset = ImageDataTest(test_mode=test_mode, sal_mode=sal_mode) 56 | data_loader = data.DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=1, 57 | pin_memory=pin) 58 | return data_loader 59 | 60 | def load_image_test(pah): 61 | if not os.path.exists(pah): 62 | print('File Not Exists') 63 | im = cv2.imread(pah) 64 | in_ = np.array(im, dtype=np.float32) 65 | im_size = tuple(in_.shape[:2]) 66 | in_ -= np.array((104.00699, 116.66877, 122.67892)) 67 | in_ = in_.transpose((2,0,1)) 68 | return in_, im_size 69 | -------------------------------------------------------------------------------- /demo/demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/demo.gif -------------------------------------------------------------------------------- /demo/images/12826985603_075fcfd119_h.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/images/12826985603_075fcfd119_h.jpg -------------------------------------------------------------------------------- /demo/images/6958051094_0ea6dac630_w.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/images/6958051094_0ea6dac630_w.jpg -------------------------------------------------------------------------------- /demo/images/ILSVRC2012_test_00002969.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/images/ILSVRC2012_test_00002969.jpg -------------------------------------------------------------------------------- /demo/images/ILSVRC2012_test_00007857.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/images/ILSVRC2012_test_00007857.jpg -------------------------------------------------------------------------------- /demo/images/ILSVRC2012_test_00043529.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/images/ILSVRC2012_test_00043529.jpg -------------------------------------------------------------------------------- /demo/img.lst: -------------------------------------------------------------------------------- 1 | ILSVRC2012_test_00007857.jpg 2 | ILSVRC2012_test_00002969.jpg 3 | ILSVRC2012_test_00043529.jpg 4 | -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00002969_edge.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00002969_edge.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00002969_sal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00002969_sal.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00002969_skel.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00002969_skel.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00007857_edge.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00007857_edge.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00007857_sal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00007857_sal.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00007857_skel.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00007857_skel.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00043529_edge.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00043529_edge.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00043529_sal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00043529_sal.png -------------------------------------------------------------------------------- /demo/predictions/ILSVRC2012_test_00043529_skel.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/backseason/DFI/213f13a47b921d8c2f2ceb038d678ebc59785c0e/demo/predictions/ILSVRC2012_test_00043529_skel.png -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | from dataset.dataset import get_loader 4 | from solver import Solver 5 | 6 | def main(config): 7 | test_loader = get_loader(test_mode=config.test_mode, sal_mode=config.sal_mode) 8 | if not os.path.exists(config.test_fold): os.mkdir(config.test_fold) 9 | test = Solver(test_loader, config) 10 | test.test(test_mode=config.test_mode) 11 | 12 | if __name__ == '__main__': 13 | parser = argparse.ArgumentParser() 14 | 15 | # Hyper-parameters 16 | parser.add_argument('--cuda', type=bool, default=True) 17 | 18 | # Testing settings 19 | parser.add_argument('--model', type=str, default='pretrained/dfi.pth') 20 | parser.add_argument('--test_fold', type=str, default='demo/predictions') 21 | parser.add_argument('--test_mode', type=int, default=3) # choose task 22 | parser.add_argument('--sal_mode', type=str, default='e') # choose dataset, details in 'dataset/dataset.py' 23 | 24 | config = parser.parse_args() 25 | main(config) 26 | -------------------------------------------------------------------------------- /networks/dfi.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | 5 | from .resnet import resnet50_ppm 6 | 7 | config = {'converter': [[64,256,512,1024,2048,256],[32,64,128,256,256,256]], 8 | 'dfims': [[[32,64,128,256,256,256], 32, 0], [[32,64,128,256,256,256], 64, 1], [[32,64,128,256,256,256], 128, 2], [[32,64,128,256,256,256], 256, 3]], 9 | 'dfims_id': [[0,1,2,3,4,5], [0,1,2,3,4,5], [0,1,2,3,4,5], [0,1,2,3,4,5]], 10 | 'tams': [32, 64, 128, 256], 11 | 'predictors': [[[32, 64, 128, 256], True], [[32, 64, 128, 256], False], [[32, 64, 128, 256], True]], 12 | 'predictors_id': [0,1,2,3] } 13 | 14 | def gn(planes, channel_per_group=4, max_groups=32): 15 | groups = planes // channel_per_group 16 | return nn.GroupNorm(min(groups, max_groups), planes) 17 | 18 | class Converter(nn.Module): 19 | def __init__(self, list_k): 20 | super(Converter, self).__init__() 21 | up = [] 22 | for i in range(len(list_k[0])): 23 | up.append(nn.Sequential( 24 | nn.Conv2d(list_k[0][i], list_k[1][i], 1, 1, bias=False), 25 | gn(list_k[1][i]), 26 | nn.ReLU(inplace=True), 27 | )) 28 | self.convert = nn.ModuleList(up) 29 | 30 | def forward(self, x): 31 | out = [] 32 | for i in range(len(x)): 33 | out.append(self.convert[i](x[i])) 34 | return out 35 | 36 | class DFIM(nn.Module): 37 | def __init__(self, list_k, k, size_id, modes=3): 38 | super(DFIM, self).__init__() 39 | self.len = len(list_k) 40 | self.size_id = size_id 41 | up = [] 42 | for i in range(len(list_k)): 43 | up.append(nn.Sequential(nn.Conv2d(list_k[i], k, 1, 1, bias=False), gn(k))) 44 | self.merge = nn.ModuleList(up) 45 | merge_convs, fcs, convs = [], [], [] 46 | for m in range(modes): 47 | merge_convs.append(nn.Sequential( 48 | nn.Conv2d(k, k//4, 1, 1, bias=False), 49 | gn(k//4), 50 | nn.ReLU(inplace=True), 51 | nn.Conv2d(k//4, k, 1, 1, bias=False), 52 | gn(k), 53 | )) 54 | fcs.append(nn.Sequential( 55 | nn.Linear(k, k // 4, bias=False), 56 | nn.ReLU(inplace=True), 57 | nn.Linear(k // 4, self.len, bias=False), 58 | )) 59 | convs.append(nn.Sequential(nn.Conv2d(k, k, 3, 1, 1, bias=False), gn(k), nn.ReLU(inplace=True))) 60 | self.merge_convs = nn.ModuleList(merge_convs) 61 | self.fcs = nn.ModuleList(fcs) 62 | self.convs = nn.ModuleList(convs) 63 | self.gap = nn.AdaptiveAvgPool2d(1) 64 | self.softmax = nn.Softmax(dim=1) 65 | self.relu =nn.ReLU(inplace=True) 66 | 67 | def forward(self, list_x, mode=3): 68 | x_size = list_x[self.size_id].size() 69 | feas = [] 70 | for i in range(len(list_x)): 71 | feas.append(self.merge[i](F.interpolate(list_x[i], x_size[2:], mode='bilinear', align_corners=True)).unsqueeze(dim=1)) 72 | feas = torch.cat(feas, dim=1) # Nx6xCxHxW 73 | fea_sum = torch.sum(feas, dim=1) # NxCxHxW 74 | 75 | if mode == 3: 76 | outs = [] 77 | for mode_ in range(3): 78 | fea_u = self.merge_convs[mode_](fea_sum) 79 | fea_s = self.gap(fea_u).squeeze(-1).squeeze(-1) # NxC 80 | fea_z = self.fcs[mode_](fea_s) # Nx6 81 | selects = self.softmax(fea_z) # Nx6 82 | feas_f = selects.reshape(x_size[0], self.len, 1, 1, 1).expand_as(feas) * feas # Nx6xCxHxW 83 | _, index = torch.topk(selects, 3, dim=1) # Nx3 84 | selected = [] 85 | for i in range(x_size[0]): 86 | selected.append(torch.index_select(feas_f, dim=1, index=index[i])) 87 | selected = torch.cat(selected, dim=0) 88 | fea_v = selected.sum(dim=1) 89 | outs.append(self.convs[mode_](self.relu(fea_v))) 90 | return torch.cat(outs, dim=0) 91 | else: 92 | fea_u = self.merge_convs[mode](fea_sum) 93 | fea_s = self.gap(fea_u).squeeze(-1).squeeze(-1) # NxC 94 | fea_z = self.fcs[mode](fea_s) # Nx6 95 | selects = self.softmax(fea_z) # Nx6 96 | feas_f = selects.reshape(x_size[0], self.len, 1, 1, 1).expand_as(feas) * feas # Nx6xCxHxW 97 | _, index = torch.topk(selects, 3, dim=1) # Nx3 98 | selected = [] 99 | for i in range(x_size[0]): 100 | selected.append(torch.index_select(feas_f, dim=1, index=index[i])) 101 | selected = torch.cat(selected, dim=0) 102 | fea_v = selected.sum(dim=1) 103 | return self.convs[mode](self.relu(fea_v)) 104 | 105 | class TAM(nn.Module): # TAM 106 | reduction = 4 107 | def __init__(self, k): 108 | super(TAM, self).__init__() 109 | k_mid = int(k // self.reduction) 110 | self.attention = nn.Sequential( 111 | nn.Conv2d(k, k_mid, 1, 1, bias=False), 112 | gn(k_mid), 113 | nn.ReLU(inplace=True), 114 | nn.Conv2d(k_mid, k, 1, 1, bias=False), 115 | gn(k), 116 | nn.Sigmoid(), 117 | ) 118 | self.block = nn.Sequential(nn.Conv2d(k, k, 3, 1, 1, bias=False), gn(k), nn.ReLU(inplace=True)) 119 | 120 | def forward(self, x): 121 | out = self.attention(x) 122 | out = torch.add(x, torch.mul(x, out)) 123 | out = self.block(out) 124 | return out 125 | 126 | class Predictor(nn.Module): 127 | def __init__(self, list_k, deep_sup): 128 | super(Predictor, self).__init__() 129 | self.trans = nn.ModuleList() 130 | for i in range(len(list_k)): 131 | self.trans.append(nn.Conv2d(list_k[i], 1, 1, 1)) 132 | self.fuse = nn.Conv2d(len(list_k), 1, 1, 1) 133 | self.deep_sup = deep_sup 134 | 135 | def forward(self, list_x, x_size=None): 136 | up_x = [] 137 | for i, i_x in enumerate(list_x): 138 | up_x.append(F.interpolate(self.trans[i](i_x), x_size[2:], mode='bilinear', align_corners=True)) 139 | fuse = self.fuse(torch.cat(up_x, dim = 1)) 140 | if self.deep_sup: 141 | return [fuse, up_x] 142 | else: 143 | return [fuse] 144 | 145 | def extra_layer(base): 146 | converter, dfims, tams, predictors = [], [], [], [] 147 | converter = Converter(config['converter']) 148 | 149 | for k in config['dfims']: 150 | dfims += [DFIM(k[0], k[1], k[2])] 151 | 152 | for k in config['tams']: 153 | tams += [TAM(k)] 154 | 155 | for k in config['predictors']: 156 | predictors += [Predictor(k[0], k[1])] 157 | 158 | return base, converter, dfims, tams, predictors 159 | 160 | 161 | class DFI(nn.Module): 162 | def __init__(self, base, converter, dfims, tams, predictors): 163 | super(DFI, self).__init__() 164 | self.dfims_id = config['dfims_id'] 165 | self.predictors_id = config['predictors_id'] 166 | 167 | self.base = base 168 | self.converter = converter 169 | self.dfims = nn.ModuleList(dfims) 170 | self.tams = nn.ModuleList(tams) 171 | self.predictors = nn.ModuleList(predictors) 172 | 173 | def forward(self, x, mode = 3): 174 | x_size = x.size() 175 | x = self.converter(self.base(x)) 176 | 177 | # DFIM 178 | dfims = [] 179 | for k in range(len(self.dfims)): 180 | dfims.append(self.dfims[k]([x[i] for i in self.dfims_id[k]], mode=mode)) 181 | 182 | # TAM 183 | tams = [] 184 | for k in range(len(self.tams)): 185 | if k in self.predictors_id: 186 | tams.append(self.tams[k](dfims[k])) 187 | 188 | # Prediction 189 | predictions = [] 190 | if mode == 3: 191 | for mode_ in range(mode): 192 | predictions.append(self.predictors[mode_]([tam[mode_:mode_+1] for tam in tams], x_size)) 193 | else: 194 | predictions = self.predictors[mode](tams, x_size) 195 | return predictions 196 | 197 | def build_model(): 198 | return DFI(*extra_layer(resnet50_ppm())) 199 | -------------------------------------------------------------------------------- /networks/resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class Bottleneck(nn.Module): 6 | expansion = 4 7 | 8 | def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None): 9 | super(Bottleneck, self).__init__() 10 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) 11 | self.bn1 = nn.BatchNorm2d(planes) 12 | for i in self.bn1.parameters(): 13 | i.requires_grad = False 14 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, 15 | padding=dilation, dilation=dilation, bias=False) 16 | self.bn2 = nn.BatchNorm2d(planes) 17 | for i in self.bn2.parameters(): 18 | i.requires_grad = False 19 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 20 | self.bn3 = nn.BatchNorm2d(planes * 4) 21 | for i in self.bn3.parameters(): 22 | i.requires_grad = False 23 | self.relu = nn.ReLU(inplace=True) 24 | self.downsample = downsample 25 | self.stride = stride 26 | 27 | def forward(self, x): 28 | residual = x 29 | 30 | out = self.conv1(x) 31 | out = self.bn1(out) 32 | out = self.relu(out) 33 | 34 | out = self.conv2(out) 35 | out = self.bn2(out) 36 | out = self.relu(out) 37 | 38 | out = self.conv3(out) 39 | out = self.bn3(out) 40 | 41 | if self.downsample is not None: 42 | residual = self.downsample(x) 43 | 44 | out += residual 45 | out = self.relu(out) 46 | 47 | return out 48 | 49 | class ResNet(nn.Module): 50 | def __init__(self, block, layers): 51 | self.freeze_bn = True 52 | self.inplanes = 64 53 | super(ResNet, self).__init__() 54 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 55 | bias=False) 56 | self.bn1 = nn.BatchNorm2d(64) 57 | for i in self.bn1.parameters(): 58 | i.requires_grad = False 59 | self.relu = nn.ReLU(inplace=True) 60 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True) 61 | self.layer1 = self._make_layer(block, 64, layers[0]) 62 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 63 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 64 | self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=2) 65 | 66 | def _make_layer(self, block, planes, blocks, stride=1, dilation=1): 67 | downsample = None 68 | if stride != 1 or self.inplanes != planes * block.expansion or dilation == 2 or dilation == 4: 69 | downsample = nn.Sequential( 70 | nn.Conv2d(self.inplanes, planes * block.expansion, 71 | kernel_size=1, stride=stride, bias=False), 72 | nn.BatchNorm2d(planes * block.expansion), 73 | ) 74 | for i in downsample._modules['1'].parameters(): 75 | i.requires_grad = False 76 | layers = [] 77 | layers.append(block(self.inplanes, planes, stride,dilation=dilation, downsample=downsample)) 78 | self.inplanes = planes * block.expansion 79 | for i in range(1, blocks): 80 | layers.append(block(self.inplanes, planes,dilation=dilation)) 81 | 82 | return nn.Sequential(*layers) 83 | 84 | def forward(self, x): 85 | tmp_x = [] 86 | x = self.conv1(x) 87 | x = self.bn1(x) 88 | x = self.relu(x) 89 | tmp_x.append(x) 90 | x = self.maxpool(x) 91 | 92 | x = self.layer1(x) 93 | tmp_x.append(x) 94 | x = self.layer2(x) 95 | tmp_x.append(x) 96 | x = self.layer3(x) 97 | tmp_x.append(x) 98 | x = self.layer4(x) 99 | tmp_x.append(x) 100 | 101 | return tmp_x 102 | 103 | 104 | class ResNet_PPM(nn.Module): 105 | def __init__(self, block, layers): 106 | super(ResNet_PPM,self).__init__() 107 | self.resnet = ResNet(block, layers) 108 | 109 | self.in_planes = 256 110 | 111 | self.ppm_pre = nn.Sequential( 112 | nn.Conv2d(2048, self.in_planes, 1, 1, bias=False), 113 | nn.GroupNorm(32, self.in_planes), 114 | ) 115 | ppms = [] 116 | for ii in [1, 3, 5]: 117 | ppms.append(nn.Sequential( 118 | nn.AdaptiveAvgPool2d(ii), 119 | nn.Conv2d(self.in_planes, self.in_planes, 1, 1, bias=False), 120 | nn.GroupNorm(32, self.in_planes), 121 | )) 122 | self.ppms = nn.ModuleList(ppms) 123 | 124 | self.ppm_cat = nn.Sequential( 125 | nn.Conv2d(self.in_planes, self.in_planes, 3, 1, 1, bias=False), 126 | nn.GroupNorm(32, self.in_planes), 127 | ) 128 | self.relu = nn.ReLU(inplace=True) 129 | 130 | def forward(self, x): 131 | x = self.resnet(x) 132 | 133 | x_pre = self.ppm_pre(x[-1]) 134 | x_ppm = x_pre 135 | for k in range(len(self.ppms)): 136 | x_ppm = torch.add(x_ppm, F.interpolate(self.ppms[k](x_pre), x_pre.size()[2:], mode='bilinear', align_corners=True)) 137 | x_ppm = self.ppm_cat(self.relu(x_ppm)) 138 | x.append(x_ppm) 139 | 140 | return x 141 | 142 | def resnet50_ppm(): 143 | model = ResNet_PPM(Bottleneck, [3, 4, 6, 3]) 144 | return model 145 | -------------------------------------------------------------------------------- /solver.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn import utils, functional as F 3 | from networks.dfi import build_model 4 | import numpy as np 5 | import os 6 | import cv2 7 | 8 | class Solver(object): 9 | def __init__(self, data_loader, config): 10 | self.data_loader = data_loader 11 | self.config = config 12 | self.net = build_model() 13 | if self.config.cuda: 14 | self.net = self.net.cuda() 15 | print('Loading pre-trained model from %s...' % self.config.model) 16 | self.net.load_state_dict(torch.load(self.config.model)) 17 | self.net.eval() 18 | 19 | def test(self, test_mode=0): 20 | mode_name = ['edge', 'sal', 'skel'] 21 | EPSILON = 1e-8 22 | img_num = len(self.data_loader) 23 | for i, data_batch in enumerate(self.data_loader): 24 | images, name, im_size = data_batch['image'], data_batch['name'][0], np.asarray(data_batch['size']) 25 | if test_mode == 0: # edge task 26 | images = images.numpy()[0].transpose((1,2,0)) 27 | scale = [0.5, 1, 1.5, 2] # multi-scale testing as commonly done 28 | multi_fuse = np.zeros(im_size, np.float32) 29 | for k in range(0, len(scale)): 30 | im_ = cv2.resize(images, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR) 31 | im_ = im_.transpose((2, 0, 1)) 32 | im_ = torch.Tensor(im_[np.newaxis, ...]) 33 | 34 | with torch.no_grad(): 35 | if self.config.cuda: 36 | im_ = im_.cuda() 37 | preds = self.net(im_, mode=test_mode) 38 | preds_i = [] 39 | for p in preds[1]: 40 | preds_i.append(np.squeeze(torch.sigmoid(p).cpu().data.numpy())) 41 | pred_fuse = np.squeeze(torch.sigmoid(preds[0]).cpu().data.numpy()) 42 | pred = (pred_fuse + sum(preds_i)) / (1.0 + len(preds_i)) 43 | 44 | pred = (pred - np.min(pred) + EPSILON) / (np.max(pred) - np.min(pred) + EPSILON) 45 | 46 | pred = cv2.resize(pred, (im_size[1], im_size[0]), interpolation=cv2.INTER_LINEAR) 47 | multi_fuse += pred 48 | 49 | multi_fuse /= len(scale) 50 | multi_fuse = 255 * (1 - multi_fuse) 51 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[test_mode] + '.png'), multi_fuse) 52 | 53 | elif test_mode == 1: # saliency task 54 | with torch.no_grad(): 55 | if self.config.cuda: 56 | images = images.cuda() 57 | preds = self.net(images, mode=test_mode) 58 | pred = np.squeeze(torch.sigmoid(preds[0]).cpu().data.numpy()) 59 | 60 | multi_fuse = 255 * pred 61 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[test_mode] + '.png'), multi_fuse) 62 | 63 | elif test_mode == 2: # skeleton task 64 | images = images.numpy()[0].transpose((1,2,0)) 65 | scale = [0.5, 1, 1.5] # multi-scale testing as commonly done 66 | multi_fuse = np.zeros(im_size, np.float32) 67 | for k in range(0, len(scale)): 68 | im_ = cv2.resize(images, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR) 69 | im_ = im_.transpose((2, 0, 1)) 70 | im_ = torch.Tensor(im_[np.newaxis, ...]) 71 | 72 | with torch.no_grad(): 73 | if self.config.cuda: 74 | im_ = im_.cuda() 75 | preds = self.net(im_, mode=test_mode) 76 | pred_fuse = np.squeeze(torch.sigmoid(preds[0]).cpu().data.numpy()) 77 | 78 | pred = pred_fuse 79 | pred = (pred - np.min(pred) + EPSILON) / (np.max(pred) - np.min(pred) + EPSILON) 80 | 81 | pred = cv2.resize(pred, (im_size[1], im_size[0]), interpolation=cv2.INTER_LINEAR) 82 | multi_fuse += pred 83 | 84 | multi_fuse /= len(scale) 85 | multi_fuse = 255 * (1 - multi_fuse) 86 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[test_mode] + '.png'), multi_fuse) 87 | elif test_mode == 3: # all tasks 88 | with torch.no_grad(): 89 | if self.config.cuda: 90 | images = images.cuda() 91 | preds = self.net(images, mode=test_mode) 92 | pred_edge = np.squeeze(torch.sigmoid(preds[0][0]).cpu().data.numpy()) 93 | pred_sal = np.squeeze(torch.sigmoid(preds[1][0]).cpu().data.numpy()) 94 | pred_skel = np.squeeze(torch.sigmoid(preds[2][0]).cpu().data.numpy()) 95 | 96 | pred_edge = (pred_edge - np.min(pred_edge) + EPSILON) / (np.max(pred_edge) - np.min(pred_edge) + EPSILON) 97 | pred_skel = (pred_skel - np.min(pred_skel) + EPSILON) / (np.max(pred_skel) - np.min(pred_skel) + EPSILON) 98 | 99 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[0] + '.png'), 255 * (1 - pred_edge)) 100 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[1] + '.png'), 255 * pred_sal) 101 | cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '_' + mode_name[2] + '.png'), 255 * (1 - pred_skel)) 102 | 103 | print('Testing Finished.') 104 | 105 | --------------------------------------------------------------------------------