├── .DS_Store
├── .idea
├── .gitignore
├── TBSN.iml
├── deployment.xml
├── inspectionProfiles
│ ├── Project_Default.xml
│ └── profiles_settings.xml
├── misc.xml
├── modules.xml
├── vcs.xml
└── webServers.xml
├── LICENSE
├── benchmark
├── dnd.py
├── ensemble_wrapper.py
└── sidd.py
├── dataset
├── __init__.py
├── base_function.py
├── dnd.py
└── sidd.py
├── environment.yaml
├── exercise.py
├── model
├── __init__.py
├── apbsn.py
├── apbsn_distill.py
└── base.py
├── network
├── __init__.py
├── base_function.py
├── calculate_param_flops.py
├── tbsn.py
└── unet.py
├── option
├── knowledge_distillation_sidd.json
├── tbsn_dnd.json
└── tbsn_sidd.json
├── pretrained_models
├── model_dnd.pth
└── model_sidd.pth
├── readme.md
├── slurm.sh
├── train.sh
├── train
├── base.py
└── experiment_name.py
├── util
├── build.py
├── io.py
└── option.py
└── validate
├── base.py
├── base_function.py
├── visualization.py
└── visualize_receptive_field.py
/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nagejacob/TBSN/678e93dd4ab816d5c736179b01f7398373db45ef/.DS_Store
--------------------------------------------------------------------------------
/.idea/.gitignore:
--------------------------------------------------------------------------------
1 | # Default ignored files
2 | /shelf/
3 | /workspace.xml
4 | # Editor-based HTTP Client requests
5 | /httpRequests/
6 | # Datasource local storage ignored files
7 | /dataSources/
8 | /dataSources.local.xml
9 |
--------------------------------------------------------------------------------
/.idea/TBSN.iml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/deployment.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
--------------------------------------------------------------------------------
/.idea/inspectionProfiles/Project_Default.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
26 |
27 |
28 |
--------------------------------------------------------------------------------
/.idea/inspectionProfiles/profiles_settings.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/webServers.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
16 |
17 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | GNU GENERAL PUBLIC LICENSE
2 | Version 3, 29 June 2007
3 |
4 | Copyright (C) 2007 Free Software Foundation, Inc.
5 | Everyone is permitted to copy and distribute verbatim copies
6 | of this license document, but changing it is not allowed.
7 |
8 | Preamble
9 |
10 | The GNU General Public License is a free, copyleft license for
11 | software and other kinds of works.
12 |
13 | The licenses for most software and other practical works are designed
14 | to take away your freedom to share and change the works. By contrast,
15 | the GNU General Public License is intended to guarantee your freedom to
16 | share and change all versions of a program--to make sure it remains free
17 | software for all its users. We, the Free Software Foundation, use the
18 | GNU General Public License for most of our software; it applies also to
19 | any other work released this way by its authors. You can apply it to
20 | your programs, too.
21 |
22 | When we speak of free software, we are referring to freedom, not
23 | price. Our General Public Licenses are designed to make sure that you
24 | have the freedom to distribute copies of free software (and charge for
25 | them if you wish), that you receive source code or can get it if you
26 | want it, that you can change the software or use pieces of it in new
27 | free programs, and that you know you can do these things.
28 |
29 | To protect your rights, we need to prevent others from denying you
30 | these rights or asking you to surrender the rights. Therefore, you have
31 | certain responsibilities if you distribute copies of the software, or if
32 | you modify it: responsibilities to respect the freedom of others.
33 |
34 | For example, if you distribute copies of such a program, whether
35 | gratis or for a fee, you must pass on to the recipients the same
36 | freedoms that you received. You must make sure that they, too, receive
37 | or can get the source code. And you must show them these terms so they
38 | know their rights.
39 |
40 | Developers that use the GNU GPL protect your rights with two steps:
41 | (1) assert copyright on the software, and (2) offer you this License
42 | giving you legal permission to copy, distribute and/or modify it.
43 |
44 | For the developers' and authors' protection, the GPL clearly explains
45 | that there is no warranty for this free software. For both users' and
46 | authors' sake, the GPL requires that modified versions be marked as
47 | changed, so that their problems will not be attributed erroneously to
48 | authors of previous versions.
49 |
50 | Some devices are designed to deny users access to install or run
51 | modified versions of the software inside them, although the manufacturer
52 | can do so. This is fundamentally incompatible with the aim of
53 | protecting users' freedom to change the software. The systematic
54 | pattern of such abuse occurs in the area of products for individuals to
55 | use, which is precisely where it is most unacceptable. Therefore, we
56 | have designed this version of the GPL to prohibit the practice for those
57 | products. If such problems arise substantially in other domains, we
58 | stand ready to extend this provision to those domains in future versions
59 | of the GPL, as needed to protect the freedom of users.
60 |
61 | Finally, every program is threatened constantly by software patents.
62 | States should not allow patents to restrict development and use of
63 | software on general-purpose computers, but in those that do, we wish to
64 | avoid the special danger that patents applied to a free program could
65 | make it effectively proprietary. To prevent this, the GPL assures that
66 | patents cannot be used to render the program non-free.
67 |
68 | The precise terms and conditions for copying, distribution and
69 | modification follow.
70 |
71 | TERMS AND CONDITIONS
72 |
73 | 0. Definitions.
74 |
75 | "This License" refers to version 3 of the GNU General Public License.
76 |
77 | "Copyright" also means copyright-like laws that apply to other kinds of
78 | works, such as semiconductor masks.
79 |
80 | "The Program" refers to any copyrightable work licensed under this
81 | License. Each licensee is addressed as "you". "Licensees" and
82 | "recipients" may be individuals or organizations.
83 |
84 | To "modify" a work means to copy from or adapt all or part of the work
85 | in a fashion requiring copyright permission, other than the making of an
86 | exact copy. The resulting work is called a "modified version" of the
87 | earlier work or a work "based on" the earlier work.
88 |
89 | A "covered work" means either the unmodified Program or a work based
90 | on the Program.
91 |
92 | To "propagate" a work means to do anything with it that, without
93 | permission, would make you directly or secondarily liable for
94 | infringement under applicable copyright law, except executing it on a
95 | computer or modifying a private copy. Propagation includes copying,
96 | distribution (with or without modification), making available to the
97 | public, and in some countries other activities as well.
98 |
99 | To "convey" a work means any kind of propagation that enables other
100 | parties to make or receive copies. Mere interaction with a user through
101 | a computer network, with no transfer of a copy, is not conveying.
102 |
103 | An interactive user interface displays "Appropriate Legal Notices"
104 | to the extent that it includes a convenient and prominently visible
105 | feature that (1) displays an appropriate copyright notice, and (2)
106 | tells the user that there is no warranty for the work (except to the
107 | extent that warranties are provided), that licensees may convey the
108 | work under this License, and how to view a copy of this License. If
109 | the interface presents a list of user commands or options, such as a
110 | menu, a prominent item in the list meets this criterion.
111 |
112 | 1. Source Code.
113 |
114 | The "source code" for a work means the preferred form of the work
115 | for making modifications to it. "Object code" means any non-source
116 | form of a work.
117 |
118 | A "Standard Interface" means an interface that either is an official
119 | standard defined by a recognized standards body, or, in the case of
120 | interfaces specified for a particular programming language, one that
121 | is widely used among developers working in that language.
122 |
123 | The "System Libraries" of an executable work include anything, other
124 | than the work as a whole, that (a) is included in the normal form of
125 | packaging a Major Component, but which is not part of that Major
126 | Component, and (b) serves only to enable use of the work with that
127 | Major Component, or to implement a Standard Interface for which an
128 | implementation is available to the public in source code form. A
129 | "Major Component", in this context, means a major essential component
130 | (kernel, window system, and so on) of the specific operating system
131 | (if any) on which the executable work runs, or a compiler used to
132 | produce the work, or an object code interpreter used to run it.
133 |
134 | The "Corresponding Source" for a work in object code form means all
135 | the source code needed to generate, install, and (for an executable
136 | work) run the object code and to modify the work, including scripts to
137 | control those activities. However, it does not include the work's
138 | System Libraries, or general-purpose tools or generally available free
139 | programs which are used unmodified in performing those activities but
140 | which are not part of the work. For example, Corresponding Source
141 | includes interface definition files associated with source files for
142 | the work, and the source code for shared libraries and dynamically
143 | linked subprograms that the work is specifically designed to require,
144 | such as by intimate data communication or control flow between those
145 | subprograms and other parts of the work.
146 |
147 | The Corresponding Source need not include anything that users
148 | can regenerate automatically from other parts of the Corresponding
149 | Source.
150 |
151 | The Corresponding Source for a work in source code form is that
152 | same work.
153 |
154 | 2. Basic Permissions.
155 |
156 | All rights granted under this License are granted for the term of
157 | copyright on the Program, and are irrevocable provided the stated
158 | conditions are met. This License explicitly affirms your unlimited
159 | permission to run the unmodified Program. The output from running a
160 | covered work is covered by this License only if the output, given its
161 | content, constitutes a covered work. This License acknowledges your
162 | rights of fair use or other equivalent, as provided by copyright law.
163 |
164 | You may make, run and propagate covered works that you do not
165 | convey, without conditions so long as your license otherwise remains
166 | in force. You may convey covered works to others for the sole purpose
167 | of having them make modifications exclusively for you, or provide you
168 | with facilities for running those works, provided that you comply with
169 | the terms of this License in conveying all material for which you do
170 | not control copyright. Those thus making or running the covered works
171 | for you must do so exclusively on your behalf, under your direction
172 | and control, on terms that prohibit them from making any copies of
173 | your copyrighted material outside their relationship with you.
174 |
175 | Conveying under any other circumstances is permitted solely under
176 | the conditions stated below. Sublicensing is not allowed; section 10
177 | makes it unnecessary.
178 |
179 | 3. Protecting Users' Legal Rights From Anti-Circumvention Law.
180 |
181 | No covered work shall be deemed part of an effective technological
182 | measure under any applicable law fulfilling obligations under article
183 | 11 of the WIPO copyright treaty adopted on 20 December 1996, or
184 | similar laws prohibiting or restricting circumvention of such
185 | measures.
186 |
187 | When you convey a covered work, you waive any legal power to forbid
188 | circumvention of technological measures to the extent such circumvention
189 | is effected by exercising rights under this License with respect to
190 | the covered work, and you disclaim any intention to limit operation or
191 | modification of the work as a means of enforcing, against the work's
192 | users, your or third parties' legal rights to forbid circumvention of
193 | technological measures.
194 |
195 | 4. Conveying Verbatim Copies.
196 |
197 | You may convey verbatim copies of the Program's source code as you
198 | receive it, in any medium, provided that you conspicuously and
199 | appropriately publish on each copy an appropriate copyright notice;
200 | keep intact all notices stating that this License and any
201 | non-permissive terms added in accord with section 7 apply to the code;
202 | keep intact all notices of the absence of any warranty; and give all
203 | recipients a copy of this License along with the Program.
204 |
205 | You may charge any price or no price for each copy that you convey,
206 | and you may offer support or warranty protection for a fee.
207 |
208 | 5. Conveying Modified Source Versions.
209 |
210 | You may convey a work based on the Program, or the modifications to
211 | produce it from the Program, in the form of source code under the
212 | terms of section 4, provided that you also meet all of these conditions:
213 |
214 | a) The work must carry prominent notices stating that you modified
215 | it, and giving a relevant date.
216 |
217 | b) The work must carry prominent notices stating that it is
218 | released under this License and any conditions added under section
219 | 7. This requirement modifies the requirement in section 4 to
220 | "keep intact all notices".
221 |
222 | c) You must license the entire work, as a whole, under this
223 | License to anyone who comes into possession of a copy. This
224 | License will therefore apply, along with any applicable section 7
225 | additional terms, to the whole of the work, and all its parts,
226 | regardless of how they are packaged. This License gives no
227 | permission to license the work in any other way, but it does not
228 | invalidate such permission if you have separately received it.
229 |
230 | d) If the work has interactive user interfaces, each must display
231 | Appropriate Legal Notices; however, if the Program has interactive
232 | interfaces that do not display Appropriate Legal Notices, your
233 | work need not make them do so.
234 |
235 | A compilation of a covered work with other separate and independent
236 | works, which are not by their nature extensions of the covered work,
237 | and which are not combined with it such as to form a larger program,
238 | in or on a volume of a storage or distribution medium, is called an
239 | "aggregate" if the compilation and its resulting copyright are not
240 | used to limit the access or legal rights of the compilation's users
241 | beyond what the individual works permit. Inclusion of a covered work
242 | in an aggregate does not cause this License to apply to the other
243 | parts of the aggregate.
244 |
245 | 6. Conveying Non-Source Forms.
246 |
247 | You may convey a covered work in object code form under the terms
248 | of sections 4 and 5, provided that you also convey the
249 | machine-readable Corresponding Source under the terms of this License,
250 | in one of these ways:
251 |
252 | a) Convey the object code in, or embodied in, a physical product
253 | (including a physical distribution medium), accompanied by the
254 | Corresponding Source fixed on a durable physical medium
255 | customarily used for software interchange.
256 |
257 | b) Convey the object code in, or embodied in, a physical product
258 | (including a physical distribution medium), accompanied by a
259 | written offer, valid for at least three years and valid for as
260 | long as you offer spare parts or customer support for that product
261 | model, to give anyone who possesses the object code either (1) a
262 | copy of the Corresponding Source for all the software in the
263 | product that is covered by this License, on a durable physical
264 | medium customarily used for software interchange, for a price no
265 | more than your reasonable cost of physically performing this
266 | conveying of source, or (2) access to copy the
267 | Corresponding Source from a network server at no charge.
268 |
269 | c) Convey individual copies of the object code with a copy of the
270 | written offer to provide the Corresponding Source. This
271 | alternative is allowed only occasionally and noncommercially, and
272 | only if you received the object code with such an offer, in accord
273 | with subsection 6b.
274 |
275 | d) Convey the object code by offering access from a designated
276 | place (gratis or for a charge), and offer equivalent access to the
277 | Corresponding Source in the same way through the same place at no
278 | further charge. You need not require recipients to copy the
279 | Corresponding Source along with the object code. If the place to
280 | copy the object code is a network server, the Corresponding Source
281 | may be on a different server (operated by you or a third party)
282 | that supports equivalent copying facilities, provided you maintain
283 | clear directions next to the object code saying where to find the
284 | Corresponding Source. Regardless of what server hosts the
285 | Corresponding Source, you remain obligated to ensure that it is
286 | available for as long as needed to satisfy these requirements.
287 |
288 | e) Convey the object code using peer-to-peer transmission, provided
289 | you inform other peers where the object code and Corresponding
290 | Source of the work are being offered to the general public at no
291 | charge under subsection 6d.
292 |
293 | A separable portion of the object code, whose source code is excluded
294 | from the Corresponding Source as a System Library, need not be
295 | included in conveying the object code work.
296 |
297 | A "User Product" is either (1) a "consumer product", which means any
298 | tangible personal property which is normally used for personal, family,
299 | or household purposes, or (2) anything designed or sold for incorporation
300 | into a dwelling. In determining whether a product is a consumer product,
301 | doubtful cases shall be resolved in favor of coverage. For a particular
302 | product received by a particular user, "normally used" refers to a
303 | typical or common use of that class of product, regardless of the status
304 | of the particular user or of the way in which the particular user
305 | actually uses, or expects or is expected to use, the product. A product
306 | is a consumer product regardless of whether the product has substantial
307 | commercial, industrial or non-consumer uses, unless such uses represent
308 | the only significant mode of use of the product.
309 |
310 | "Installation Information" for a User Product means any methods,
311 | procedures, authorization keys, or other information required to install
312 | and execute modified versions of a covered work in that User Product from
313 | a modified version of its Corresponding Source. The information must
314 | suffice to ensure that the continued functioning of the modified object
315 | code is in no case prevented or interfered with solely because
316 | modification has been made.
317 |
318 | If you convey an object code work under this section in, or with, or
319 | specifically for use in, a User Product, and the conveying occurs as
320 | part of a transaction in which the right of possession and use of the
321 | User Product is transferred to the recipient in perpetuity or for a
322 | fixed term (regardless of how the transaction is characterized), the
323 | Corresponding Source conveyed under this section must be accompanied
324 | by the Installation Information. But this requirement does not apply
325 | if neither you nor any third party retains the ability to install
326 | modified object code on the User Product (for example, the work has
327 | been installed in ROM).
328 |
329 | The requirement to provide Installation Information does not include a
330 | requirement to continue to provide support service, warranty, or updates
331 | for a work that has been modified or installed by the recipient, or for
332 | the User Product in which it has been modified or installed. Access to a
333 | network may be denied when the modification itself materially and
334 | adversely affects the operation of the network or violates the rules and
335 | protocols for communication across the network.
336 |
337 | Corresponding Source conveyed, and Installation Information provided,
338 | in accord with this section must be in a format that is publicly
339 | documented (and with an implementation available to the public in
340 | source code form), and must require no special password or key for
341 | unpacking, reading or copying.
342 |
343 | 7. Additional Terms.
344 |
345 | "Additional permissions" are terms that supplement the terms of this
346 | License by making exceptions from one or more of its conditions.
347 | Additional permissions that are applicable to the entire Program shall
348 | be treated as though they were included in this License, to the extent
349 | that they are valid under applicable law. If additional permissions
350 | apply only to part of the Program, that part may be used separately
351 | under those permissions, but the entire Program remains governed by
352 | this License without regard to the additional permissions.
353 |
354 | When you convey a copy of a covered work, you may at your option
355 | remove any additional permissions from that copy, or from any part of
356 | it. (Additional permissions may be written to require their own
357 | removal in certain cases when you modify the work.) You may place
358 | additional permissions on material, added by you to a covered work,
359 | for which you have or can give appropriate copyright permission.
360 |
361 | Notwithstanding any other provision of this License, for material you
362 | add to a covered work, you may (if authorized by the copyright holders of
363 | that material) supplement the terms of this License with terms:
364 |
365 | a) Disclaiming warranty or limiting liability differently from the
366 | terms of sections 15 and 16 of this License; or
367 |
368 | b) Requiring preservation of specified reasonable legal notices or
369 | author attributions in that material or in the Appropriate Legal
370 | Notices displayed by works containing it; or
371 |
372 | c) Prohibiting misrepresentation of the origin of that material, or
373 | requiring that modified versions of such material be marked in
374 | reasonable ways as different from the original version; or
375 |
376 | d) Limiting the use for publicity purposes of names of licensors or
377 | authors of the material; or
378 |
379 | e) Declining to grant rights under trademark law for use of some
380 | trade names, trademarks, or service marks; or
381 |
382 | f) Requiring indemnification of licensors and authors of that
383 | material by anyone who conveys the material (or modified versions of
384 | it) with contractual assumptions of liability to the recipient, for
385 | any liability that these contractual assumptions directly impose on
386 | those licensors and authors.
387 |
388 | All other non-permissive additional terms are considered "further
389 | restrictions" within the meaning of section 10. If the Program as you
390 | received it, or any part of it, contains a notice stating that it is
391 | governed by this License along with a term that is a further
392 | restriction, you may remove that term. If a license document contains
393 | a further restriction but permits relicensing or conveying under this
394 | License, you may add to a covered work material governed by the terms
395 | of that license document, provided that the further restriction does
396 | not survive such relicensing or conveying.
397 |
398 | If you add terms to a covered work in accord with this section, you
399 | must place, in the relevant source files, a statement of the
400 | additional terms that apply to those files, or a notice indicating
401 | where to find the applicable terms.
402 |
403 | Additional terms, permissive or non-permissive, may be stated in the
404 | form of a separately written license, or stated as exceptions;
405 | the above requirements apply either way.
406 |
407 | 8. Termination.
408 |
409 | You may not propagate or modify a covered work except as expressly
410 | provided under this License. Any attempt otherwise to propagate or
411 | modify it is void, and will automatically terminate your rights under
412 | this License (including any patent licenses granted under the third
413 | paragraph of section 11).
414 |
415 | However, if you cease all violation of this License, then your
416 | license from a particular copyright holder is reinstated (a)
417 | provisionally, unless and until the copyright holder explicitly and
418 | finally terminates your license, and (b) permanently, if the copyright
419 | holder fails to notify you of the violation by some reasonable means
420 | prior to 60 days after the cessation.
421 |
422 | Moreover, your license from a particular copyright holder is
423 | reinstated permanently if the copyright holder notifies you of the
424 | violation by some reasonable means, this is the first time you have
425 | received notice of violation of this License (for any work) from that
426 | copyright holder, and you cure the violation prior to 30 days after
427 | your receipt of the notice.
428 |
429 | Termination of your rights under this section does not terminate the
430 | licenses of parties who have received copies or rights from you under
431 | this License. If your rights have been terminated and not permanently
432 | reinstated, you do not qualify to receive new licenses for the same
433 | material under section 10.
434 |
435 | 9. Acceptance Not Required for Having Copies.
436 |
437 | You are not required to accept this License in order to receive or
438 | run a copy of the Program. Ancillary propagation of a covered work
439 | occurring solely as a consequence of using peer-to-peer transmission
440 | to receive a copy likewise does not require acceptance. However,
441 | nothing other than this License grants you permission to propagate or
442 | modify any covered work. These actions infringe copyright if you do
443 | not accept this License. Therefore, by modifying or propagating a
444 | covered work, you indicate your acceptance of this License to do so.
445 |
446 | 10. Automatic Licensing of Downstream Recipients.
447 |
448 | Each time you convey a covered work, the recipient automatically
449 | receives a license from the original licensors, to run, modify and
450 | propagate that work, subject to this License. You are not responsible
451 | for enforcing compliance by third parties with this License.
452 |
453 | An "entity transaction" is a transaction transferring control of an
454 | organization, or substantially all assets of one, or subdividing an
455 | organization, or merging organizations. If propagation of a covered
456 | work results from an entity transaction, each party to that
457 | transaction who receives a copy of the work also receives whatever
458 | licenses to the work the party's predecessor in interest had or could
459 | give under the previous paragraph, plus a right to possession of the
460 | Corresponding Source of the work from the predecessor in interest, if
461 | the predecessor has it or can get it with reasonable efforts.
462 |
463 | You may not impose any further restrictions on the exercise of the
464 | rights granted or affirmed under this License. For example, you may
465 | not impose a license fee, royalty, or other charge for exercise of
466 | rights granted under this License, and you may not initiate litigation
467 | (including a cross-claim or counterclaim in a lawsuit) alleging that
468 | any patent claim is infringed by making, using, selling, offering for
469 | sale, or importing the Program or any portion of it.
470 |
471 | 11. Patents.
472 |
473 | A "contributor" is a copyright holder who authorizes use under this
474 | License of the Program or a work on which the Program is based. The
475 | work thus licensed is called the contributor's "contributor version".
476 |
477 | A contributor's "essential patent claims" are all patent claims
478 | owned or controlled by the contributor, whether already acquired or
479 | hereafter acquired, that would be infringed by some manner, permitted
480 | by this License, of making, using, or selling its contributor version,
481 | but do not include claims that would be infringed only as a
482 | consequence of further modification of the contributor version. For
483 | purposes of this definition, "control" includes the right to grant
484 | patent sublicenses in a manner consistent with the requirements of
485 | this License.
486 |
487 | Each contributor grants you a non-exclusive, worldwide, royalty-free
488 | patent license under the contributor's essential patent claims, to
489 | make, use, sell, offer for sale, import and otherwise run, modify and
490 | propagate the contents of its contributor version.
491 |
492 | In the following three paragraphs, a "patent license" is any express
493 | agreement or commitment, however denominated, not to enforce a patent
494 | (such as an express permission to practice a patent or covenant not to
495 | sue for patent infringement). To "grant" such a patent license to a
496 | party means to make such an agreement or commitment not to enforce a
497 | patent against the party.
498 |
499 | If you convey a covered work, knowingly relying on a patent license,
500 | and the Corresponding Source of the work is not available for anyone
501 | to copy, free of charge and under the terms of this License, through a
502 | publicly available network server or other readily accessible means,
503 | then you must either (1) cause the Corresponding Source to be so
504 | available, or (2) arrange to deprive yourself of the benefit of the
505 | patent license for this particular work, or (3) arrange, in a manner
506 | consistent with the requirements of this License, to extend the patent
507 | license to downstream recipients. "Knowingly relying" means you have
508 | actual knowledge that, but for the patent license, your conveying the
509 | covered work in a country, or your recipient's use of the covered work
510 | in a country, would infringe one or more identifiable patents in that
511 | country that you have reason to believe are valid.
512 |
513 | If, pursuant to or in connection with a single transaction or
514 | arrangement, you convey, or propagate by procuring conveyance of, a
515 | covered work, and grant a patent license to some of the parties
516 | receiving the covered work authorizing them to use, propagate, modify
517 | or convey a specific copy of the covered work, then the patent license
518 | you grant is automatically extended to all recipients of the covered
519 | work and works based on it.
520 |
521 | A patent license is "discriminatory" if it does not include within
522 | the scope of its coverage, prohibits the exercise of, or is
523 | conditioned on the non-exercise of one or more of the rights that are
524 | specifically granted under this License. You may not convey a covered
525 | work if you are a party to an arrangement with a third party that is
526 | in the business of distributing software, under which you make payment
527 | to the third party based on the extent of your activity of conveying
528 | the work, and under which the third party grants, to any of the
529 | parties who would receive the covered work from you, a discriminatory
530 | patent license (a) in connection with copies of the covered work
531 | conveyed by you (or copies made from those copies), or (b) primarily
532 | for and in connection with specific products or compilations that
533 | contain the covered work, unless you entered into that arrangement,
534 | or that patent license was granted, prior to 28 March 2007.
535 |
536 | Nothing in this License shall be construed as excluding or limiting
537 | any implied license or other defenses to infringement that may
538 | otherwise be available to you under applicable patent law.
539 |
540 | 12. No Surrender of Others' Freedom.
541 |
542 | If conditions are imposed on you (whether by court order, agreement or
543 | otherwise) that contradict the conditions of this License, they do not
544 | excuse you from the conditions of this License. If you cannot convey a
545 | covered work so as to satisfy simultaneously your obligations under this
546 | License and any other pertinent obligations, then as a consequence you may
547 | not convey it at all. For example, if you agree to terms that obligate you
548 | to collect a royalty for further conveying from those to whom you convey
549 | the Program, the only way you could satisfy both those terms and this
550 | License would be to refrain entirely from conveying the Program.
551 |
552 | 13. Use with the GNU Affero General Public License.
553 |
554 | Notwithstanding any other provision of this License, you have
555 | permission to link or combine any covered work with a work licensed
556 | under version 3 of the GNU Affero General Public License into a single
557 | combined work, and to convey the resulting work. The terms of this
558 | License will continue to apply to the part which is the covered work,
559 | but the special requirements of the GNU Affero General Public License,
560 | section 13, concerning interaction through a network will apply to the
561 | combination as such.
562 |
563 | 14. Revised Versions of this License.
564 |
565 | The Free Software Foundation may publish revised and/or new versions of
566 | the GNU General Public License from time to time. Such new versions will
567 | be similar in spirit to the present version, but may differ in detail to
568 | address new problems or concerns.
569 |
570 | Each version is given a distinguishing version number. If the
571 | Program specifies that a certain numbered version of the GNU General
572 | Public License "or any later version" applies to it, you have the
573 | option of following the terms and conditions either of that numbered
574 | version or of any later version published by the Free Software
575 | Foundation. If the Program does not specify a version number of the
576 | GNU General Public License, you may choose any version ever published
577 | by the Free Software Foundation.
578 |
579 | If the Program specifies that a proxy can decide which future
580 | versions of the GNU General Public License can be used, that proxy's
581 | public statement of acceptance of a version permanently authorizes you
582 | to choose that version for the Program.
583 |
584 | Later license versions may give you additional or different
585 | permissions. However, no additional obligations are imposed on any
586 | author or copyright holder as a result of your choosing to follow a
587 | later version.
588 |
589 | 15. Disclaimer of Warranty.
590 |
591 | THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
592 | APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
593 | HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
594 | OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
595 | THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
596 | PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
597 | IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
598 | ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
599 |
600 | 16. Limitation of Liability.
601 |
602 | IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
603 | WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
604 | THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
605 | GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
606 | USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
607 | DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
608 | PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
609 | EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
610 | SUCH DAMAGES.
611 |
612 | 17. Interpretation of Sections 15 and 16.
613 |
614 | If the disclaimer of warranty and limitation of liability provided
615 | above cannot be given local legal effect according to their terms,
616 | reviewing courts shall apply local law that most closely approximates
617 | an absolute waiver of all civil liability in connection with the
618 | Program, unless a warranty or assumption of liability accompanies a
619 | copy of the Program in return for a fee.
620 |
621 | END OF TERMS AND CONDITIONS
622 |
623 | How to Apply These Terms to Your New Programs
624 |
625 | If you develop a new program, and you want it to be of the greatest
626 | possible use to the public, the best way to achieve this is to make it
627 | free software which everyone can redistribute and change under these terms.
628 |
629 | To do so, attach the following notices to the program. It is safest
630 | to attach them to the start of each source file to most effectively
631 | state the exclusion of warranty; and each file should have at least
632 | the "copyright" line and a pointer to where the full notice is found.
633 |
634 |
635 | Copyright (C)
636 |
637 | This program is free software: you can redistribute it and/or modify
638 | it under the terms of the GNU General Public License as published by
639 | the Free Software Foundation, either version 3 of the License, or
640 | (at your option) any later version.
641 |
642 | This program is distributed in the hope that it will be useful,
643 | but WITHOUT ANY WARRANTY; without even the implied warranty of
644 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
645 | GNU General Public License for more details.
646 |
647 | You should have received a copy of the GNU General Public License
648 | along with this program. If not, see .
649 |
650 | Also add information on how to contact you by electronic and paper mail.
651 |
652 | If the program does terminal interaction, make it output a short
653 | notice like this when it starts in an interactive mode:
654 |
655 | Copyright (C)
656 | This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
657 | This is free software, and you are welcome to redistribute it
658 | under certain conditions; type `show c' for details.
659 |
660 | The hypothetical commands `show w' and `show c' should show the appropriate
661 | parts of the General Public License. Of course, your program's commands
662 | might be different; for a GUI interface, you would use an "about box".
663 |
664 | You should also get your employer (if you work as a programmer) or school,
665 | if any, to sign a "copyright disclaimer" for the program, if necessary.
666 | For more information on this, and how to apply and follow the GNU GPL, see
667 | .
668 |
669 | The GNU General Public License does not permit incorporating your program
670 | into proprietary programs. If your program is a subroutine library, you
671 | may consider it more useful to permit linking proprietary applications with
672 | the library. If this is what you want to do, use the GNU Lesser General
673 | Public License instead of this License. But first, please read
674 | .
675 |
--------------------------------------------------------------------------------
/benchmark/dnd.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | import argparse
4 | from benchmark.ensemble_wrapper import EnsembleWrapper
5 | from dataset.dnd import DNDBenchmarkPNGANDataset
6 | import numpy as np
7 | import os
8 | import scipy.io as sio
9 | import shutil
10 | from skimage.metrics import peak_signal_noise_ratio
11 | import torch
12 | from torch.utils.data import DataLoader
13 | from tqdm import tqdm
14 | from util.option import parse, recursive_print
15 |
16 | def bundle_submissions_srgb(submission_folder):
17 | '''
18 | Bundles submission data for sRGB denoising
19 |
20 | submission_folder Folder where denoised images reside
21 |
22 | Output is written to /bundled/. Please submit
23 | the content of this folder.
24 | '''
25 | out_folder = os.path.join(submission_folder, "bundled/")
26 | try:
27 | os.mkdir(out_folder)
28 | except:
29 | pass
30 | israw = False
31 | eval_version = "1.0"
32 |
33 | for i in range(50):
34 | Idenoised = np.zeros((20,), dtype=object)
35 | for bb in range(20):
36 | filename = '%04d_%02d.mat' % (i + 1, bb + 1)
37 | s = sio.loadmat(os.path.join(submission_folder, filename))
38 | Idenoised_crop = s["Idenoised_crop"]
39 | Idenoised[bb] = Idenoised_crop
40 | filename = '%04d.mat' % (i + 1)
41 | sio.savemat(os.path.join(out_folder, filename),
42 | {"Idenoised": Idenoised,
43 | "israw": israw,
44 | "eval_version": eval_version},
45 | )
46 |
47 | def main(opt):
48 | test_set = DNDBenchmarkPNGANDataset()
49 | test_loader = DataLoader(test_set, batch_size=1)
50 |
51 | if os.path.exists(opt['mat_dir']):
52 | shutil.rmtree(opt['mat_dir'])
53 | os.makedirs(opt['mat_dir'])
54 |
55 | Model = getattr(__import__('model'), opt['model'])
56 | model = Model(opt)
57 | model.data_parallel()
58 | if 'resume_from' in opt:
59 | model.load_model(opt['resume_from'])
60 | if opt['ensemble']:
61 | model = EnsembleWrapper(model)
62 |
63 | psnrs, count = 0, 0
64 | for data in tqdm(test_loader):
65 | output = model.validation_step(data)
66 | output = torch.clamp(output, 0, 1)
67 |
68 | output = output.permute(0, 2, 3, 1).cpu().numpy()
69 | i = count // 20
70 | k = count % 20
71 | save_file = os.path.join(opt['mat_dir'], '%04d_%02d.mat' % (i + 1, k + 1))
72 | sio.savemat(save_file, {'Idenoised_crop': output})
73 | count += 1
74 |
75 | output = output[0]
76 | gt = data['H'].squeeze(0).permute(1, 2, 0).numpy()
77 | psnr = peak_signal_noise_ratio(output, gt, data_range=1)
78 | psnrs += psnr
79 |
80 | print('%s, psnr: %6.4f' % (test_set.__class__.__name__, psnrs / count))
81 |
82 | if __name__ == '__main__':
83 | parser = argparse.ArgumentParser(description="Train the denoiser")
84 | parser.add_argument("--config_file", type=str, default='../option/tbsn_sidd.json')
85 | argspar = parser.parse_args()
86 |
87 | opt = parse(argspar.config_file)
88 | opt['mat_dir'] = 'dnd_mat'
89 | opt['ensemble'] = True
90 | recursive_print(opt)
91 |
92 | main(opt)
93 | bundle_submissions_srgb(opt['mat_dir'])
--------------------------------------------------------------------------------
/benchmark/ensemble_wrapper.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | class EnsembleWrapper():
4 | def __init__(self, model):
5 | self.model = model
6 |
7 | def validation_step(self, data):
8 | input = data['L']
9 | outputs = []
10 | for i in range(8):
11 | aug_input = input.clone()
12 | if i >= 4:
13 | aug_input = torch.flip(aug_input, [2])
14 | if i % 4 > 1:
15 | aug_input = torch.flip(aug_input, [3])
16 | if (i % 4) % 2 == 1:
17 | aug_input = torch.rot90(aug_input, 1, [2, 3])
18 |
19 | aug_output = self.model.validation_step({'L': aug_input})
20 |
21 | if (i % 4) % 2 == 1:
22 | aug_output = torch.rot90(aug_output, 3, [2, 3])
23 | if i % 4 > 1:
24 | aug_output = torch.flip(aug_output, [3])
25 | if i >= 4:
26 | aug_output = torch.flip(aug_output, [2])
27 | outputs.append(aug_output)
28 | output = torch.stack(outputs, dim=0)
29 | output = torch.mean(output, dim=0, keepdim=False)
30 | return output
--------------------------------------------------------------------------------
/benchmark/sidd.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | import argparse
4 | from dataset.sidd import SIDDBenchmarkDataset
5 | import numpy as np
6 | import os
7 | import scipy.io as sio
8 | from benchmark.ensemble_wrapper import EnsembleWrapper
9 | import torch
10 | from torch.utils.data import DataLoader
11 | from tqdm import tqdm
12 | from util.option import parse, recursive_print
13 |
14 | def main(opt):
15 | test_set = SIDDBenchmarkDataset()
16 | test_loader = DataLoader(test_set, batch_size=1)
17 |
18 | if os.path.exists(opt['mat_path']):
19 | os.remove(opt['mat_path'])
20 |
21 | Model = getattr(__import__('model'), opt['model'])
22 | model = Model(opt)
23 | model.data_parallel()
24 | if opt['ensemble']:
25 | model = EnsembleWrapper(model)
26 |
27 | count = 0
28 | denoised_block = np.zeros_like(test_set.noisy_block)
29 | for data in tqdm(test_loader):
30 | output = model.validation_step(data)
31 | output = torch.floor(output * 255. + 0.5)
32 | output = torch.clamp(output, 0, 255)
33 | output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
34 |
35 | index_n = count // test_set.noisy_block.shape[1]
36 | index_k = count % test_set.noisy_block.shape[1]
37 | output = np.uint8(output)
38 | denoised_block[index_n, index_k] = output
39 | count += 1
40 |
41 | save_dict = {}
42 | save_dict['__header__'] = b'MATLAB 5.0 MAT-file, Platform: PCWIN64, Created on: Thu Jan 10 13:08:11 2019'
43 | save_dict['__version__'] = 1.0
44 | save_dict['__globals__'] = []
45 | save_dict['DenoisedBlocksSrgb'] = denoised_block
46 | sio.savemat(opt['mat_path'], save_dict)
47 |
48 |
49 | if __name__ == '__main__':
50 | parser = argparse.ArgumentParser(description="Train the denoiser")
51 | parser.add_argument("--config_file", type=str, default='../option/tbsn_sidd.json')
52 | argspar = parser.parse_args()
53 |
54 | opt = parse(argspar.config_file)
55 | opt['mat_path'] = 'SubmitSrgb.mat'
56 | opt['ensemble'] = True
57 | recursive_print(opt)
58 |
59 | main(opt)
--------------------------------------------------------------------------------
/dataset/__init__.py:
--------------------------------------------------------------------------------
1 | from dataset.dnd import DNDBenchmarkTrainDataset, DNDBenchmarkPNGANDataset
2 | from dataset.sidd import SIDDMediumTrainDataset, SIDDValidationDataset, SIDDBenchmarkDataset
--------------------------------------------------------------------------------
/dataset/base_function.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import random
3 | import socket
4 | import torch
5 |
6 | hostname = socket.gethostname()
7 | if 'ubuntu' == hostname: # ubuntu
8 | dataset_path = '/home/nagejacob/Documents/datasets'
9 | else: # hpc
10 | dataset_path = '/mnt/ssd0/anaconda3/lijunyi/datasets'
11 |
12 | # c, h, w numpy
13 | def _aug_np3(img, flip_h, flip_w, transpose):
14 | if flip_h:
15 | img = img[:, ::-1, :]
16 | if flip_w:
17 | img = img[:, :, ::-1]
18 | if transpose:
19 | img = np.transpose(img, (0, 2, 1))
20 | return img
21 |
22 | def _aug_torch3(img, flip_h, flip_w, transpose):
23 | if flip_h:
24 | img = torch.flip(img, dims=[1])
25 | if flip_w:
26 | img = torch.flip(img, dims=[2])
27 | if transpose:
28 | img = img.permute(0, 2, 1)
29 | return img
30 |
31 | def _aug_3(img, flip_h, flip_w, transpose):
32 | if type(img) is np.ndarray:
33 | img = _aug_np3(img, flip_h, flip_w, transpose)
34 | elif type(img) is torch.Tensor:
35 | img = _aug_torch3(img, flip_h, flip_w, transpose)
36 | else:
37 | raise TypeError('img is neither np.ndarray nor torch.Tensor !')
38 | return img
39 |
40 | def aug_3(img_L, img_H=None):
41 | flip_h = random.random() > 0.5
42 | flip_w = random.random() > 0.5
43 | transpose = random.random() > 0.5
44 |
45 | img_L = _aug_3(img_L, flip_h, flip_w, transpose)
46 | if img_H is not None:
47 | img_H = _aug_3(img_H, flip_h, flip_w, transpose)
48 | return img_L, img_H
49 | else:
50 | return img_L
51 |
52 | def crop_3(patch_size, img_L, img_H=None):
53 | C, H, W = img_L.shape
54 | position_H = random.randint(0, H - patch_size)
55 | position_W = random.randint(0, W - patch_size)
56 | patch_L = img_L[:, position_H:position_H+patch_size, position_W:position_W+patch_size]
57 | if img_H is not None:
58 | patch_H = img_H[:, position_H:position_H+patch_size, position_W:position_W+patch_size]
59 | return patch_L, patch_H
60 | else:
61 | return patch_L
--------------------------------------------------------------------------------
/dataset/dnd.py:
--------------------------------------------------------------------------------
1 | from dataset.base_function import dataset_path, crop_3
2 | import h5py
3 | import glob
4 | import numpy as np
5 | import os
6 | import scipy.io as sio
7 | from torch.utils.data import Dataset
8 |
9 | dnd_path = os.path.join(dataset_path, 'DND')
10 |
11 | class DNDBenchmarkTrainDataset(Dataset):
12 | def __init__(self, patch_size, pin_memory=True):
13 | super().__init__()
14 | self.patch_size = patch_size
15 | self.pin_memory = pin_memory
16 |
17 | self._img_paths = self._get_img_paths()
18 | if self.pin_memory:
19 | self._imgs = self._open_images()
20 |
21 | def __getitem__(self, index):
22 | index = index % len(self._img_paths)
23 |
24 | if self.pin_memory:
25 | img_L = self.imgs[index]['L']
26 | else:
27 | img_path = self._img_paths[index]
28 | img_L = self._open_image(img_path['L'])
29 |
30 | patch_L = crop_3(self.patch_size, img_L)
31 |
32 | return {'L': patch_L.copy()}
33 |
34 | def __len__(self):
35 | return len(self._img_paths) * 100
36 |
37 | def _get_img_paths(self):
38 | L_pattern = os.path.join(dnd_path, 'images_srgb/00*.mat')
39 | L_paths = sorted(glob.glob(L_pattern))
40 |
41 | img_paths = []
42 | for L_path in L_paths:
43 | img_paths.append({'L':L_path})
44 | return img_paths
45 |
46 | def _open_images(self):
47 | self.imgs = []
48 | for img_path in self._img_paths:
49 | img_L = self._open_image(img_path['L'])
50 | self.imgs.append({'L': img_L})
51 |
52 | def _open_image(self, path):
53 | img = h5py.File(path)
54 | img = np.float32(np.array(img['InoisySRGB']).T)
55 | img = np.transpose(img, (2, 0, 1))
56 | return img
57 |
58 | # use PNGAN results as pseudo GT
59 | class DNDBenchmarkPNGANDataset(Dataset):
60 | def __init__(self, length=None):
61 | super(DNDBenchmarkPNGANDataset, self).__init__()
62 | self.length = length
63 | self.imgs = []
64 | infos = h5py.File(os.path.join(dnd_path, 'info.mat'), 'r')
65 | info = infos['info']
66 | bb = info['boundingboxes']
67 | for i in range(50):
68 | filename = os.path.join(dnd_path, 'images_srgb', '%04d.mat' % (i + 1))
69 | img = h5py.File(filename, 'r')
70 | Inoisy = np.float32(np.array(img['InoisySRGB']).T)
71 | ref = bb[0][i]
72 | boxes = np.array(info[ref]).T
73 | for k in range(20):
74 | idx = [int(boxes[k, 0] - 1), int(boxes[k, 2]), int(boxes[k, 1] - 1), int(boxes[k, 3])]
75 | Inoisy_crop = Inoisy[idx[0]:idx[1], idx[2]:idx[3], :].copy()
76 | Inoisy_crop = np.transpose(Inoisy_crop, (2, 0, 1))
77 | Iclean_crop = sio.loadmat(os.path.join(dnd_path, 'PNGAN/dnd_srgb_mat', '%04d_%02d.mat' % (i + 1, k + 1)))['Idenoised_crop']
78 | Iclean_crop = np.transpose(Iclean_crop[0], (2, 0, 1))
79 |
80 | self.imgs.append({'L':Inoisy_crop, 'H': Iclean_crop})
81 |
82 | def __getitem__(self, index):
83 | return self.imgs[index]
84 |
85 | def __len__(self):
86 | if self.length is not None:
87 | return self.length
88 | return 1000
--------------------------------------------------------------------------------
/dataset/sidd.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | from dataset.base_function import dataset_path, crop_3
4 | import glob
5 | import imageio
6 | import numpy as np
7 | import os
8 | from PIL import Image
9 | import scipy.io as sio
10 | import torch
11 | from torch.utils.data import Dataset
12 | from tqdm import tqdm
13 |
14 | sidd_path = os.path.join(dataset_path, 'SIDD')
15 |
16 | def open_image_SIDD(path):
17 | img = Image.open(path)
18 | img_np = np.asarray(img)
19 | img.close()
20 |
21 | img_np = np.transpose(img_np, (2, 0, 1))
22 | return img_np
23 |
24 | class SIDDMediumTrainDataset(Dataset):
25 | def __init__(self, pin_memory, patch_size):
26 | super().__init__()
27 | self.pin_memory = pin_memory
28 | self.patch_size = patch_size
29 |
30 | self._img_paths = self._get_img_paths()
31 | if self.pin_memory:
32 | self._imgs = self._open_images()
33 |
34 | def __getitem__(self, index):
35 | index = index % len(self._img_paths)
36 |
37 | if self.pin_memory:
38 | img_L = self._imgs[index]['L']
39 | img_H = self._imgs[index]['H']
40 | else:
41 | img_path = self._img_paths[index]
42 | img_L = self._open_image(img_path['L'])
43 | img_H = self._open_image(img_path['H'])
44 |
45 | img_L, img_H = crop_3(self.patch_size, img_L, img_H)
46 | img_L, img_H = np.float32(img_L) / 255., np.float32(img_H) / 255.
47 | img_L, img_H = torch.from_numpy(img_L), torch.from_numpy(img_H)
48 |
49 | return {'L': img_L, 'H': img_H}
50 |
51 | def __len__(self):
52 | length = len(self._img_paths)
53 | if length <= 10000:
54 | return (10000 // length) * length
55 | else:
56 | return length
57 |
58 | def _get_img_paths(self):
59 | img_paths = []
60 | L_pattern = os.path.join(sidd_path, 'SIDD_Medium_Srgb/Data/*/*_NOISY_SRGB_*.PNG')
61 | L_paths = sorted(glob.glob(L_pattern))
62 | for L_path in L_paths:
63 | img_paths.append({'L': L_path, 'H': L_path.replace('NOISY', 'GT')})
64 | return img_paths
65 |
66 | def _open_images(self):
67 | imgs = []
68 | for img_path in self._img_paths:
69 | img_L = self._open_image(img_path['L'])
70 | img_H = self._open_image(img_path['H'])
71 | imgs.append({'L': img_L, 'H': img_H})
72 | return imgs
73 |
74 | def _open_image(self, path):
75 | return open_image_SIDD(path)
76 |
77 |
78 | class SIDDValidationDataset(Dataset):
79 | def __init__(self):
80 | super().__init__()
81 | self._open_images(sidd_path)
82 | self.n = self.noisy_block.shape[0]
83 | self.k = self.noisy_block.shape[1]
84 |
85 | def __getitem__(self, index):
86 | index_n = index // self.k
87 | index_k = index % self.k
88 |
89 | img_H = self.gt_block[index_n, index_k]
90 | img_H = np.float32(img_H) / 255.
91 | img_H = np.transpose(img_H, (2, 0, 1))
92 | img_H = img_H
93 |
94 | img_L = self.noisy_block[index_n, index_k]
95 | img_L = np.float32(img_L) / 255.
96 | img_L = np.transpose(img_L, (2, 0, 1))
97 | img_L = img_L
98 |
99 | return {'H':img_H, 'L':img_L}
100 |
101 | def __len__(self):
102 | return self.n * self.k
103 |
104 | def _open_images(self, path):
105 | mat = sio.loadmat(os.path.join(path, 'SIDD_Validation/ValidationNoisyBlocksSrgb.mat'))
106 | self.noisy_block = mat['ValidationNoisyBlocksSrgb']
107 | mat = sio.loadmat(os.path.join(path, 'SIDD_Validation/ValidationGtBlocksSrgb.mat'))
108 | self.gt_block = mat['ValidationGtBlocksSrgb']
109 |
110 |
111 | class SIDDBenchmarkDataset(Dataset):
112 | def __init__(self):
113 | super().__init__()
114 | self._open_images(sidd_path)
115 | self.n = self.noisy_block.shape[0]
116 | self.k = self.noisy_block.shape[1]
117 |
118 | def __getitem__(self, index):
119 | index_n = index // self.k
120 | index_k = index % self.k
121 |
122 | img_L = self.noisy_block[index_n, index_k]
123 | img_L = np.float32(img_L) / 255.
124 | img_L = np.transpose(img_L, (2, 0, 1))
125 | img_L = img_L
126 |
127 | return {'L': img_L}
128 |
129 | def __len__(self):
130 | return self.n * self.k
131 |
132 | def _open_images(self, path):
133 | mat = sio.loadmat(os.path.join(path, 'SIDD_Benchmark/BenchmarkNoisyBlocksSrgb.mat'))
134 | self.noisy_block = mat['BenchmarkNoisyBlocksSrgb']
135 |
136 |
137 | # extract sidd validation images
138 | if __name__ == '__main__':
139 | dest_dir = os.path.join(sidd_path, 'SIDD_Validation/ValidationGTImagesSrgb')
140 | os.makedirs(dest_dir, exist_ok=True)
141 |
142 | mat = sio.loadmat(os.path.join(sidd_path, 'SIDD_Validation/ValidationNoisyBlocksSrgb.mat'))
143 | noisy_block = mat['ValidationNoisyBlocksSrgb']
144 | mat = sio.loadmat(os.path.join(sidd_path, 'SIDD_Validation/ValidationGtBlocksSrgb.mat'))
145 | gt_block = mat['ValidationGtBlocksSrgb']
146 | n = noisy_block.shape[0]
147 | k = noisy_block.shape[1]
148 | for i in tqdm(range(n)):
149 | for j in range(k):
150 | gt_img = gt_block[i, j]
151 | imageio.imwrite(os.path.join(dest_dir, '%02d_%02d.png' % (i, j)), gt_img)
--------------------------------------------------------------------------------
/environment.yaml:
--------------------------------------------------------------------------------
1 | name: tbsn
2 | channels:
3 | - pytorch
4 | - nvidia
5 | - defaults
6 | dependencies:
7 | - python==3.8.18
8 | - pytorch==2.0.1
9 | - pytorch-cuda=11.8
10 | - torchvision==0.15.2
11 | - torchaudio==2.0.2
12 | - h5py
13 | - scikit-image
14 | - tqdm
15 | - pip:
16 | - einops
17 | - opencv-python
18 |
19 | prefix: tbsn
--------------------------------------------------------------------------------
/exercise.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 |
4 | if __name__ == '__main__':
5 | tbsn_psnrs = torch.zeros((1280))
6 | unet_psnrs = torch.zeros((1280))
7 |
8 | with open('validate/tbsn.txt', 'r') as f:
9 | for i in range(1280):
10 | psnr = float(f.readline())
11 | tbsn_psnrs[i] = psnr
12 |
13 | with open('validate/unet.txt', 'r') as f:
14 | for i in range(1280):
15 | psnr = float(f.readline())
16 | unet_psnrs[i] = psnr - 0.1
17 |
18 | print('mean: ', torch.mean(tbsn_psnrs), torch.mean(unet_psnrs))
19 | print('std: ', torch.std(tbsn_psnrs - unet_psnrs))
20 |
21 | a = torch.abs(tbsn_psnrs - unet_psnrs)
22 | count = 0
23 | for i in range(1280):
24 | if a[i] < 0.2:
25 | count += 1
26 | print(count)
27 |
28 | indices = [(8, 20), (8, 23), (9, 5), (9, 8), (11, 29), (12, 19), (35, 0-21), (36, 0-21)]
29 | indices = [276, 279, 293, 296, 381, 403, 1120, 1121, 1122, 1123,
30 | 1124, 1125, 1126, 1127, 1128, 1129, 1130, 1131, 1132, 1133,
31 | 1134, 1135, 1136, 1137, 1138, 1139, 1140, 1141, 1152, 1153,
32 | 1154, 1155, 1156, 1157, 1158, 1159, 1160, 1161, 1162, 1163,
33 | 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173]
--------------------------------------------------------------------------------
/model/__init__.py:
--------------------------------------------------------------------------------
1 | from model.apbsn import APBSNModel
2 | from model.apbsn_distill import APBSNDistillModel
--------------------------------------------------------------------------------
/model/apbsn.py:
--------------------------------------------------------------------------------
1 | from model.base import BaseModel
2 | import torch
3 | import torch.nn as nn
4 | import torch.nn.functional as F
5 |
6 | def pixel_shuffle_down_sampling(x:torch.Tensor, f:int, train=False):
7 | b,c,w,h = x.shape
8 | unshuffled = F.pixel_unshuffle(x, f)
9 | unshuffled = unshuffled.view(b, c, f * f, w // f, h // f).permute(0, 2, 1, 3, 4)
10 | if train:
11 | unshuffled = unshuffled[:, 0]
12 | unshuffled = unshuffled.reshape(b, c, w // f, h // f)
13 | else:
14 | unshuffled = unshuffled.reshape(b * f * f, c, w // f, h // f)
15 | return unshuffled
16 |
17 | def pixel_shuffle_up_sampling(x:torch.Tensor, f:int):
18 | b,c,h,w = x.shape
19 | before_shuffle = x.view(b // (f * f), f * f, c, h, w).permute(0, 2, 1, 3, 4)
20 | before_shuffle = before_shuffle.reshape(b // (f * f), c * f * f, h, w)
21 | return F.pixel_shuffle(before_shuffle, f)
22 |
23 | class APBSNModel(BaseModel):
24 | def __init__(self, opt):
25 | super(APBSNModel, self).__init__(opt)
26 | self.pd_a = opt['pd_a']
27 | self.pd_b = opt['pd_b']
28 | self.R3 = opt['R3']
29 | self.R3_T = opt['R3_T']
30 | self.R3_p = opt['R3_p']
31 | self.criteron = nn.L1Loss(reduction='mean')
32 | self.optimizer = torch.optim.AdamW(self.networks['bsn'].parameters(), lr=opt['lr'])
33 |
34 | milestones = [int(opt['num_iters'] * 0.4), int(opt['num_iters'] * 0.8)]
35 | self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=milestones, gamma=0.1)
36 |
37 | def train_step(self, data):
38 | input = data['L'].to(self.device) * 255.
39 |
40 | self.networks['bsn'].train()
41 | input = pixel_shuffle_down_sampling(input, f=self.pd_a, train=True)
42 | output = self.networks['bsn'](input)
43 |
44 | self.loss = self.criteron(output, input)
45 | self.optimizer.zero_grad()
46 | self.loss.backward()
47 | self.optimizer.step()
48 | self.scheduler.step()
49 | self.iter += 1
50 |
51 | def validation_step(self, data):
52 | input = data['L'].to(self.device) * 255.
53 | b, c, h, w = input.shape
54 | input_pd = pixel_shuffle_down_sampling(input, f=self.pd_b)
55 |
56 | self.networks['bsn'].eval()
57 | with torch.no_grad():
58 | output_pd = self.networks['bsn'](input_pd)
59 |
60 | output = pixel_shuffle_up_sampling(output_pd, f=self.pd_b)
61 | if not self.R3:
62 | ''' Directly return the result (w/o R3) '''
63 | denoised = output[:, :, :h, :w]
64 | else:
65 | denoised = torch.empty(*(input.shape), self.R3_T, device=input.device)
66 | # torch.manual_seed(0)
67 | for t in range(self.R3_T):
68 | indice = torch.rand_like(input)
69 | mask = indice < self.R3_p
70 |
71 | tmp_input = torch.clone(output).detach()
72 | tmp_input[mask] = input[mask]
73 | with torch.no_grad():
74 | tmp_output = self.networks['bsn'](tmp_input)
75 | denoised[..., t] = tmp_output
76 |
77 | denoised = torch.mean(denoised, dim=-1)
78 |
79 | return denoised / 255.
80 |
--------------------------------------------------------------------------------
/model/apbsn_distill.py:
--------------------------------------------------------------------------------
1 | from model.base import BaseModel
2 | from model.apbsn import pixel_shuffle_down_sampling, pixel_shuffle_up_sampling
3 | import torch
4 | import torch.nn as nn
5 |
6 |
7 | class APBSNDistillModel(BaseModel):
8 | def __init__(self, opt):
9 | super(APBSNDistillModel, self).__init__(opt)
10 | self.pd_a = opt['pd_a']
11 | self.pd_b = opt['pd_b']
12 | self.R3 = opt['R3']
13 | self.R3_T = opt['R3_T']
14 | self.R3_p = opt['R3_p']
15 | self.criteron = nn.L1Loss(reduction='mean')
16 | self.optimizer = torch.optim.AdamW(self.networks['network'].parameters(), lr=opt['lr'])
17 |
18 | milestones = [int(opt['num_iters'] * 0.4), int(opt['num_iters'] * 0.8)]
19 | self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=milestones, gamma=0.1)
20 |
21 |
22 | def train_step(self, data):
23 |
24 | input = data['L'].to(self.device) * 255.
25 | b, c, h, w = input.shape
26 | input_pd = pixel_shuffle_down_sampling(input, f=self.pd_b)
27 |
28 | self.networks['bsn'].eval()
29 | with torch.no_grad():
30 | output_pd = self.networks['bsn'](input_pd)
31 |
32 | output = pixel_shuffle_up_sampling(output_pd, f=self.pd_b)
33 | if not self.R3:
34 | ''' Directly return the result (w/o R3) '''
35 | denoised = output[:, :, :h, :w]
36 | else:
37 | denoised = torch.empty(*(input.shape), self.R3_T, device=input.device)
38 | # torch.manual_seed(0)
39 | for t in range(self.R3_T):
40 | indice = torch.rand_like(input)
41 | mask = indice < self.R3_p
42 |
43 | tmp_input = torch.clone(output).detach()
44 | tmp_input[mask] = input[mask]
45 | with torch.no_grad():
46 | tmp_output = self.networks['bsn'](tmp_input)
47 | denoised[..., t] = tmp_output
48 |
49 | denoised = torch.mean(denoised, dim=-1)
50 |
51 | target = denoised / 255.
52 |
53 | input = data['L'].to(self.device)
54 | self.networks['network'].train()
55 | output = self.networks['network'](input)
56 |
57 | self.loss = self.criteron(output, target)
58 | self.optimizer.zero_grad()
59 | self.loss.backward()
60 | self.optimizer.step()
61 | self.scheduler.step()
62 | self.iter += 1
63 |
64 |
65 | def validation_step(self, data):
66 | input = data['L'].to(self.device)
67 |
68 | self.networks['network'].eval()
69 | with torch.no_grad():
70 | output = self.networks['network'](input)
71 |
72 | return output
--------------------------------------------------------------------------------
/model/base.py:
--------------------------------------------------------------------------------
1 | from abc import abstractmethod
2 | import os
3 | import torch
4 | from torch.nn.parallel import DataParallel, DistributedDataParallel
5 | from util.build import build
6 | from util.io import log
7 |
8 | class BaseModel():
9 | def __init__(self, opt):
10 | self.opt = opt
11 | self.iter = 0 if 'iter' not in opt else opt['iter']
12 | self.networks = {}
13 | for network_opt in opt['networks']:
14 | Net = getattr(__import__('network'), network_opt['type'])
15 | net = build(Net, network_opt['args'])
16 | if 'path' in network_opt.keys():
17 | self.load_net(net, network_opt['path'])
18 | self.networks[network_opt['name']] = net
19 |
20 | @abstractmethod
21 | def train_step(self, data):
22 | pass
23 |
24 | @abstractmethod
25 | def validation_step(self, data):
26 | pass
27 |
28 | def data_parallel(self):
29 | self.device = torch.device('cuda')
30 | for name in self.networks.keys():
31 | net = self.networks[name]
32 | net = net.cuda()
33 | net = DataParallel(net)
34 | self.networks[name] = net
35 |
36 | def distributed_parallel(self, rank):
37 | self.device = torch.device('cuda:%d' % rank)
38 | torch.cuda.set_device(rank)
39 | for name in self.networks.keys():
40 | net = self.networks[name]
41 | net = net.to(torch.device('cuda', rank))
42 | net = DistributedDataParallel(net, device_ids=[rank], output_device=rank)
43 | self.networks[name] = net
44 |
45 | def save_net(self):
46 | for name, net in self.networks.items():
47 | if isinstance(net, (DataParallel, DistributedDataParallel)):
48 | net = net.module
49 | torch.save(net.state_dict(), os.path.join(self.opt['log_dir'], '%s_iter_%08d.pth' % (name, self.iter)))
50 |
51 | def load_net(self, net, path):
52 | state_dict = torch.load(path, map_location='cpu')
53 | if 'model_weight' in state_dict:
54 | state_dict = state_dict['model_weight']['denoiser']
55 | if 'bsn' in list(state_dict.keys())[0]:
56 | for key in list(state_dict.keys()):
57 | state_dict[key.replace('bsn.', '')] = state_dict.pop(key)
58 | net.load_state_dict(state_dict)
59 |
60 |
61 | def log(self):
62 | with torch.no_grad():
63 | log(self.opt['log_file'], 'iter: %d, loss: %f\n' % (self.iter, self.loss.item()))
64 |
--------------------------------------------------------------------------------
/network/__init__.py:
--------------------------------------------------------------------------------
1 | from network.tbsn import TBSN
2 | from network.unet import Unet
--------------------------------------------------------------------------------
/network/base_function.py:
--------------------------------------------------------------------------------
1 | from einops import rearrange
2 | import numbers
3 | import torch
4 | from torch import einsum
5 | import torch.nn as nn
6 |
7 |
8 | def to(x):
9 | return {'device': x.device, 'dtype': x.dtype}
10 |
11 | def pair(x):
12 | return (x, x) if not isinstance(x, tuple) else x
13 |
14 | def expand_dim(t, dim, k):
15 | t = t.unsqueeze(dim = dim)
16 | expand_shape = [-1] * len(t.shape)
17 | expand_shape[dim] = k
18 | return t.expand(*expand_shape)
19 |
20 | def rel_to_abs(x):
21 | b, l, m = x.shape
22 | r = (m + 1) // 2
23 |
24 | col_pad = torch.zeros((b, l, 1), **to(x))
25 | x = torch.cat((x, col_pad), dim = 2)
26 | flat_x = rearrange(x, 'b l c -> b (l c)')
27 | flat_pad = torch.zeros((b, m - l), **to(x))
28 | flat_x_padded = torch.cat((flat_x, flat_pad), dim = 1)
29 | final_x = flat_x_padded.reshape(b, l + 1, m)
30 | final_x = final_x[:, :l, -r:]
31 | return final_x
32 |
33 | def relative_logits_1d(q, rel_k):
34 | b, h, w, _ = q.shape
35 | r = (rel_k.shape[0] + 1) // 2
36 |
37 | logits = einsum('b x y d, r d -> b x y r', q, rel_k)
38 | logits = rearrange(logits, 'b x y r -> (b x) y r')
39 | logits = rel_to_abs(logits)
40 |
41 | logits = logits.reshape(b, h, w, r)
42 | logits = expand_dim(logits, dim = 2, k = r)
43 | return logits
44 |
45 | class RelPosEmb(nn.Module):
46 | def __init__(
47 | self,
48 | block_size,
49 | rel_size,
50 | dim_head
51 | ):
52 | super().__init__()
53 | height = width = rel_size
54 | scale = dim_head ** -0.5
55 |
56 | self.block_size = block_size
57 | self.rel_height = nn.Parameter(torch.randn(height * 2 - 1, dim_head) * scale)
58 | self.rel_width = nn.Parameter(torch.randn(width * 2 - 1, dim_head) * scale)
59 |
60 | def forward(self, q):
61 | block = self.block_size
62 |
63 | q = rearrange(q, 'b (x y) c -> b x y c', x = block)
64 | rel_logits_w = relative_logits_1d(q, self.rel_width)
65 | rel_logits_w = rearrange(rel_logits_w, 'b x i y j-> b (x y) (i j)')
66 |
67 | q = rearrange(q, 'b x y d -> b y x d')
68 | rel_logits_h = relative_logits_1d(q, self.rel_height)
69 | rel_logits_h = rearrange(rel_logits_h, 'b x i y j -> b (y x) (j i)')
70 | return rel_logits_w + rel_logits_h
71 |
72 |
73 | class FixedPosEmb(nn.Module):
74 | def __init__(self, window_size, overlap_window_size):
75 | super().__init__()
76 | self.window_size = window_size
77 | self.overlap_window_size = overlap_window_size
78 |
79 | attention_mask_table = torch.zeros((window_size + overlap_window_size - 1), (window_size + overlap_window_size - 1))
80 | attention_mask_table[0::2, :] = float('-inf')
81 | attention_mask_table[:, 0::2] = float('-inf')
82 | attention_mask_table = attention_mask_table.view((window_size + overlap_window_size - 1) * (window_size + overlap_window_size - 1))
83 |
84 | # get pair-wise relative position index for each token inside the window
85 | coords_h = torch.arange(self.window_size)
86 | coords_w = torch.arange(self.window_size)
87 | coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
88 | coords_flatten_1 = torch.flatten(coords, 1) # 2, Wh*Ww
89 | coords_h = torch.arange(self.overlap_window_size)
90 | coords_w = torch.arange(self.overlap_window_size)
91 | coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
92 | coords_flatten_2 = torch.flatten(coords, 1)
93 |
94 | relative_coords = coords_flatten_1[:, :, None] - coords_flatten_2[:, None, :] # 2, Wh*Ww, Wh*Ww
95 | relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
96 | relative_coords[:, :, 0] += self.overlap_window_size - 1 # shift to start from 0
97 | relative_coords[:, :, 1] += self.overlap_window_size - 1
98 | relative_coords[:, :, 0] *= self.window_size + self.overlap_window_size - 1
99 | relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
100 | self.attention_mask = nn.Parameter(attention_mask_table[relative_position_index.view(-1)].view(
101 | 1, self.window_size ** 2, self.overlap_window_size ** 2
102 | ), requires_grad=False)
103 | def forward(self):
104 | return self.attention_mask
105 |
106 |
107 | ##########################################################################
108 | ## Layer Norm
109 |
110 | def to_3d(x):
111 | return rearrange(x, 'b c h w -> b (h w) c')
112 |
113 | def to_4d(x,h,w):
114 | return rearrange(x, 'b (h w) c -> b c h w',h=h,w=w)
115 |
116 | class BiasFree_LayerNorm(nn.Module):
117 | def __init__(self, normalized_shape):
118 | super(BiasFree_LayerNorm, self).__init__()
119 | if isinstance(normalized_shape, numbers.Integral):
120 | normalized_shape = (normalized_shape,)
121 | normalized_shape = torch.Size(normalized_shape)
122 |
123 | assert len(normalized_shape) == 1
124 |
125 | self.weight = nn.Parameter(torch.ones(normalized_shape))
126 | self.normalized_shape = normalized_shape
127 |
128 | def forward(self, x):
129 | sigma = x.var(-1, keepdim=True, unbiased=False)
130 | return x / torch.sqrt(sigma+1e-5) * self.weight
131 |
132 | class WithBias_LayerNorm(nn.Module):
133 | def __init__(self, normalized_shape):
134 | super(WithBias_LayerNorm, self).__init__()
135 | if isinstance(normalized_shape, numbers.Integral):
136 | normalized_shape = (normalized_shape,)
137 | normalized_shape = torch.Size(normalized_shape)
138 |
139 | assert len(normalized_shape) == 1
140 |
141 | self.weight = nn.Parameter(torch.ones(normalized_shape))
142 | self.bias = nn.Parameter(torch.zeros(normalized_shape))
143 | self.normalized_shape = normalized_shape
144 |
145 | def forward(self, x):
146 | mu = x.mean(-1, keepdim=True)
147 | sigma = x.var(-1, keepdim=True, unbiased=False)
148 | return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias
149 |
150 | class LayerNorm(nn.Module):
151 | def __init__(self, dim, LayerNorm_type):
152 | super(LayerNorm, self).__init__()
153 | if LayerNorm_type =='BiasFree':
154 | self.body = BiasFree_LayerNorm(dim)
155 | else:
156 | self.body = WithBias_LayerNorm(dim)
157 |
158 | def forward(self, x):
159 | h, w = x.shape[-2:]
160 | return to_4d(self.body(to_3d(x)), h, w)
--------------------------------------------------------------------------------
/network/calculate_param_flops.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | from deepspeed.accelerator import get_accelerator
4 | from deepspeed.profiling.flops_profiler import get_model_profile
5 | import time
6 | import torch
7 |
8 |
9 | if __name__ == '__main__':
10 | with get_accelerator().device(0):
11 | Model = getattr(__import__('network'), 'TBSN')
12 | model = Model().cuda()
13 |
14 | # t = time.time()
15 | # input = torch.zeros(1, 3, 256, 256).cuda()
16 | # for i in range(1000):
17 | # with torch.no_grad():
18 | # output = model(input)
19 | # print(time.time() - t), exit()
20 |
21 |
22 | flops, macs, params = get_model_profile(model=model, # model
23 | input_shape=(1, 3, 256, 256),
24 | # input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument.
25 | args=None, # list of positional arguments to the model.
26 | kwargs=None, # dictionary of keyword arguments to the model.
27 | print_profile=False,
28 | # prints the model graph with the measured profile attached to each module
29 | detailed=False, # print the detailed profile
30 | module_depth=-1,
31 | # depth into the nested modules, with -1 being the inner most modules
32 | top_modules=1, # the number of top modules to print aggregated profile
33 | warm_up=0,
34 | # the number of warm-ups before measuring the time of each module
35 | as_string=False,
36 | # print raw numbers (e.g. 1000) or as human-readable strings (e.g. 1k)
37 | output_file=None,
38 | # path to the output file. If None, the profiler prints to stdout.
39 | ignore_modules=None) # the list of modules to ignore in the profiling
40 | params = params / (1000000)
41 | flops = flops / (1000000000)
42 | print('params(M): ', params, 'flops(G): ', flops)
43 |
--------------------------------------------------------------------------------
/network/tbsn.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | from einops import rearrange
5 | import math
6 | from network.base_function import RelPosEmb, FixedPosEmb, LayerNorm
7 |
8 |
9 | class PatchUnshuffle(nn.Module):
10 | def __init__(self, p=2, s=2):
11 | super().__init__()
12 | self.p = p
13 | self.s = s
14 |
15 | def forward(self, x):
16 | n, c, h, w = x.shape
17 | x = nn.functional.pixel_unshuffle(x, self.p)
18 | x = nn.functional.pixel_unshuffle(x, self.s)
19 | x = x.view(n, c, self.p * self.p, self.s * self.s, h//self.p//self.s, w//self.p//self.s).permute(0, 1, 3, 2, 4, 5)
20 | x = x.contiguous().view(n, c * (self.p**2) * (self.s**2), h//self.p//self.s, w//self.p//self.s)
21 | x = nn.functional.pixel_shuffle(x, self.p)
22 | return x
23 |
24 | class PatchShuffle(nn.Module):
25 | def __init__(self, p=2, s=2):
26 | super().__init__()
27 | self.p = p
28 | self.s = s
29 |
30 | def forward(self, x):
31 | n, c, h, w = x.shape
32 | x = nn.functional.pixel_unshuffle(x, self.p)
33 | x = x.view(n, c//(self.s**2), (self.s**2), (self.p**2), h//self.p, w//self.p).permute(0, 1, 3, 2, 4, 5)
34 | x = x.contiguous().view(n, c * (self.p**2), h//self.p, w//self.p)
35 | x = nn.functional.pixel_shuffle(x, self.s)
36 | x = nn.functional.pixel_shuffle(x, self.p)
37 | return x
38 |
39 |
40 | class CentralMaskedConv2d(nn.Conv2d):
41 | def __init__(self, *args, **kwargs):
42 | super().__init__(*args, **kwargs)
43 |
44 | self.register_buffer('mask', self.weight.data.clone())
45 | _, _, kH, kW = self.weight.size()
46 | self.mask.fill_(1)
47 | self.mask[:, :, kH // 2, kH // 2] = 0
48 |
49 | def forward(self, x):
50 | self.weight.data *= self.mask
51 | return super().forward(x)
52 |
53 |
54 | class DilatedMDTA(nn.Module):
55 | def __init__(self, dim, num_heads, bias):
56 | super(DilatedMDTA, self).__init__()
57 | self.num_heads = num_heads
58 | self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
59 |
60 | self.qkv = nn.Conv2d(dim, dim*3, kernel_size=1, bias=bias)
61 | self.qkv_dwconv = nn.Conv2d(dim*3, dim*3, kernel_size=3, stride=1, dilation=2, padding=2, groups=dim*3, bias=bias)
62 | self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
63 |
64 | def forward(self, x):
65 | b,c,h,w = x.shape
66 |
67 | qkv = self.qkv_dwconv(self.qkv(x))
68 | q,k,v = qkv.chunk(3, dim=1)
69 |
70 | q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
71 | k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
72 | v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
73 |
74 | q = torch.nn.functional.normalize(q, dim=-1)
75 | k = torch.nn.functional.normalize(k, dim=-1)
76 |
77 | attn = (q @ k.transpose(-2, -1)) * self.temperature
78 | attn = attn.softmax(dim=-1)
79 |
80 | out = (attn @ v)
81 |
82 | out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
83 |
84 | out = self.project_out(out)
85 | return out
86 |
87 |
88 | class DilatedOCA(nn.Module):
89 | def __init__(self, dim, window_size, overlap_ratio, num_heads, dim_head, bias):
90 | super(DilatedOCA, self).__init__()
91 | self.num_spatial_heads = num_heads
92 | self.dim = dim
93 | self.window_size = window_size
94 | self.overlap_win_size = int(window_size * overlap_ratio) + window_size
95 | self.dim_head = dim_head
96 | self.inner_dim = self.dim_head * self.num_spatial_heads
97 | self.scale = self.dim_head**-0.5
98 |
99 | self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size), stride=window_size, padding=(self.overlap_win_size-window_size)//2)
100 | self.qkv = nn.Conv2d(self.dim, self.inner_dim*3, kernel_size=1, bias=bias)
101 | self.project_out = nn.Conv2d(self.inner_dim, dim, kernel_size=1, bias=bias)
102 | self.rel_pos_emb = RelPosEmb(
103 | block_size = window_size,
104 | rel_size = window_size + (self.overlap_win_size - window_size),
105 | dim_head = self.dim_head
106 | )
107 | self.fixed_pos_emb = FixedPosEmb(window_size, self.overlap_win_size)
108 | def forward(self, x):
109 | b, c, h, w = x.shape
110 | qkv = self.qkv(x)
111 | qs, ks, vs = qkv.chunk(3, dim=1)
112 |
113 | # spatial attention
114 | qs = rearrange(qs, 'b c (h p1) (w p2) -> (b h w) (p1 p2) c', p1 = self.window_size, p2 = self.window_size)
115 | ks, vs = map(lambda t: self.unfold(t), (ks, vs))
116 | ks, vs = map(lambda t: rearrange(t, 'b (c j) i -> (b i) j c', c = self.inner_dim), (ks, vs))
117 |
118 | # print(f'qs.shape:{qs.shape}, ks.shape:{ks.shape}, vs.shape:{vs.shape}')
119 | #split heads
120 | qs, ks, vs = map(lambda t: rearrange(t, 'b n (head c) -> (b head) n c', head = self.num_spatial_heads), (qs, ks, vs))
121 |
122 | # attention
123 | qs = qs * self.scale
124 | spatial_attn = (qs @ ks.transpose(-2, -1))
125 | spatial_attn += self.rel_pos_emb(qs)
126 | spatial_attn += self.fixed_pos_emb()
127 | spatial_attn = spatial_attn.softmax(dim=-1)
128 |
129 | out = (spatial_attn @ vs)
130 |
131 | out = rearrange(out, '(b h w head) (p1 p2) c -> b (head c) (h p1) (w p2)', head = self.num_spatial_heads, h = h // self.window_size, w = w // self.window_size, p1 = self.window_size, p2 = self.window_size)
132 |
133 | # merge spatial and channel
134 | out = self.project_out(out)
135 |
136 | return out
137 |
138 |
139 | class FeedForward(nn.Module):
140 | def __init__(self, dim, ffn_expansion_factor, bias):
141 | super(FeedForward, self).__init__()
142 |
143 | hidden_features = int(dim * ffn_expansion_factor)
144 |
145 | self.project_in = nn.Conv2d(dim, hidden_features, kernel_size=3, stride=1, dilation=2, padding=2, bias=bias)
146 | self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=3, stride=1, dilation=2, padding=2, bias=bias)
147 |
148 | def forward(self, x):
149 | x = self.project_in(x)
150 | x = F.gelu(x)
151 | x = self.project_out(x)
152 | return x
153 |
154 |
155 | class TransformerBlock(nn.Module):
156 | def __init__(self, dim, window_size, overlap_ratio, num_channel_heads, num_spatial_heads, spatial_dim_head, ffn_expansion_factor, bias, LayerNorm_type):
157 | super(TransformerBlock, self).__init__()
158 |
159 |
160 | self.spatial_attn = DilatedOCA(dim, window_size, overlap_ratio, num_spatial_heads, spatial_dim_head, bias)
161 | self.channel_attn = DilatedMDTA(dim, num_channel_heads, bias)
162 |
163 | self.norm1 = LayerNorm(dim, LayerNorm_type)
164 | self.norm2 = LayerNorm(dim, LayerNorm_type)
165 | self.norm3 = LayerNorm(dim, LayerNorm_type)
166 | self.norm4 = LayerNorm(dim, LayerNorm_type)
167 |
168 | self.channel_ffn = FeedForward(dim, ffn_expansion_factor, bias)
169 | self.spatial_ffn = FeedForward(dim, ffn_expansion_factor, bias)
170 |
171 |
172 | def forward(self, x):
173 | x = x + self.channel_attn(self.norm1(x))
174 | x = x + self.channel_ffn(self.norm2(x))
175 | x = x + self.spatial_attn(self.norm3(x))
176 | x = x + self.spatial_ffn(self.norm4(x))
177 | return x
178 |
179 | ##########################################################################
180 | ## Overlapped image patch embedding with 3x3 Conv
181 | class OverlapPatchEmbed(nn.Module):
182 | def __init__(self, in_c=3, embed_dim=48, bias=False):
183 | super(OverlapPatchEmbed, self).__init__()
184 |
185 | self.proj = CentralMaskedConv2d(in_c, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias)
186 |
187 | def forward(self, x):
188 | x = self.proj(x)
189 |
190 | return x
191 |
192 | ##########################################################################
193 | ## Resizing modules
194 | class Downsample(nn.Module):
195 | def __init__(self, n_feat):
196 | super(Downsample, self).__init__()
197 |
198 | self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat//2, kernel_size=1, bias=False),
199 | PatchUnshuffle())
200 |
201 | def forward(self, x):
202 | return self.body(x)
203 |
204 |
205 | class Upsample(nn.Module):
206 | def __init__(self, n_feat):
207 | super(Upsample, self).__init__()
208 |
209 | self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat*2, kernel_size=1, bias=False),
210 | PatchShuffle())
211 |
212 | def forward(self, x):
213 | return self.body(x)
214 |
215 |
216 | class SR_Upsample(nn.Sequential):
217 | """SR_Upsample module.
218 | Args:
219 | scale (int): Scale factor. Supported scales: 2^n and 3.
220 | num_feat (int): Channel number of features.
221 | """
222 |
223 | def __init__(self, scale, num_feat):
224 | m = []
225 |
226 | if (scale & (scale - 1)) == 0: # scale = 2^n
227 | for _ in range(int(math.log(scale, 2))):
228 | m.append(nn.Conv2d(num_feat, 4 * num_feat, kernel_size = 3, stride = 1, padding = 1))
229 | m.append(nn.PixelShuffle(2))
230 | elif scale == 3:
231 | m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
232 | m.append(nn.PixelShuffle(3))
233 | else:
234 | raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
235 | super(SR_Upsample, self).__init__(*m)
236 |
237 | ##########################################################################
238 |
239 | class TBSN(nn.Module):
240 | def __init__(self,
241 | in_ch=3,
242 | out_ch=3,
243 | dim = 48,
244 | num_blocks = [4,4,4],
245 | num_refinement_blocks = 4,
246 | channel_heads = [1,2,4],
247 | spatial_heads = [2,2,3],
248 | overlap_ratio=[0.5, 0.5, 0.5],
249 | window_size = 8,
250 | spatial_dim_head = 16,
251 | bias = False,
252 | ffn_expansion_factor = 1,
253 | LayerNorm_type = 'BiasFree',
254 | scale = 1
255 | ):
256 |
257 | super(TBSN, self).__init__()
258 | self.scale = scale
259 |
260 | self.patch_embed = OverlapPatchEmbed(in_ch, dim)
261 | self.encoder_level1 = nn.Sequential(*[TransformerBlock(dim=dim, window_size = window_size, overlap_ratio=overlap_ratio[0], num_channel_heads=channel_heads[0], num_spatial_heads=spatial_heads[0], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
262 |
263 | self.down1_2 = Downsample(dim) ## From Level 1 to Level 2
264 | self.encoder_level2 = nn.Sequential(*[TransformerBlock(dim=int(dim*2**1), window_size = window_size, overlap_ratio=overlap_ratio[1], num_channel_heads=channel_heads[1], num_spatial_heads=spatial_heads[1], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[1])])
265 |
266 | self.down2_3 = Downsample(int(dim*2**1)) ## From Level 2 to Level 3
267 | self.latent = nn.Sequential(*[TransformerBlock(dim=int(dim*2**2), window_size = window_size, overlap_ratio=overlap_ratio[2], num_channel_heads=channel_heads[2], num_spatial_heads=spatial_heads[2], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[2])])
268 |
269 | self.up3_2 = Upsample(int(dim*2**2)) ## From Level 3 to Level 2
270 | self.reduce_chan_level2 = nn.Conv2d(int(dim*2**2), int(dim*2**1), kernel_size=1, bias=bias)
271 | self.decoder_level2 = nn.Sequential(*[TransformerBlock(dim=int(dim*2**1), window_size = window_size, overlap_ratio=overlap_ratio[1], num_channel_heads=channel_heads[1], num_spatial_heads=spatial_heads[1], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[1])])
272 |
273 | self.up2_1 = Upsample(int(dim*2**1)) ## From Level 2 to Level 1 (NO 1x1 conv to reduce channels)
274 |
275 | self.decoder_level1 = nn.Sequential(*[TransformerBlock(dim=int(dim*2**1), window_size = window_size, overlap_ratio=overlap_ratio[0], num_channel_heads=channel_heads[0], num_spatial_heads=spatial_heads[0], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
276 |
277 | self.refinement = nn.Sequential(*[TransformerBlock(dim=int(dim*2**1), window_size = window_size, overlap_ratio=overlap_ratio[0], num_channel_heads=channel_heads[0], num_spatial_heads=spatial_heads[0], spatial_dim_head = spatial_dim_head, ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_refinement_blocks)])
278 |
279 | self.output = nn.Conv2d(int(dim*2**1), out_ch, kernel_size=1, bias=bias)
280 |
281 | def forward(self, inp_img):
282 |
283 | if self.scale > 1:
284 | inp_img = F.interpolate(inp_img, scale_factor=self.scale, mode='bilinear', align_corners=False)
285 |
286 | inp_enc_level1 = self.patch_embed(inp_img)
287 | out_enc_level1 = self.encoder_level1(inp_enc_level1)
288 |
289 | inp_enc_level2 = self.down1_2(out_enc_level1)
290 | out_enc_level2 = self.encoder_level2(inp_enc_level2)
291 |
292 | inp_enc_level3 = self.down2_3(out_enc_level2)
293 | latent = self.latent(inp_enc_level3)
294 |
295 | inp_dec_level2 = self.up3_2(latent)
296 | inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 1)
297 | inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
298 | out_dec_level2 = self.decoder_level2(inp_dec_level2)
299 |
300 | inp_dec_level1 = self.up2_1(out_dec_level2)
301 | inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 1)
302 | out_dec_level1 = self.decoder_level1(inp_dec_level1)
303 |
304 | out_dec_level1 = self.refinement(out_dec_level1)
305 | out_dec_level1 = self.output(out_dec_level1)
306 |
307 | return out_dec_level1
308 |
309 |
--------------------------------------------------------------------------------
/network/unet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from typing import Tuple
4 |
5 | '''
6 | https://github.com/COMP6248-Reproducability-Challenge/selfsupervised-denoising/blob/master-with-report/ssdn/ssdn/models/noise_network.py
7 | '''
8 | class Unet(nn.Module):
9 | """Custom U-Net architecture for Self Supervised Denoising (SSDN) and Noise2Noise (N2N).
10 | Base N2N implementation was made with reference to @joeylitalien's N2N implementation.
11 | Changes made are removal of weight sharing when blocks are reused. Usage of LeakyReLu
12 | over standard ReLu and incorporation of blindspot functionality.
13 |
14 | Unlike other typical U-Net implementations dropout is not used when the model is trained.
15 |
16 | When in blindspot mode the following behaviour changes occur:
17 |
18 | * Input batches are duplicated for rotations: 0, 90, 180, 270. This increases the
19 | batch size by 4x. After the encode-decode stage the rotations are undone and
20 | concatenated on the channel axis with the associated original image. This 4x
21 | increase in channel count is collapsed to the standard channel count in the
22 | first 1x1 kernel convolution.
23 |
24 | * To restrict the receptive field into the upward direction a shift is used for
25 | convolutions (see nn.Conv2d) and downsampling. Downsampling uses a single
26 | pixel shift prior to max pooling as dictated by Laine et al. This is equivalent
27 | to applying a shift on the upsample.
28 |
29 | Args:
30 | in_channels (int, optional): Number of input channels, this will typically be either
31 | 1 (Mono) or 3 (RGB) but can be more. Defaults to 3.
32 | out_channels (int, optional): Number of channels the final convolution should output.
33 | Defaults to 3.
34 | blindspot (bool, optional): Whether to enable the network blindspot. This will
35 | add in rotation stages and shift stages while max pooling and during convolutions.
36 | A futher shift will occur after upsample. Defaults to False.
37 | zero_output_weights (bool, optional): Whether to initialise the weights of
38 | `nin_c` to zero. This is not mentioned in literature but is done as part
39 | of the tensorflow implementation for the parameter estimation network.
40 | Defaults to False.
41 | """
42 |
43 | def __init__(self, in_ch=3, out_ch=3, zero_output=False, dim=48):
44 | super(Unet, self).__init__()
45 | self.zero_output = zero_output
46 | in_channels = in_ch
47 | out_channels = out_ch
48 |
49 | ####################################
50 | # Encode Blocks
51 | ####################################
52 |
53 | # Layers: enc_conv0, enc_conv1, pool1
54 | self.encode_block_1 = nn.Sequential(
55 | nn.Conv2d(in_channels, dim, 3, stride=1, padding=1),
56 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
57 | nn.Conv2d(dim, dim, 3, padding=1),
58 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
59 | nn.MaxPool2d(2)
60 | )
61 |
62 | # Layers: enc_conv(i), pool(i); i=2..5
63 | def _encode_block_2_3_4_5() -> nn.Module:
64 | return nn.Sequential(
65 | nn.Conv2d(dim, dim, 3, stride=1, padding=1),
66 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
67 | nn.MaxPool2d(2)
68 | )
69 |
70 | # Separate instances of same encode module definition created
71 | self.encode_block_2 = _encode_block_2_3_4_5()
72 | self.encode_block_3 = _encode_block_2_3_4_5()
73 | self.encode_block_4 = _encode_block_2_3_4_5()
74 | self.encode_block_5 = _encode_block_2_3_4_5()
75 |
76 | # Layers: enc_conv6
77 | self.encode_block_6 = nn.Sequential(
78 | nn.Conv2d(dim, dim, 3, stride=1, padding=1),
79 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
80 | )
81 |
82 | ####################################
83 | # Decode Blocks
84 | ####################################
85 | # Layers: upsample5
86 | self.decode_block_6 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"))
87 |
88 | # Layers: dec_conv5a, dec_conv5b, upsample4
89 | self.decode_block_5 = nn.Sequential(
90 | nn.Conv2d(dim * 2, dim * 2, 3, stride=1, padding=1),
91 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
92 | nn.Conv2d(dim * 2, dim * 2, 3, stride=1, padding=1),
93 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
94 | nn.Upsample(scale_factor=2, mode="nearest"),
95 | )
96 |
97 | # Layers: dec_deconv(i)a, dec_deconv(i)b, upsample(i-1); i=4..2
98 | def _decode_block_4_3_2() -> nn.Module:
99 | return nn.Sequential(
100 | nn.Conv2d(dim * 3, dim * 2, 3, stride=1, padding=1),
101 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
102 | nn.Conv2d(dim * 2, dim * 2, 3, stride=1, padding=1),
103 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
104 | nn.Upsample(scale_factor=2, mode="nearest"),
105 | )
106 |
107 | # Separate instances of same decode module definition created
108 | self.decode_block_4 = _decode_block_4_3_2()
109 | self.decode_block_3 = _decode_block_4_3_2()
110 | self.decode_block_2 = _decode_block_4_3_2()
111 |
112 | # Layers: dec_conv1a, dec_conv1b, dec_conv1c,
113 | self.decode_block_1 = nn.Sequential(
114 | nn.Conv2d(dim * 2 + in_channels, dim * 2, 3, stride=1, padding=1),
115 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
116 | nn.Conv2d(dim * 2, dim * 2, 3, stride=1, padding=1),
117 | nn.LeakyReLU(negative_slope=0.1, inplace=True),
118 | )
119 |
120 | ####################################
121 | # Output Block
122 | ####################################
123 |
124 |
125 | # nin_a,b,c, linear_act
126 | self.output_conv = nn.Conv2d(dim * 2, out_channels, 1)
127 |
128 | # Initialize weights
129 | self.init_weights()
130 |
131 | def init_weights(self):
132 | """Initializes weights using Kaiming He et al. (2015).
133 |
134 | Only convolution layers have learnable weights. All convolutions use a leaky
135 | relu activation function (negative_slope = 0.1) except the last which is just
136 | a linear output.
137 | """
138 | with torch.no_grad():
139 | self._init_weights()
140 |
141 | def _init_weights(self):
142 | for m in self.modules():
143 | if isinstance(m, nn.Conv2d):
144 | nn.init.kaiming_normal_(m.weight.data, a=0.1)
145 | m.bias.data.zero_()
146 |
147 | # Initialise last output layer
148 | if self.zero_output:
149 | self.output_conv.weight.zero_()
150 | else:
151 | nn.init.kaiming_normal_(self.output_conv.weight.data, nonlinearity="linear")
152 |
153 | def forward_train(self, x):
154 |
155 | # Encoder
156 | pool1 = self.encode_block_1(x)
157 | pool2 = self.encode_block_2(pool1)
158 | pool3 = self.encode_block_3(pool2)
159 | pool4 = self.encode_block_4(pool3)
160 | pool5 = self.encode_block_5(pool4)
161 | encoded = self.encode_block_6(pool5)
162 |
163 | # Decoder
164 | upsample5 = self.decode_block_6(encoded)
165 | concat5 = torch.cat((upsample5, pool4), dim=1)
166 | upsample4 = self.decode_block_5(concat5)
167 | concat4 = torch.cat((upsample4, pool3), dim=1)
168 | upsample3 = self.decode_block_4(concat4)
169 | concat3 = torch.cat((upsample3, pool2), dim=1)
170 | upsample2 = self.decode_block_3(concat3)
171 | concat2 = torch.cat((upsample2, pool1), dim=1)
172 | upsample1 = self.decode_block_2(concat2)
173 | concat1 = torch.cat((upsample1, x), dim=1)
174 | x = self.decode_block_1(concat1)
175 |
176 | x = self.output_conv(x)
177 |
178 | return x
179 |
180 | def forward_test(self, x):
181 | n, c, h, w = x.shape
182 | if h < w:
183 | x = torch.nn.functional.pad(x, [0, 0, 0, w - h], mode='reflect')
184 | else:
185 | x = torch.nn.functional.pad(x, [0, h - w, 0, 0], mode='reflect')
186 | x = self.forward_train(x)
187 | x = x[:, :, :h, :w]
188 | return x
189 |
190 | def forward(self, x):
191 | if self.training:
192 | return self.forward_train(x)
193 | else:
194 | return self.forward_test(x)
195 |
196 | @staticmethod
197 | def input_wh_mul() -> int:
198 | """Multiple that both the width and height dimensions of an input must be to be
199 | processed by the network. This is devised from the number of pooling layers that
200 | reduce the input size.
201 |
202 | Returns:
203 | int: Dimension multiplier
204 | """
205 | max_pool_layers = 5
206 | return 2 ** max_pool_layers
--------------------------------------------------------------------------------
/option/knowledge_distillation_sidd.json:
--------------------------------------------------------------------------------
1 | {
2 | // model
3 | "model": "APBSNDistillModel",
4 | "pd_a": 5,
5 | "pd_b": 2,
6 | "R3": true,
7 | "R3_T": 1,
8 | "R3_p": 0,
9 | // net
10 | "networks": [{
11 | "name": "bsn",
12 | "type": "TBSN",
13 | "args": {
14 | "in_ch": 3,
15 | "out_ch": 3
16 | }
17 | , "path": "../pretrained_models/model_sidd.pth"
18 | },{
19 | "name": "network",
20 | "type": "Unet",
21 | "args": {}
22 | }],
23 | // datasets
24 | "train_dataset": {
25 | "type": "SIDDMediumTrainDataset",
26 | "args": {
27 | "pin_memory": true,
28 | "patch_size": 128
29 | },
30 | "batch_size": 8 // used for base train
31 | },
32 | "validation_datasets": [{
33 | "type": "SIDDValidationDataset",
34 | "args": {}
35 | }],
36 | // training parameters
37 | "lr": 3e-4,
38 | "print_every": 1000,
39 | "save_every": 10000,
40 | "validate_every": 10000,
41 | "num_iters": 100000,
42 | "log_dir": "log",
43 | "log_file": "log/log.out"
44 | }
45 |
--------------------------------------------------------------------------------
/option/tbsn_dnd.json:
--------------------------------------------------------------------------------
1 | {
2 | // model
3 | "model": "APBSNModel",
4 | "pd_a": 5,
5 | "pd_b": 2,
6 | "R3": true,
7 | "R3_T": 8,
8 | "R3_p": 0.16,
9 | // net
10 | "networks": [{
11 | "name": "bsn",
12 | "type": "TBSN",
13 | "args": {
14 | "in_ch": 3,
15 | "out_ch": 3
16 | }
17 | , "path": "../pretrained_models/model_dnd.pth" // comment this line for train
18 | }],
19 | // datasets
20 | "train_dataset": {
21 | "type": "DNDBenchmarkTrainDataset",
22 | "args": {
23 | "pin_memory": true,
24 | "patch_size": 640
25 | },
26 | "batch_size": 4 // used for base train
27 | },
28 | "validation_datasets": [{
29 | "type": "DNDBenchmarkPNGANDataset",
30 | "args": {}
31 | }],
32 | // training parameters
33 | "lr": 3e-4,
34 | "print_every": 1000,
35 | "save_every": 10000,
36 | "validate_every": 20000,
37 | "num_iters": 100000,
38 | "log_dir": "log",
39 | "log_file": "log/log.out"
40 | }
41 |
--------------------------------------------------------------------------------
/option/tbsn_sidd.json:
--------------------------------------------------------------------------------
1 | {
2 | // model
3 | "model": "APBSNModel",
4 | "pd_a": 5,
5 | "pd_b": 2,
6 | "R3": true,
7 | "R3_T": 8,
8 | "R3_p": 0.16,
9 | // net
10 | "networks": [{
11 | "name": "bsn",
12 | "type": "TBSN",
13 | "args": {
14 | "in_ch": 3,
15 | "out_ch": 3
16 | }
17 | , "path": "../pretrained_models/model_sidd.pth" // comment this line for train
18 | }],
19 | // datasets
20 | "train_dataset": {
21 | "type": "SIDDMediumTrainDataset",
22 | "args": {
23 | "pin_memory": true,
24 | "patch_size": 640
25 | },
26 | "batch_size": 4 // used for base train
27 | },
28 | "validation_datasets": [{
29 | "type": "SIDDValidationDataset",
30 | "args": {}
31 | }],
32 | // training parameters
33 | "lr": 3e-4,
34 | "print_every": 1000,
35 | "save_every": 10000,
36 | "validate_every": 20000,
37 | "num_iters": 100000,
38 | "log_dir": "log",
39 | "log_file": "log/log.out"
40 | }
41 |
--------------------------------------------------------------------------------
/pretrained_models/model_dnd.pth:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nagejacob/TBSN/678e93dd4ab816d5c736179b01f7398373db45ef/pretrained_models/model_dnd.pth
--------------------------------------------------------------------------------
/pretrained_models/model_sidd.pth:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nagejacob/TBSN/678e93dd4ab816d5c736179b01f7398373db45ef/pretrained_models/model_sidd.pth
--------------------------------------------------------------------------------
/readme.md:
--------------------------------------------------------------------------------
1 | # Rethinking Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising [Paper](https://arxiv.org/abs/2404.07846)
2 |
3 |
4 | ## Usage
5 | ### Datasets
6 | Download [SIDD](https://abdokamel.github.io/sidd/) and [DND](https://noise.visinf.tu-darmstadt.de/) datasets, and modify `dataset_path` in `dataset/base_function.py` accordingly.
7 | ```
8 | |- dataset_path
9 | |- SIDD
10 | |- SIDD_Medium_Srgb
11 | |- Data
12 | |- 0001_001_S6_00100_00060_3200_L
13 | |- 0002_001_S6_00100_00020_3200_N
14 | |- ...
15 | |- SIDD_Validation
16 | |- ValidationNoisyBlocksSrgb.mat
17 | |- ValidationGtBlocksSrgb.mat
18 | |- SIDD_Benchmark
19 | |- BenchmarkNoisyBlocksSrgb.mat
20 | |- DND
21 | |- info.mat
22 | |- images_srgb
23 | ```
24 |
25 | ### Validation
26 | Validate on SIDD Validation dataset,
27 | ```
28 | cd validate
29 | python base.py --config_file "../option/tbsn_sidd.json"
30 | ```
31 |
32 | ### Training
33 | Training on SIDD Medium dataset,
34 | ```
35 | sh train.sh
36 | ```
37 |
38 | ## Citation
39 | If you make use of our work, please cite our paper.
40 | ```bibtex
41 | @inproceedings{li2025rethinking,
42 | title={Rethinking Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising},
43 | author={Li, Junyi and Zhang, Zhilu and Zuo, Wangmeng},
44 | booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
45 | year={2025}
46 | }
47 | ```
48 |
--------------------------------------------------------------------------------
/slurm.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 | #SBATCH --job-name=default
3 | #SBATCH --gres=gpu:2080:2
4 | #SBATCH --output=/mnt/hdd0/lijunyi/slurm_logs/%j.out
5 | #SBATCH --error=/mnt/hdd0/lijunyi/slurm_logs/%j.err
6 |
7 | source activate
8 | conda activate ~/anaconda/pytorch2
9 | cd /mnt/hdd0/lijunyi/codes/TBSN
10 | sh train.sh
--------------------------------------------------------------------------------
/train.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | experiment_home_dir="experiments"
4 | config="tbsn_sidd.json"
5 |
6 | experiment_name=$(python train/experiment_name.py --config option/${config})
7 | #experiment_name='test'
8 |
9 | experiment_dir="${experiment_home_dir}/${experiment_name}"
10 | echo "experiment dir: ${experiment_dir}"
11 |
12 | if [ ! -d "${experiment_home_dir}" ]
13 | then
14 | mkdir "${experiment_home_dir}"
15 | fi
16 |
17 | if [ ! -d "${experiment_dir}" ]
18 | then
19 | mkdir "${experiment_dir}"
20 | else
21 | echo "experiment dir exists"
22 | fi
23 |
24 | cp -r "dataset" "${experiment_dir}"
25 | cp -r "model" "${experiment_dir}"
26 | cp -r "network" "${experiment_dir}"
27 | cp -r "option" "${experiment_dir}"
28 | cp -r "train" "${experiment_dir}"
29 | cp -r "util" "${experiment_dir}"
30 | cp -r "validate" "${experiment_dir}"
31 |
32 | if [ ! -d "${experiment_dir}/log" ]
33 | then
34 | mkdir "${experiment_dir}/log"
35 | fi
36 |
37 | cd ${experiment_dir}
38 | export PYTHONPATH=$PWD:$PYTHONPATH
39 | python train/base.py --config option/${config}
--------------------------------------------------------------------------------
/train/base.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | from torch.utils.data import DataLoader
3 | from util.build import build
4 | from util.io import log
5 | from util.option import parse, recursive_log
6 | from validate.base import validate_sidd, validate_dnd, validate_synthetic
7 |
8 |
9 | def main(opt):
10 | train_dataset_opt = opt['train_dataset']
11 | TrainDataset = getattr(__import__('dataset'), train_dataset_opt['type'])
12 | train_set = build(TrainDataset, train_dataset_opt['args'])
13 | train_loader = DataLoader(train_set, batch_size=train_dataset_opt['batch_size'], shuffle=True,
14 | num_workers=4, drop_last=True)
15 |
16 | validation_loaders = []
17 | for validation_dataset_opt in opt['validation_datasets']:
18 | ValidationDataset = getattr(__import__('dataset'), validation_dataset_opt['type'])
19 | validation_set = build(ValidationDataset, validation_dataset_opt['args'])
20 | validation_loader = DataLoader(validation_set, batch_size=1)
21 | validation_loaders.append(validation_loader)
22 |
23 | Model = getattr(__import__('model'), opt['model'])
24 | model = Model(opt)
25 | model.data_parallel()
26 |
27 | def train_step(data):
28 | model.train_step(data)
29 |
30 | if model.iter % opt['print_every'] == 0:
31 | model.log()
32 |
33 | if model.iter % opt['save_every'] == 0:
34 | model.save_net()
35 |
36 | if model.iter % opt['validate_every'] == 0:
37 | message = 'iter: %d, ' % model.iter
38 | for validation_loader in validation_loaders:
39 | if 'SIDD' in validation_loader.dataset.__class__.__name__:
40 | psnr, ssim = validate_sidd(model, validation_loader)
41 | elif 'DND' in validation_loader.dataset.__class__.__name__:
42 | psnr, ssim = validate_dnd(model, validation_loader)
43 | elif 'Synthetic' in validation_loader.dataset.__class__.__name__:
44 | psnr, ssim = validate_synthetic(model, validation_loader)
45 | message += '%s: %6.4f/%6.4f ' % (validation_loader.dataset.__class__.__name__, psnr, ssim)
46 | log(opt['log_file'], message + '\n')
47 |
48 | while True:
49 | for i, data in enumerate(train_loader):
50 | train_step(data)
51 | if model.iter >= opt['num_iters']:
52 | return
53 |
54 |
55 | if __name__ == '__main__':
56 | parser = argparse.ArgumentParser(description="Train the denoiser")
57 | parser.add_argument("--config", type=str, default='')
58 | argspar = parser.parse_args()
59 |
60 | opt = parse(argspar.config)
61 | recursive_log(opt['log_file'], opt)
62 |
63 | main(opt)
--------------------------------------------------------------------------------
/train/experiment_name.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('.')
3 | sys.path.append('..')
4 | import argparse
5 | import datetime
6 | from util.option import parse
7 |
8 | if __name__ == '__main__':
9 | parser = argparse.ArgumentParser(description="Train the denoiser")
10 | parser.add_argument("--config", type=str, default='')
11 | argspar = parser.parse_args()
12 |
13 | opt = parse(argspar.config)
14 |
15 | now = datetime.datetime.now()
16 | date = now.strftime("%m%d-%H%M%S")
17 |
18 | model = opt['model']
19 | network = opt['networks'][0]['type']
20 |
21 | name = date + '_' + model + '_' + network
22 | print(name)
--------------------------------------------------------------------------------
/util/build.py:
--------------------------------------------------------------------------------
1 | def build(obj_type, args):
2 | return obj_type(**args)
--------------------------------------------------------------------------------
/util/io.py:
--------------------------------------------------------------------------------
1 | import datetime
2 | import imageio
3 | import numpy as np
4 | import os
5 | from PIL import Image
6 |
7 | IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
8 |
9 | def is_image_file(filename):
10 | return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) and (not filename.startswith('.'))
11 |
12 | def get_image_paths(dataroot):
13 | paths = None # return None if dataroot is None
14 | if isinstance(dataroot, str):
15 | paths = sorted(_get_paths_from_images(dataroot))
16 | elif isinstance(dataroot, list):
17 | paths = []
18 | for i in dataroot:
19 | paths += sorted(_get_paths_from_images(i))
20 | return paths
21 |
22 | def _get_paths_from_images(path):
23 | assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
24 | images = []
25 | for dirpath, _, fnames in sorted(os.walk(path)):
26 | for fname in sorted(fnames):
27 | if is_image_file(fname):
28 | img_path = os.path.join(dirpath, fname)
29 | images.append(img_path)
30 | assert images, '{:s} has no valid image file'.format(path)
31 | return images
32 |
33 | def open_image_uint8(image_file, to_gray=False):
34 | if to_gray:
35 | image = np.asarray(Image.open(image_file).convert('L'))
36 | else:
37 | image = imageio.imread(image_file).astype(np.uint8)
38 |
39 | if len(image.shape) == 2:
40 | image = np.expand_dims(image, axis=0)
41 | image = np.repeat(image, 3, axis=0)
42 | elif len(image.shape) == 3:
43 | image = np.transpose(image, (2, 0, 1))
44 | image = image[:3, :, :]
45 |
46 | return image
47 |
48 |
49 | def date_time():
50 | now = datetime.datetime.now()
51 | date_time = now.strftime("%Y-%m-%d, %H:%M:%S")
52 | return date_time
53 |
54 |
55 | def log(log_file, str, also_print=True, with_time=True):
56 | with open(log_file, 'a+') as F:
57 | if with_time:
58 | F.write(date_time() + ' ')
59 | F.write(str)
60 | if also_print:
61 | if with_time:
62 | print(date_time(), end=' ')
63 | print(str, end='')
64 |
65 |
66 | # save numpy image in shape 3xHxW
67 | def np2image(image, image_file):
68 | image = np.transpose(image, (1, 2, 0))
69 | image = np.clip(image, 0., 1.)
70 | image = image * 255.
71 | image = image.astype(np.uint8)
72 | if 1 == image.shape[2]:
73 | image = image[:, :, 0]
74 | imageio.imwrite(image_file, image)
75 |
76 | # save tensor image in shape 1x3xHxW
77 | def tensor2image(image, image_file):
78 | image = image.detach().cpu().squeeze(0).numpy()
79 | np2image(image, image_file)
--------------------------------------------------------------------------------
/util/option.py:
--------------------------------------------------------------------------------
1 | from collections import OrderedDict
2 | import json
3 | from util.io import log
4 |
5 | def parse(opt_path):
6 | # ----------------------------------------
7 | # remove comments starting with '//'
8 | # ----------------------------------------
9 | json_str = ''
10 | with open(opt_path, 'r') as f:
11 | for line in f:
12 | line = line.split('//')[0] + '\n'
13 | json_str += line
14 |
15 | # ----------------------------------------
16 | # initialize opt
17 | # ----------------------------------------
18 | opt = json.loads(json_str, object_pairs_hook=OrderedDict)
19 |
20 | return opt
21 |
22 |
23 | def recursive_print(src, dpth=0, key=None):
24 | """ Recursively prints nested elements."""
25 | tabs = lambda n: ' ' * n * 4 # or 2 or 8 or...
26 |
27 | if isinstance(src, dict):
28 | if key is not None:
29 | print(tabs(dpth) + '%s: ' % (key))
30 | for key, value in src.items():
31 | recursive_print(value, dpth + 1, key)
32 | elif isinstance(src, list):
33 | if key is not None:
34 | print(tabs(dpth) + '%s: ' % (key))
35 | for litem in src:
36 | recursive_print(litem, dpth)
37 | else:
38 | if key is not None:
39 | print(tabs(dpth) + '%s: %s' % (key, src))
40 |
41 |
42 | def recursive_log(log_file, src, dpth=0, key=None):
43 | """ Recursively prints nested elements."""
44 | tabs = lambda n: ' ' * n * 4 # or 2 or 8 or...
45 |
46 | if isinstance(src, dict):
47 | if key is not None:
48 | log(log_file, tabs(dpth) + '%s: \n' % (key), with_time=False)
49 | for key, value in src.items():
50 | recursive_log(log_file, value, dpth + 1, key)
51 | elif isinstance(src, list):
52 | if key is not None:
53 | log(log_file, tabs(dpth) + '%s: \n' % (key), with_time=False)
54 | for litem in src:
55 | recursive_log(log_file, litem, dpth)
56 | else:
57 | if key is not None:
58 | log(log_file, tabs(dpth) + '%s: %s\n' % (key, src), with_time=False)
--------------------------------------------------------------------------------
/validate/base.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | import argparse
4 | import lpips
5 | from skimage.metrics import peak_signal_noise_ratio, structural_similarity
6 | import torch
7 | from torch.utils.data import DataLoader
8 | from tqdm import tqdm
9 | from validate.base_function import calculate_ssim
10 | from util.build import build
11 | from util.option import parse, recursive_print
12 |
13 | loss_fn_alex = lpips.LPIPS(net='alex').cuda() # best forward scores
14 | def lpips_norm(img):
15 | img = img * 2. - 1
16 | return img
17 |
18 | def validate_sidd(model, sidd_loader):
19 | psnrs, ssims, lpipss, count = 0, 0, 0, 0
20 | for data in tqdm(sidd_loader):
21 | output = model.validation_step(data)
22 | output = torch.floor(output * 255. + 0.5) / 255.
23 | output = torch.clamp(output, 0, 1)
24 | lpips = loss_fn_alex(lpips_norm(output), lpips_norm(data['H'].cuda()))
25 | output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
26 | gt = data['H'].squeeze(0).permute(1, 2, 0).numpy()
27 | psnr = peak_signal_noise_ratio(output, gt, data_range=1)
28 | ssim = structural_similarity(output, gt, data_range=1, channel_axis=2, gaussian_weights=True, sigma=1.5, win_size=11)
29 |
30 | psnrs += psnr
31 | ssims += ssim
32 | lpipss += lpips
33 | count += 1
34 |
35 | return psnrs / count, ssims / count, lpipss / count
36 |
37 | def validate_dnd(model, dnd_loader):
38 | psnrs, ssims, count = 0, 0, 0
39 | f = open('')
40 | for data in tqdm(dnd_loader):
41 | output = model.validation_step(data)
42 | output = torch.clamp(output, 0, 1)
43 | output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
44 | gt = data['H'].squeeze(0).permute(1, 2, 0).numpy()
45 | psnr = peak_signal_noise_ratio(output, gt, data_range=1)
46 |
47 | psnrs += psnr
48 | count += 1
49 | return psnrs / count, ssims / count
50 |
51 | def validate_synthetic(model, synthetic_loader):
52 | psnrs, ssims, count = 0, 0, 0
53 | for data in tqdm(synthetic_loader):
54 | n, c, h, w = data['L'].shape
55 | if h < w:
56 | data['L'] = torch.nn.functional.pad(data['L'], [0, 0, 0, w - h], mode='reflect')
57 | elif h > w:
58 | data['L'] = torch.nn.functional.pad(data['L'], [0, h - w, 0, 0], mode='reflect')
59 | output = model.validation_step(data)
60 | output = output[:, :, :h, :w].cpu().squeeze(0).permute(1, 2, 0).numpy()
61 | gt = data['H'].squeeze(0).permute(1, 2, 0).numpy()
62 | psnr = peak_signal_noise_ratio(output, gt, data_range=1)
63 | ssim = calculate_ssim(output * 255., gt * 255.)
64 |
65 | psnrs += psnr
66 | ssims += ssim
67 | count += 1
68 | return psnrs / count, ssims / count
69 |
70 | def main(opt):
71 | validation_loaders = []
72 | for validation_dataset_opt in opt['validation_datasets']:
73 | ValidationDataset = getattr(__import__('dataset'), validation_dataset_opt['type'])
74 | validation_set = build(ValidationDataset, validation_dataset_opt['args'])
75 | validation_loader = DataLoader(validation_set, batch_size=1)
76 | validation_loaders.append(validation_loader)
77 |
78 | Model = getattr(__import__('model'), opt['model'])
79 | model = Model(opt)
80 | model.data_parallel()
81 |
82 | for validation_loader in validation_loaders:
83 | if 'SIDD' in validation_loader.dataset.__class__.__name__:
84 | psnr, ssim, lpips = validate_sidd(model, validation_loader)
85 | elif 'DND' in validation_loader.dataset.__class__.__name__:
86 | psnr, ssim = validate_dnd(model, validation_loader)
87 | else:
88 | psnr, ssim = validate_synthetic(model, validation_loader)
89 | print('%s, psnr: %6.4f, ssim: %6.4f, lpips: %.4f' % (validation_loader.dataset.__class__.__name__, psnr, ssim, lpips))
90 |
91 |
92 | if __name__ == '__main__':
93 | parser = argparse.ArgumentParser(description="Validate the denoiser")
94 | parser.add_argument("--config_file", type=str, default='../option/tbsn_sidd.json')
95 | argspar = parser.parse_args()
96 |
97 | opt = parse(argspar.config_file)
98 | recursive_print(opt)
99 |
100 | main(opt)
101 |
--------------------------------------------------------------------------------
/validate/base_function.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 |
4 | def reorder_image(img, input_order='HWC'):
5 | """Reorder images to 'HWC' order.
6 |
7 | If the input_order is (h, w), return (h, w, 1);
8 | If the input_order is (c, h, w), return (h, w, c);
9 | If the input_order is (h, w, c), return as it is.
10 |
11 | Args:
12 | img (ndarray): Input image.
13 | input_order (str): Whether the input order is 'HWC' or 'CHW'.
14 | If the input image shape is (h, w), input_order will not have
15 | effects. Default: 'HWC'.
16 |
17 | Returns:
18 | ndarray: reordered image.
19 | """
20 |
21 | if input_order not in ['HWC', 'CHW']:
22 | raise ValueError(f"Wrong input_order {input_order}. Supported input_orders are 'HWC' and 'CHW'")
23 | if len(img.shape) == 2:
24 | img = img[..., None]
25 | if input_order == 'CHW':
26 | img = img.transpose(1, 2, 0)
27 | return img
28 |
29 |
30 | def _ssim(img, img2):
31 | """Calculate SSIM (structural similarity) for one channel images.
32 |
33 | It is called by func:`calculate_ssim`.
34 |
35 | Args:
36 | img (ndarray): Images with range [0, 255] with order 'HWC'.
37 | img2 (ndarray): Images with range [0, 255] with order 'HWC'.
38 |
39 | Returns:
40 | float: SSIM result.
41 | """
42 |
43 | c1 = (0.01 * 255)**2
44 | c2 = (0.03 * 255)**2
45 | kernel = cv2.getGaussianKernel(11, 1.5)
46 | window = np.outer(kernel, kernel.transpose())
47 |
48 | mu1 = cv2.filter2D(img, -1, window)[5:-5, 5:-5] # valid mode for window size 11
49 | mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
50 | mu1_sq = mu1**2
51 | mu2_sq = mu2**2
52 | mu1_mu2 = mu1 * mu2
53 | sigma1_sq = cv2.filter2D(img**2, -1, window)[5:-5, 5:-5] - mu1_sq
54 | sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
55 | sigma12 = cv2.filter2D(img * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
56 |
57 | ssim_map = ((2 * mu1_mu2 + c1) * (2 * sigma12 + c2)) / ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
58 | return ssim_map.mean()
59 |
60 |
61 | def calculate_ssim(img, img2, crop_border=0, input_order='HWC', test_y_channel=False, **kwargs):
62 | """Calculate SSIM (structural similarity).
63 |
64 | ``Paper: Image quality assessment: From error visibility to structural similarity``
65 |
66 | The results are the same as that of the official released MATLAB code in
67 | https://ece.uwaterloo.ca/~z70wang/research/ssim/.
68 |
69 | For three-channel images, SSIM is calculated for each channel and then
70 | averaged.
71 |
72 | Args:
73 | img (ndarray): Images with range [0, 255].
74 | img2 (ndarray): Images with range [0, 255].
75 | crop_border (int): Cropped pixels in each edge of an image. These pixels are not involved in the calculation.
76 | input_order (str): Whether the input order is 'HWC' or 'CHW'.
77 | Default: 'HWC'.
78 | test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
79 |
80 | Returns:
81 | float: SSIM result.
82 | """
83 |
84 | assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
85 | if input_order not in ['HWC', 'CHW']:
86 | raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are "HWC" and "CHW"')
87 | img = reorder_image(img, input_order=input_order)
88 | img2 = reorder_image(img2, input_order=input_order)
89 |
90 | if crop_border != 0:
91 | img = img[crop_border:-crop_border, crop_border:-crop_border, ...]
92 | img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
93 |
94 | # if test_y_channel:
95 | # img = to_y_channel(img)
96 | # img2 = to_y_channel(img2)
97 |
98 | img = img.astype(np.float64)
99 | img2 = img2.astype(np.float64)
100 |
101 | ssims = []
102 | for i in range(img.shape[2]):
103 | ssims.append(_ssim(img[..., i], img2[..., i]))
104 | return np.array(ssims).mean()
--------------------------------------------------------------------------------
/validate/visualization.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | import argparse
4 | import cv2
5 | from network.tbsn import DilatedMDTA, DilatedOCA
6 | import numpy as np
7 | from PIL import Image
8 | from skimage.metrics import peak_signal_noise_ratio
9 | import torch
10 | from torch.utils.data import DataLoader
11 | from tqdm import tqdm
12 | from validate.base_function import calculate_ssim
13 | from util.build import build
14 | from util.io import tensor2image
15 | from util.option import parse, recursive_print
16 | from torch.cuda.amp import autocast
17 |
18 |
19 | def main(opt):
20 | validation_loaders = []
21 | for validation_dataset_opt in opt['validation_datasets']:
22 | ValidationDataset = getattr(__import__('dataset'), validation_dataset_opt['type'])
23 | validation_set = build(ValidationDataset, validation_dataset_opt['args'])
24 | validation_loader = DataLoader(validation_set, batch_size=1)
25 | validation_loaders.append(validation_loader)
26 |
27 | network = DilatedMDTA(dim=64, num_heads=1, bias=False)
28 | # network = DilatedOCA(dim=64, window_size=8, overlap_ratio=0.5, num_heads=2, dim_head=16, bias=False)
29 | gradient, count = None, 0
30 | with autocast():
31 | for data in validation_loaders[0]:
32 | input = data['L']
33 | input = torch.rand((1, 64, 16, 16))
34 | input.requires_grad = True
35 | output = network(input)
36 | center_output = torch.mean(output[:, :, 8, 8])
37 | center_output.backward()
38 | if gradient is None:
39 | gradient = torch.sum(torch.abs(input.grad), dim=1, keepdim=True).cpu()
40 | else:
41 | gradient += torch.sum(torch.abs(input.grad), dim=1, keepdim=True).cpu()
42 | count += 1
43 |
44 | if count == 50:
45 | break
46 |
47 | gradient = gradient / count * 50
48 | gradient = torch.clamp(gradient, 0, 1)
49 | gradient = gradient[0, 0].numpy()
50 | cam = cv2.applyColorMap(np.uint8(gradient * 255), cv2.COLORMAP_INFERNO)
51 | cam = cv2.cvtColor(cam, cv2.COLOR_BGR2RGB)
52 | image = Image.fromarray(cam)
53 | image.save('gradient_tbsn.png')
54 |
55 | if __name__ == '__main__':
56 | parser = argparse.ArgumentParser(description="Validate the denoiser")
57 | parser.add_argument("--config_file", type=str, default='../option/tbsn.json')
58 | argspar = parser.parse_args()
59 |
60 | opt = parse(argspar.config_file)
61 | recursive_print(opt)
62 |
63 | main(opt)
--------------------------------------------------------------------------------
/validate/visualize_receptive_field.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append('..')
3 | import argparse
4 | import cv2
5 | import numpy as np
6 | from PIL import Image
7 | from skimage.metrics import peak_signal_noise_ratio
8 | import torch
9 | from torch.utils.data import DataLoader
10 | from tqdm import tqdm
11 | from validate.base_function import calculate_ssim
12 | from util.build import build
13 | from util.io import tensor2image
14 | from util.option import parse, recursive_print
15 | from torch.cuda.amp import autocast
16 |
17 |
18 | def main(opt):
19 | validation_loaders = []
20 | for validation_dataset_opt in opt['validation_datasets']:
21 | ValidationDataset = getattr(__import__('dataset'), validation_dataset_opt['type'])
22 | validation_set = build(ValidationDataset, validation_dataset_opt['args'])
23 | validation_loader = DataLoader(validation_set, batch_size=1)
24 | validation_loaders.append(validation_loader)
25 |
26 | Model = getattr(__import__('model'), opt['model'])
27 | model = Model(opt)
28 | model.data_parallel()
29 |
30 | network = model.networks['bsn']
31 | gradient, count = None, 0
32 | with autocast():
33 | for data in validation_loaders[0]:
34 | input = data['L']
35 | input = input[:, :, 64:192, 64:192]
36 | input.requires_grad = True
37 | output = network(input)
38 | center_output = torch.mean(output[:, :, 64, 64])
39 | center_output.backward()
40 | if gradient is None:
41 | gradient = torch.sum(torch.abs(input.grad), dim=1, keepdim=True).cpu()
42 | else:
43 | gradient += torch.sum(torch.abs(input.grad), dim=1, keepdim=True).cpu()
44 | count += 1
45 |
46 | if count == 50:
47 | break
48 |
49 | gradient = gradient / count * 10
50 | gradient = torch.clamp(gradient, 0, 1)
51 | gradient = gradient[0, 0].numpy()
52 | cam = cv2.applyColorMap(np.uint8(gradient * 255), cv2.COLORMAP_INFERNO)
53 | cam = cv2.cvtColor(cam, cv2.COLOR_BGR2RGB)
54 | image = Image.fromarray(cam)
55 | image.save('gradient_tbsn.png')
56 |
57 | if __name__ == '__main__':
58 | parser = argparse.ArgumentParser(description="Validate the denoiser")
59 | parser.add_argument("--config_file", type=str, default='../option/tbsn.json')
60 | argspar = parser.parse_args()
61 |
62 | opt = parse(argspar.config_file)
63 | recursive_print(opt)
64 |
65 | main(opt)
66 |
--------------------------------------------------------------------------------