diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..ac26bbc
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,163 @@
+*/.vscode/*
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..f288702
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,674 @@
+ GNU GENERAL PUBLIC LICENSE
+ Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc.
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+ Preamble
+
+ The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+ The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works. By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users. We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors. You can apply it to
+your programs, too.
+
+ When we speak of free software, we are referring to freedom, not
+price. Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+ To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights. Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+ For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received. You must make sure that they, too, receive
+or can get the source code. And you must show them these terms so they
+know their rights.
+
+ Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+ For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software. For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+ Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so. This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software. The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable. Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products. If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+ Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary. To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+ The precise terms and conditions for copying, distribution and
+modification follow.
+
+ TERMS AND CONDITIONS
+
+ 0. Definitions.
+
+ "This License" refers to version 3 of the GNU General Public License.
+
+ "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+ "The Program" refers to any copyrightable work licensed under this
+License. Each licensee is addressed as "you". "Licensees" and
+"recipients" may be individuals or organizations.
+
+ To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy. The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+ A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+ To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy. Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+ To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies. Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+ An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License. If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+ 1. Source Code.
+
+ The "source code" for a work means the preferred form of the work
+for making modifications to it. "Object code" means any non-source
+form of a work.
+
+ A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+ The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form. A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+ The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities. However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work. For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+ The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+ The Corresponding Source for a work in source code form is that
+same work.
+
+ 2. Basic Permissions.
+
+ All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met. This License explicitly affirms your unlimited
+permission to run the unmodified Program. The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work. This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+ You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force. You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright. Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+ Conveying under any other circumstances is permitted solely under
+the conditions stated below. Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+ 3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+ No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+ When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+ 4. Conveying Verbatim Copies.
+
+ You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+ You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+ 5. Conveying Modified Source Versions.
+
+ You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+ a) The work must carry prominent notices stating that you modified
+ it, and giving a relevant date.
+
+ b) The work must carry prominent notices stating that it is
+ released under this License and any conditions added under section
+ 7. This requirement modifies the requirement in section 4 to
+ "keep intact all notices".
+
+ c) You must license the entire work, as a whole, under this
+ License to anyone who comes into possession of a copy. This
+ License will therefore apply, along with any applicable section 7
+ additional terms, to the whole of the work, and all its parts,
+ regardless of how they are packaged. This License gives no
+ permission to license the work in any other way, but it does not
+ invalidate such permission if you have separately received it.
+
+ d) If the work has interactive user interfaces, each must display
+ Appropriate Legal Notices; however, if the Program has interactive
+ interfaces that do not display Appropriate Legal Notices, your
+ work need not make them do so.
+
+ A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit. Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+ 6. Conveying Non-Source Forms.
+
+ You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+ a) Convey the object code in, or embodied in, a physical product
+ (including a physical distribution medium), accompanied by the
+ Corresponding Source fixed on a durable physical medium
+ customarily used for software interchange.
+
+ b) Convey the object code in, or embodied in, a physical product
+ (including a physical distribution medium), accompanied by a
+ written offer, valid for at least three years and valid for as
+ long as you offer spare parts or customer support for that product
+ model, to give anyone who possesses the object code either (1) a
+ copy of the Corresponding Source for all the software in the
+ product that is covered by this License, on a durable physical
+ medium customarily used for software interchange, for a price no
+ more than your reasonable cost of physically performing this
+ conveying of source, or (2) access to copy the
+ Corresponding Source from a network server at no charge.
+
+ c) Convey individual copies of the object code with a copy of the
+ written offer to provide the Corresponding Source. This
+ alternative is allowed only occasionally and noncommercially, and
+ only if you received the object code with such an offer, in accord
+ with subsection 6b.
+
+ d) Convey the object code by offering access from a designated
+ place (gratis or for a charge), and offer equivalent access to the
+ Corresponding Source in the same way through the same place at no
+ further charge. You need not require recipients to copy the
+ Corresponding Source along with the object code. If the place to
+ copy the object code is a network server, the Corresponding Source
+ may be on a different server (operated by you or a third party)
+ that supports equivalent copying facilities, provided you maintain
+ clear directions next to the object code saying where to find the
+ Corresponding Source. Regardless of what server hosts the
+ Corresponding Source, you remain obligated to ensure that it is
+ available for as long as needed to satisfy these requirements.
+
+ e) Convey the object code using peer-to-peer transmission, provided
+ you inform other peers where the object code and Corresponding
+ Source of the work are being offered to the general public at no
+ charge under subsection 6d.
+
+ A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+ A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling. In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage. For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product. A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+ "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source. The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+ If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information. But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+ The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed. Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+ Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+ 7. Additional Terms.
+
+ "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law. If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+ When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it. (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.) You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+ Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+ a) Disclaiming warranty or limiting liability differently from the
+ terms of sections 15 and 16 of this License; or
+
+ b) Requiring preservation of specified reasonable legal notices or
+ author attributions in that material or in the Appropriate Legal
+ Notices displayed by works containing it; or
+
+ c) Prohibiting misrepresentation of the origin of that material, or
+ requiring that modified versions of such material be marked in
+ reasonable ways as different from the original version; or
+
+ d) Limiting the use for publicity purposes of names of licensors or
+ authors of the material; or
+
+ e) Declining to grant rights under trademark law for use of some
+ trade names, trademarks, or service marks; or
+
+ f) Requiring indemnification of licensors and authors of that
+ material by anyone who conveys the material (or modified versions of
+ it) with contractual assumptions of liability to the recipient, for
+ any liability that these contractual assumptions directly impose on
+ those licensors and authors.
+
+ All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10. If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term. If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+ If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+ Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+ 8. Termination.
+
+ You may not propagate or modify a covered work except as expressly
+provided under this License. Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+ However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+ Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+ Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License. If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+ 9. Acceptance Not Required for Having Copies.
+
+ You are not required to accept this License in order to receive or
+run a copy of the Program. Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance. However,
+nothing other than this License grants you permission to propagate or
+modify any covered work. These actions infringe copyright if you do
+not accept this License. Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+ 10. Automatic Licensing of Downstream Recipients.
+
+ Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License. You are not responsible
+for enforcing compliance by third parties with this License.
+
+ An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations. If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+ You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License. For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+ 11. Patents.
+
+ A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based. The
+work thus licensed is called the contributor's "contributor version".
+
+ A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version. For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+ Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+ In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement). To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+ If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients. "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+ If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+ A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License. You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+ Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+ 12. No Surrender of Others' Freedom.
+
+ If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License. If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all. For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+ 13. Use with the GNU Affero General Public License.
+
+ Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work. The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+ 14. Revised Versions of this License.
+
+ The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time. Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+ Each version is given a distinguishing version number. If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation. If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+ If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+ Later license versions may give you additional or different
+permissions. However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+ 15. Disclaimer of Warranty.
+
+ THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+ 16. Limitation of Liability.
+
+ IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+ 17. Interpretation of Sections 15 and 16.
+
+ If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+ END OF TERMS AND CONDITIONS
+
+ How to Apply These Terms to Your New Programs
+
+ If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+ To do so, attach the following notices to the program. It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+
+ Copyright (C)
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+
+Also add information on how to contact you by electronic and paper mail.
+
+ If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+ Copyright (C)
+ This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+ This is free software, and you are welcome to redistribute it
+ under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License. Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+ You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+.
+
+ The GNU General Public License does not permit incorporating your program
+into proprietary programs. If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library. If this is what you want to do, use the GNU Lesser General
+Public License instead of this License. But first, please read
+.
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..8ae2102
--- /dev/null
+++ b/README.md
@@ -0,0 +1,60 @@
+# GelSight SDK
+[](https://www.gnu.org/licenses/gpl-3.0) [
](https://rpl.ri.cmu.edu/)
+
+This repository is a modified version of the [the official gsrobotics implementation](https://github.com/gelsightinc/gsrobotics), offering improvements in usability and sensor compatibility.
+
+## Key Features
+
+- **Sensor Calibration**: Added functionality for users to calibrate their own GelSight sensors.
+
+- **Expanded Sensor Compatibility**: Compatible with other GelSight sensors, including Digit and lab-made sensors.
+
+- **Low-Latency Sensor Reading**: Including an enhanced image streaming pipeline for reduced latency and frame drop, especially for GelSight Mini.
+
+Authors:
+* [Hung-Jui Huang](https://joehjhuang.github.io/) (hungjuih@andrew.cmu.edu)
+* Ruihan Gao (ruihang@andrew.cmu.edu)
+
+## Support System
+* Tested on Ubuntu 22.04
+* Tested on GelSight Mini and Digit
+* Python >= 3.9
+
+## Installation
+Clone and install gs_sdk from source:
+```bash
+git clone git@github.com:joehjhuang/gs_sdk.git
+cd gs_sdk
+pip install -e .
+```
+
+## Coordinate Conventions
+The coordinate system convention in this SDK is shown below, using the GelSight Mini sensor for illustration:
+
+| 2D (sensor image) | 3D |
+| --------------------------------- | --------------------------------- |
+|
|
|
+
+## Sensor Calibration
+For more details on sensor calibration, see the [Calibration README](calibration/README.md).
+
+## Examples
+These examples show basic usage of this GelSight SDK.
+### Sensor Streaming
+Stream images from a connected GelSight Mini:
+```python
+python examples/stream_device.py
+```
+
+### Low Latency Sensor Streaming
+Stream images with low latency and without frame dropping from a connected GelSight Mini:
+```python
+python examples/fast_stream_device.py
+```
+
+### Reconstruct Touched Surface
+Reconstruct a touched surface using the calibration model. Calibration steps are detailed in the [Calibration README](calibration/README.md).
+```python
+python examples/reconstruct.py
+```
+The reconstructed surface will be displayed and saved in `examples/data`.
\ No newline at end of file
diff --git a/assets/ball_image.png b/assets/ball_image.png
new file mode 100644
index 0000000..28760ad
Binary files /dev/null and b/assets/ball_image.png differ
diff --git a/assets/ball_indenter.jpg b/assets/ball_indenter.jpg
new file mode 100644
index 0000000..147bf54
Binary files /dev/null and b/assets/ball_indenter.jpg differ
diff --git a/assets/gsmini_frame_2D.png b/assets/gsmini_frame_2D.png
new file mode 100644
index 0000000..290743b
Binary files /dev/null and b/assets/gsmini_frame_2D.png differ
diff --git a/assets/gsmini_frame_3D.png b/assets/gsmini_frame_3D.png
new file mode 100644
index 0000000..822221c
Binary files /dev/null and b/assets/gsmini_frame_3D.png differ
diff --git a/assets/nanogui.png b/assets/nanogui.png
new file mode 100644
index 0000000..92e2de3
Binary files /dev/null and b/assets/nanogui.png differ
diff --git a/assets/pressing.jpg b/assets/pressing.jpg
new file mode 100644
index 0000000..a694d13
Binary files /dev/null and b/assets/pressing.jpg differ
diff --git a/assets/rpl.png b/assets/rpl.png
new file mode 100644
index 0000000..962fe64
Binary files /dev/null and b/assets/rpl.png differ
diff --git a/calibration/README.md b/calibration/README.md
new file mode 100644
index 0000000..b5fe019
--- /dev/null
+++ b/calibration/README.md
@@ -0,0 +1,68 @@
+## Sensor Calibration
+We provide calibration tools for new GelSight sensors, following the method described in [1]; the process generally takes less than an hour. For lab-made sensors, refer to `examples/configs/gsmini.yaml` to create a sensor specification, which will be used as the `CONFIG_PATH` in the command below. Below are step-by-step instructions for sensor calibration. **Note: The current version does not support gel pads with markers.**
+### Calibration Data Collection
+| Ball Indenter | Collecting Data | Collected Image |
+|---------|---------|---------|
+|
|
|
|
+
+To collect calibration data, use a ball indenter of known diameter to press against the sensor. Examples of the setup and resulting images are shown above. Run the following command and allocate a location `CALIB_DIR` to save the calibration data:
+```bash
+collect_data [-b CALIB_DIR] [-d BALL_DIAMETER_IN_MM] [-c CONFIG_PATH]
+```
+* Instruction:
+ * Save background image: Press 'b'
+ * Capture ~50 tactile images by pressing the ball in various locations: Press 'w'
+ * Exit: Press 'q'
+* Tips for Optimal Calibration:
+ * Ball Size: Select a ball that appears well-sized within the sensor’s view, like the tactile image shown above; 4mm to 9mm is suitable for GelSight Mini.
+ * Pressure: Avoid pressing too hard.
+ * Coverage: Capture at least 50 images with the ball contacting different regions of the sensor.
+ * Using Multiple Balls: Use the same `CALIB_DIR` and specify distinct `BALL_DIAMETER` values if balls in different size are applied.
+
+### Label Collected Data
+| NanoGui Screenshot |
+|---------|
+|
|
+
+Run the command below to label the contact circle on the collected tactile data using NanoGUI:
+
+```bash
+label_data [-b CALIB_DIR] [-c CONFIG_PATH]
+```
+* Instruction:
+ * Click the **Open** icon to begin.
+ * Keyboard Controls for aligning the label with the contact circle:
+ * **Arrow keys (left/right/up/down)**: Adjust the circle's position.
+ * **'m' / 'p'**: Decrease / increase the circle's radius.
+ * **'f' / 'c'**: Decrease / increase the circle's movement step.
+ * Once aligned, click the **Calibrate** icon.
+ * After labeling all data, close the window to exit.
+
+### Prepare Dataset
+Run the command below to prepare the dataset for calibration model training:
+```bash
+prepare_data [-b CALIB_DIR] [-c CONFIG_PATH]
+```
+
+### Train Calibration Model
+Train the MLP model to map pixel color and location (RGBXY) to surface gradients for each pixel. Use the following command to train the model with the collected dataset:
+```bash
+train_model [-b CALIB_DIR] [-d {cpu|cuda}]
+```
+
+The trained model is saved in `CALIB_DIR/model/nnmodel.pth`.
+
+## Test the Trained Calibration Model
+Once the model is trained, connect the sensor and run the following command to stream images and perform real-time surface reconstruction using the trained calibration model:
+
+```bash
+test_model [-b CALIB_DIR] [-c CONFIG_PATH]
+```
+After starting, wait briefly for background data collection; real-time surface gradient predictions will then be displayed. Press any key to exit.
+
+
+### References
+1. S. Wang, Y. She, B. Romero, and E. H. Adelson, “Gelsight wedge:
+Measuring high-resolution 3d contact geometry with a compact robot
+finger,” in 2021 IEEE International Conference on Robotics and
+Automation (ICRA). IEEE, 2021.
\ No newline at end of file
diff --git a/calibration/__init__.py b/calibration/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/calibration/collect_data.py b/calibration/collect_data.py
new file mode 100644
index 0000000..46ae7df
--- /dev/null
+++ b/calibration/collect_data.py
@@ -0,0 +1,148 @@
+import argparse
+import os
+
+import cv2
+import numpy as np
+import yaml
+
+from gs_sdk.gs_device import Camera
+from calibration.utils import load_csv_as_dict
+
+"""
+This script collects tactile data using ball indenters for sensor calibration.
+
+Instruction:
+ 1. Connect the sensor to the computer.
+ 2. Prepare a ball indenter with known diameter.
+ 3. Runs this script, press 'b' to collect a background image.
+ 4. Press the sensor with the ball indenter at multiple locations (~50 locations preferred),
+ press 'w' to save the tactile image. When done, press 'q' to quit.
+Note:
+ If you have prepared multiple balls in different diameters, you can run this script multiple
+ times, assign the same calib_dir but different ball diameters, the system will treat it as
+ one single dataset.
+
+Usage:
+ python collect_data.py --calib_dir CALIB_DIR --ball_diameter DIAMETER [--config_path CONFIG_PATH]
+
+Arguments:
+ --calib_dir: Path to the directory where the collected data will be saved
+ --ball_diameter: Diameter of the ball indenter in mm
+ --config_path: (Optional) Path to the configuration file about the sensor dimensions.
+ If not provided, GelSight Mini is assumed.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "../examples/configs")
+
+
+def collect_data():
+ # Argument Parsers
+ parser = argparse.ArgumentParser(
+ description="Collect calibration data with ball indenters to calibrate the sensor."
+ )
+ parser.add_argument(
+ "-b",
+ "--calib_dir",
+ type=str,
+ help="path to save calibration data",
+ )
+ parser.add_argument(
+ "-d", "--ball_diameter", type=float, help="diameter of the indenter in mm"
+ )
+ parser.add_argument(
+ "-c",
+ "--config_path",
+ type=str,
+ help="path of the sensor information",
+ default=os.path.join(config_dir, "gsmini.yaml"),
+ )
+ args = parser.parse_args()
+
+ # Create the data saving directories
+ calib_dir = args.calib_dir
+ ball_diameter = args.ball_diameter
+ indenter_subdir = "%.3fmm" % (ball_diameter)
+ indenter_dir = os.path.join(calib_dir, indenter_subdir)
+ if not os.path.isdir(indenter_dir):
+ os.makedirs(indenter_dir)
+
+ # Read the configuration
+ config_path = args.config_path
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ device_name = config["device_name"]
+ imgh = config["imgh"]
+ imgw = config["imgw"]
+
+ # Create the data saving catalog
+ catalog_path = os.path.join(calib_dir, "catalog.csv")
+ if not os.path.isfile(catalog_path):
+ with open(catalog_path, "w") as f:
+ f.write("experiment_reldir,diameter(mm)\n")
+
+ # Find last data_count collected with this diameter
+ data_dict = load_csv_as_dict(catalog_path)
+ diameters = np.array([float(diameter) for diameter in data_dict["diameter(mm)"]])
+ data_idxs = np.where(np.abs(diameters - ball_diameter) < 1e-3)[0]
+ data_counts = np.array(
+ [int(os.path.basename(reldir)) for reldir in data_dict["experiment_reldir"]]
+ )
+ if len(data_idxs) == 0:
+ data_count = 0
+ else:
+ data_count = max(data_counts[data_idxs]) + 1
+
+ # Connect to the device and collect data until quit
+ device = Camera(device_name, imgh, imgw)
+ device.connect()
+ print("Press key to collect data, collect background, or quit (w/b/q)")
+ while True:
+ image = device.get_image()
+
+ # Display the image and decide record or quit
+ cv2.imshow("frame", image)
+ key = cv2.waitKey(100)
+ if key == ord("w"):
+ # Save the image
+ experiment_reldir = os.path.join(indenter_subdir, str(data_count))
+ experiment_dir = os.path.join(calib_dir, experiment_reldir)
+ if not os.path.isdir(experiment_dir):
+ os.makedirs(experiment_dir)
+ save_path = os.path.join(experiment_dir, "gelsight.png")
+ cv2.imwrite(save_path, image)
+ print("Save data to new path: %s" % save_path)
+
+ # Save to catalog
+ with open(catalog_path, "a") as f:
+ f.write(experiment_reldir + "," + str(ball_diameter))
+ f.write("\n")
+ data_count += 1
+ elif key == ord("b"):
+ print("Collecting 10 background images, please wait ...")
+ images = []
+ for _ in range(10):
+ image = device.get_image()
+ images.append(image)
+ cv2.imshow("frame", image)
+ cv2.waitKey(1)
+ image = np.mean(images, axis=0).astype(np.uint8)
+ # Save the background image
+ save_path = os.path.join(calib_dir, "background.png")
+ cv2.imwrite(save_path, image)
+ print("Save background image to %s" % save_path)
+ elif key == ord("q"):
+ # Quit
+ break
+ elif key == -1:
+ # No key pressed
+ continue
+ else:
+ print("Unrecognized key %s" % key)
+
+ device.release()
+ cv2.destroyAllWindows()
+ print("%d images collected in total." % data_count)
+
+
+if __name__ == "__main__":
+ collect_data()
diff --git a/calibration/label_data.py b/calibration/label_data.py
new file mode 100644
index 0000000..eb88c67
--- /dev/null
+++ b/calibration/label_data.py
@@ -0,0 +1,314 @@
+import gc
+import os
+import argparse
+
+import cv2
+import numpy as np
+import nanogui as ng
+from nanogui import Texture
+from nanogui import glfw
+import yaml
+
+from calibration.utils import load_csv_as_dict
+
+"""
+Rewrite from Zilin Si's code: https://github.com/Robo-Touch/Taxim
+
+This script is for manually labeling the contact circle in the tactile image for sensor calibration.
+The labeled data will save the center and radius of the circle in the tactile image.
+
+Prerequisite:
+ - Tactile images collected using ball indenters with known diameters are collected.
+Instruction:
+ 1. Runs the script,
+ 2. Mouse-press 'Open' to select the directory where the collected data are saved.
+ 3. Press left/right/up/down to control the circle's location,
+ Press m/p to decrease/increase the circle's radius,
+ Press f/c to decrease/increase the circle's moving step.
+ Mouse-press 'Calibrate' to save the labeled data.
+ 4. Repeat step 3 for all the tactile images and close Nanogui when done.
+
+Usage:
+ python label_data.py --calib_dir CALIB_DIR [--config_path CONFIG_PATH] [--display_difference] [--detect_circle]
+
+Arguments:
+ --calib_dir: Path to the directory where the collected data are saved
+ --config_path: (Optional) Path to the configuration file about the sensor dimensions.
+ If not provided, GelSight Mini is assumed.
+ --display_difference: (Store True) Display the difference between the background image.
+ --detect_circle: (Store True) Automatically detect the circle in the image.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "../examples/configs")
+
+
+class Circle:
+ """the circle drawed on the tactile image to get the contact size"""
+
+ color_circle = (128, 0, 0)
+ opacity = 0.5
+
+ def __init__(self, x, y, radius=25, increments=2):
+ self.center = [x, y]
+ self.radius = radius
+ self.increments = increments
+
+
+class CalibrateApp(ng.Screen):
+ fnames = list()
+ read_all = False # flag to indicate if all images have been read
+ load_img = True
+ change = False
+
+ def __init__(
+ self, calib_data, imgw, imgh, display_difference=False, detect_circle=False
+ ):
+ super(CalibrateApp, self).__init__((1024, 768), "Gelsight Calibration App")
+ self.imgw = imgw
+ self.imgh = imgh
+ self.display_difference = display_difference
+ self.detect_circle = detect_circle
+ # Load background
+ self.bg_img = cv2.imread(os.path.join(calib_data, "background.png"))
+ # Initialize the circle
+ self.circle = Circle(self.imgw / 2, self.imgh / 2, radius=40)
+
+ window = ng.Window(self, "IO Window")
+ window.set_position((15, 15))
+ window.set_layout(ng.GroupLayout())
+
+ ng.Label(window, "Folder dialog", "sans-bold")
+ tools = ng.Widget(window)
+ tools.set_layout(
+ ng.BoxLayout(ng.Orientation.Horizontal, ng.Alignment.Middle, 0, 6)
+ )
+
+ # Initialize the file directory and list of filenames
+ b = ng.Button(tools, "Open")
+
+ def open_cb():
+ self.parent_dir = calib_data
+ # Read the catalog and create a list of all filenames
+ catalog_dict = load_csv_as_dict(
+ os.path.join(self.parent_dir, "catalog.csv")
+ )
+ self.fnames = [
+ os.path.join(self.parent_dir, fname)
+ for fname in catalog_dict["experiment_reldir"]
+ ]
+ self.circle_radius = [
+ float(radius) for radius in catalog_dict["diameter(mm)"]
+ ]
+ print(
+ f"Selected directory = {self.parent_dir}, total {len(self.fnames)} images"
+ )
+ self.img_idx = 0
+
+ b.set_callback(open_cb)
+
+ # Initialize the image window
+ self.img_window = ng.Window(self, "Current image")
+ self.img_window.set_position((200, 15))
+ self.img_window.set_layout(ng.GroupLayout())
+
+ # Initialize the calibrate button
+ b = ng.Button(self.img_window, "Calibrate")
+
+ def calibrate_cb():
+ frame = self.orig_img
+ center = self.circle.center
+ radius = self.circle.radius
+ print(f"Frame {self.img_idx}: center = {center}, radius = is {radius}")
+ # save the data for each individual frame instead of creating a long list and save it at the end
+ # save the radius and center to npz file
+ save_dir = os.path.join(self.fnames[self.img_idx], "label.npz")
+ np.savez(save_dir, center=center, radius=radius)
+ # save the labeled image
+ labeled_img = self.overlay_circle(frame, self.circle)
+ labeled_img_path = os.path.join(self.fnames[self.img_idx], "labeled.png")
+ cv2.imwrite(labeled_img_path, labeled_img)
+
+ # Update img index
+ self.load_img = True
+ self.update_img_idx()
+
+ b.set_callback(calibrate_cb)
+
+ ###########
+ self.img_view = ng.ImageView(self.img_window)
+ self.img_tex = ng.Texture(
+ pixel_format=Texture.PixelFormat.RGB,
+ component_format=Texture.ComponentFormat.UInt8,
+ size=[imgw, imgh],
+ min_interpolation_mode=Texture.InterpolationMode.Trilinear,
+ mag_interpolation_mode=Texture.InterpolationMode.Nearest,
+ flags=Texture.TextureFlags.ShaderRead | Texture.TextureFlags.RenderTarget,
+ )
+ self.perform_layout()
+
+ def update_img_idx(self):
+ self.img_idx += 1
+ if self.img_idx == len(self.fnames) - 1:
+ self.read_all = True
+
+ def overlay_circle(self, orig_img, circle):
+ center = circle.center
+ radius = circle.radius
+ color_circle = circle.color_circle
+ opacity = circle.opacity
+
+ overlay = orig_img.copy()
+ center_tuple = (int(center[0]), int(center[1]))
+ cv2.circle(overlay, center_tuple, radius, color_circle, -1)
+ cv2.addWeighted(overlay, opacity, orig_img, 1 - opacity, 0, overlay)
+ return overlay
+
+ def draw(self, ctx):
+ self.img_window.set_size((2000, 2600))
+ self.img_view.set_size((self.imgw, self.imgh))
+
+ # load a new image
+ if self.load_img and len(self.fnames) > 0 and not self.read_all:
+ print("Loading %s" % self.fnames[self.img_idx])
+
+ # Load img
+ self.orig_img = cv2.imread(
+ os.path.join(self.fnames[self.img_idx], "gelsight.png")
+ )
+ # Initialize the circle pose
+ if self.detect_circle:
+ diff_image = self.orig_img.astype(np.float32) - self.bg_img.astype(
+ np.float32
+ )
+ color_mask = np.linalg.norm(diff_image, axis=-1) > 15
+ color_mask = cv2.dilate(
+ color_mask.astype(np.uint8), np.ones((7, 7), np.uint8)
+ )
+ color_mask = cv2.erode(
+ color_mask.astype(np.uint8), np.ones((15, 15), np.uint8)
+ )
+ contours, _ = cv2.findContours(
+ color_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+ )
+ largest_contour = max(contours, key=cv2.contourArea)
+ M = cv2.moments(largest_contour)
+ cx = int(M["m10"] / M["m00"])
+ cy = int(M["m01"] / M["m00"])
+ else:
+ cx = self.circle.center[0]
+ cy = self.circle.center[1]
+ radius = max(self.circle.radius - 13, 5)
+ self.circle = Circle(cx, cy, radius=radius)
+
+ # Add circle and add img to viewer
+ if (self.load_img and len(self.fnames) > 0) or self.change:
+ self.load_img = False
+ self.change = False
+ # Add circle
+ if self.display_difference:
+ diff_img = (
+ self.orig_img.astype(np.float32) - self.bg_img.astype(np.float32)
+ ) * 3
+ diff_img = np.clip(diff_img, -127, 128) + np.ones_like(diff_img) * 127
+ display_img = cv2.cvtColor(diff_img.astype(np.uint8), cv2.COLOR_BGR2RGB)
+ else:
+ display_img = cv2.cvtColor(self.orig_img, cv2.COLOR_BGR2RGB)
+ img = self.overlay_circle(display_img, self.circle)
+
+ if self.img_tex.channels() > 3:
+ height, width = img.shape[:2]
+ alpha = 255 * np.ones((height, width, 1), dtype=img.dtype)
+ img = np.concatenate((img, alpha), axis=2)
+
+ # Add to img view
+ self.img_tex.upload(img)
+ self.img_view.set_image(self.img_tex)
+
+ super(CalibrateApp, self).draw(ctx)
+
+ def keyboard_event(self, key, scancode, action, modifiers):
+ if super(CalibrateApp, self).keyboard_event(key, scancode, action, modifiers):
+ return True
+ if key == glfw.KEY_ESCAPE and action == glfw.PRESS:
+ self.set_visible(False)
+ return True
+ elif key == glfw.KEY_C:
+ self.circle.increments *= 2
+ elif key == glfw.KEY_F:
+ self.circle.increments /= 2
+ else:
+ self.change = True
+ if key == glfw.KEY_LEFT:
+ self.circle.center[0] -= self.circle.increments
+ elif key == glfw.KEY_RIGHT:
+ self.circle.center[0] += self.circle.increments
+ elif key == glfw.KEY_UP:
+ self.circle.center[1] -= self.circle.increments
+ elif key == glfw.KEY_DOWN:
+ self.circle.center[1] += self.circle.increments
+ elif key == glfw.KEY_M:
+ self.circle.radius -= 1
+ elif key == glfw.KEY_P:
+ self.circle.radius += 1
+
+ return False
+
+
+def label_data():
+ # Argument parser
+ parser = argparse.ArgumentParser(
+ description="Label the ball indenter data using Nanogui."
+ )
+ parser.add_argument(
+ "-b",
+ "--calib_dir",
+ type=str,
+ help="path to save calibration data",
+ )
+ parser.add_argument(
+ "-c",
+ "--config_path",
+ type=str,
+ help="path of configuring gelsight",
+ default=os.path.join(config_dir, "gsmini.yaml"),
+ )
+ parser.add_argument(
+ "-d",
+ "--display_difference",
+ action="store_true",
+ help="Display the difference between the background image",
+ )
+ parser.add_argument(
+ "-r",
+ "--detect_circle",
+ action="store_true",
+ help="Automatically detect the circle in the image",
+ )
+ args = parser.parse_args()
+
+ # Read the configuration
+ config_path = args.config_path
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ imgh = config["imgh"]
+ imgw = config["imgw"]
+
+ # Start the label process
+ ng.init()
+ app = CalibrateApp(
+ args.calib_dir,
+ imgw,
+ imgh,
+ display_difference=args.display_difference,
+ detect_circle=args.detect_circle,
+ )
+ app.draw_all()
+ app.set_visible(True)
+ ng.mainloop(refresh=1 / 60.0 * 1000)
+ del app
+ gc.collect()
+ ng.shutdown()
+
+
+if __name__ == "__main__":
+ label_data()
diff --git a/calibration/models.py b/calibration/models.py
new file mode 100644
index 0000000..cda5356
--- /dev/null
+++ b/calibration/models.py
@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+
+class BGRXYDataset(Dataset):
+ """The BGRXY Datast."""
+
+ def __init__(self, bgrxys, gxyangles):
+ self.bgrxys = bgrxys
+ self.gxyangles = gxyangles
+
+ def __len__(self):
+ return len(self.bgrxys)
+
+ def __getitem__(self, index):
+ bgrxy = torch.tensor(self.bgrxys[index], dtype=torch.float32)
+ gxyangles = torch.tensor(self.gxyangles[index], dtype=torch.float32)
+ return bgrxy, gxyangles
+
+
+class BGRXYMLPNet_(nn.Module):
+ """
+ The architecture using MLP, this is never used in test time.
+ We train with this architecture and then transfer weights to the 1-by-1 convolution architecture.
+ """
+
+ def __init__(self):
+ super(BGRXYMLPNet_, self).__init__()
+ input_size = 5
+ self.fc1 = nn.Linear(input_size, 128)
+ self.bn1 = nn.BatchNorm1d(128)
+ self.fc2 = nn.Linear(128, 32)
+ self.bn2 = nn.BatchNorm1d(32)
+ self.fc3 = nn.Linear(32, 32)
+ self.bn3 = nn.BatchNorm1d(32)
+ self.fc4 = nn.Linear(32, 2)
+
+ def forward(self, x):
+ x = F.relu(self.bn1(self.fc1(x)))
+ x = F.relu(self.bn2(self.fc2(x)))
+ x = F.relu(self.bn3(self.fc3(x)))
+ x = self.fc4(x)
+ return x
diff --git a/calibration/prepare_data.py b/calibration/prepare_data.py
new file mode 100644
index 0000000..b20c1a1
--- /dev/null
+++ b/calibration/prepare_data.py
@@ -0,0 +1,132 @@
+import argparse
+import json
+import os
+
+import cv2
+import numpy as np
+import yaml
+
+from calibration.utils import load_csv_as_dict
+from gs_sdk.gs_reconstruct import image2bgrxys
+
+"""
+This script prepares dataset for the tactile sensor calibration.
+It is based on the collected and labeled data.
+
+Prerequisite:
+ - Tactile images collected using ball indenters with known diameters are collected.
+ - Collected tactile images are labeled.
+
+Usage:
+ python prepare_data.py --calib_dir CALIB_DIR [--config_path CONFIG_PATH] [--radius_reduction RADIUS_REDUCTION]
+
+Arguments:
+ --calib_dir: Path to the directory where the collected data will be saved
+ --config_path: (Optional) Path to the configuration file about the sensor dimensions.
+ If not provided, GelSight Mini is assumed.
+ --radius_reduction: (Optional) Reduce the radius of the labeled circle. This helps guarantee all labeled pixels are indented.
+ If not provided, 4 pixels will be reduced.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "../examples/configs")
+
+
+def prepare_data():
+ # Argument Parsers
+ parser = argparse.ArgumentParser(
+ description="Use the labeled collected data to prepare the dataset files (npz)."
+ )
+ parser.add_argument(
+ "-b",
+ "--calib_dir",
+ type=str,
+ help="path of the calibration data",
+ )
+ parser.add_argument(
+ "-c",
+ "--config_path",
+ type=str,
+ help="path of configuring gelsight",
+ default=os.path.join(config_dir, "gsmini.yaml"),
+ )
+ parser.add_argument(
+ "-r",
+ "--radius_reduction",
+ type=float,
+ help="reduce the radius of the labeled circle. When not considering shadows, this helps guarantee all labeled pixels are indented. ",
+ default=4.0,
+ )
+ args = parser.parse_args()
+
+ # Load the data_dict
+ calib_dir = args.calib_dir
+ catalog_path = os.path.join(calib_dir, "catalog.csv")
+ data_dict = load_csv_as_dict(catalog_path)
+ diameters = np.array([float(diameter) for diameter in data_dict["diameter(mm)"]])
+ experiment_reldirs = np.array(data_dict["experiment_reldir"])
+
+ # Split data into train and test and save the split information
+ perm = np.random.permutation(len(experiment_reldirs))
+ n_train = 4 * len(experiment_reldirs) // 5
+ data_path = os.path.join(calib_dir, "train_test_split.json")
+ dict_to_save = {
+ "train": experiment_reldirs[perm[:n_train]].tolist(),
+ "test": experiment_reldirs[perm[n_train:]].tolist(),
+ }
+ with open(data_path, "w") as f:
+ json.dump(dict_to_save, f, indent=4)
+
+ # Read the configuration
+ config_path = args.config_path
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ ppmm = config["ppmm"]
+
+ # Extract the pixel data from each tactile image and calculate the gradients
+ for experiment_reldir, diameter in zip(experiment_reldirs, diameters):
+ experiment_dir = os.path.join(calib_dir, experiment_reldir)
+ image_path = os.path.join(experiment_dir, "gelsight.png")
+ image = cv2.imread(image_path)
+
+ # Filter the non-indented pixels
+ label_path = os.path.join(experiment_dir, "label.npz")
+ label_data = np.load(label_path)
+ center = label_data["center"]
+ radius = label_data["radius"] - args.radius_reduction
+ xys = np.dstack(
+ np.meshgrid(
+ np.arange(image.shape[1]), np.arange(image.shape[0]), indexing="xy"
+ )
+ )
+ dists = np.linalg.norm(xys - center, axis=2)
+ mask = dists < radius
+
+ # Find the gradient angles, prepare the data, and save the data
+ ball_radius = diameter / ppmm / 2.0
+ if ball_radius < radius:
+ print(experiment_reldir)
+ print("Press too deep, deeper than the ball radius")
+ continue
+ dxys = xys - center
+ dists[np.logical_not(mask)] = 0.0
+ dzs = np.sqrt(ball_radius**2 - np.square(dists))
+ gxangles = np.arctan2(dxys[:, :, 0], dzs)
+ gyangles = np.arctan2(dxys[:, :, 1], dzs)
+ gxyangles = np.stack([gxangles, gyangles], axis=-1)
+ gxyangles[np.logical_not(mask)] = np.array([0.0, 0.0])
+ bgrxys = image2bgrxys(image)
+ save_path = os.path.join(experiment_dir, "data.npz")
+ np.savez(save_path, bgrxys=bgrxys, gxyangles=gxyangles, mask=mask)
+
+ # Save the background data
+ bg_path = os.path.join(calib_dir, "background.png")
+ bg_image = cv2.imread(bg_path)
+ bgrxys = image2bgrxys(bg_image)
+ gxyangles = np.zeros((bg_image.shape[0], bg_image.shape[1], 2))
+ mask = np.ones((bg_image.shape[0], bg_image.shape[1]), dtype=np.bool_)
+ save_path = os.path.join(calib_dir, "background_data.npz")
+ np.savez(save_path, bgrxys=bgrxys, gxyangles=gxyangles, mask=mask)
+
+
+if __name__ == "__main__":
+ prepare_data()
diff --git a/calibration/test_model.py b/calibration/test_model.py
new file mode 100644
index 0000000..06659f6
--- /dev/null
+++ b/calibration/test_model.py
@@ -0,0 +1,97 @@
+import argparse
+import os
+
+import cv2
+import numpy as np
+import yaml
+
+from gs_sdk.gs_device import Camera
+from gs_sdk.gs_reconstruct import Reconstructor
+
+"""
+This script tests the calibrated model real-time reconstructing local patches with the sensor.
+
+Prerequisite:
+ - Collect data and train the calibration model.
+Instructions:
+ - Connect the sensor to the computer.
+ - Run this script, wait a bit for background collection, press any key to quit the streaming session.
+
+Usage:
+ python test.py --calib_dir CALIB_DIR [--config_path CONFIG_PATH]
+
+Arguments:
+ --calib_dir: Path to the directory where the collected data is stored.
+ --config_path: (Optional) Path to the configuration file about the sensor dimensions.
+ If not provided, GelSight Mini is assumed.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "../examples/configs")
+
+
+def test_model():
+ # Argument Parsers
+ parser = argparse.ArgumentParser(
+ description="Read image from the device and reconstruct based on the calibrated model."
+ )
+ parser.add_argument(
+ "-b",
+ "--calib_dir",
+ type=str,
+ help="place where the calibration data is stored",
+ )
+ parser.add_argument(
+ "-c",
+ "--config_path",
+ type=str,
+ help="path of the sensor information",
+ default=os.path.join(config_dir, "gsmini.yaml"),
+ )
+ args = parser.parse_args()
+
+ # Load the device configuration
+ with open(args.config_path, "r") as f:
+ config = yaml.safe_load(f)
+ device_name = config["device_name"]
+ imgh = config["imgh"]
+ imgw = config["imgw"]
+ ppmm = config["ppmm"]
+
+ # Create device and the reconstructor
+ device = Camera(device_name, imgh, imgw)
+ device.connect()
+ model_path = os.path.join(args.calib_dir, "model", "nnmodel.pth")
+ recon = Reconstructor(model_path, device="cpu")
+
+ # Collect background images
+ print("Collecting 10 background images, please wait ...")
+ bg_images = []
+ for _ in range(10):
+ image = device.get_image()
+ bg_images.append(image)
+ bg_image = np.mean(bg_images, axis=0).astype(np.uint8)
+ recon.load_bg(bg_image)
+
+ # Real-time reconstruct
+ print("\nPrss any key to quit.\n")
+ while True:
+ image = device.get_image()
+ G, H, C = recon.get_surface_info(image, ppmm)
+ # Create the image for gradient visualization
+ red = G[:, :, 0] * 255 / 3.0 + 127
+ red = np.clip(red, 0, 255)
+ blue = G[:, :, 1] * 255 / 3.0 + 127
+ blue = np.clip(blue, 0, 255)
+ grad_image = np.stack((blue, np.zeros_like(blue), red), axis=-1).astype(np.uint8)
+ # Display
+ cv2.imshow(device_name, grad_image)
+ key = cv2.waitKey(1)
+ if key != -1:
+ break
+
+ device.release()
+ cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+ test_model()
diff --git a/calibration/train_model.py b/calibration/train_model.py
new file mode 100644
index 0000000..b8959c5
--- /dev/null
+++ b/calibration/train_model.py
@@ -0,0 +1,195 @@
+import argparse
+import json
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from calibration.utils import load_csv_as_dict, transfer_weights
+from calibration.models import BGRXYDataset, BGRXYMLPNet_
+from gs_sdk.gs_reconstruct import BGRXYMLPNet
+
+"""
+This script trains the gradient prediction network.
+The network is trained as MLP taking the pixel BGRXY as input and predict the gradients gx, gy.
+
+Prerequisite:
+ - Tactile images collected using ball indenters with known diameters are collected.
+ - Collected tactile images are labeled.
+ - Labeled data are prepared into dataset.
+
+Usage:
+ python train.py --calib_dir CALIB_DIR [--n_epochs N_EPOCHS] [--lr LR] [--device {cpu, cuda}]
+
+Arguments:
+ --calib_dir: Path to the directory where the collected data will be saved
+ --n_epochs: (Optional) Number of training epochs. Default is 200.
+ --lr: (Optional) Learning rate. Default is 0.002.
+ --device: (Optional) The device to train the network. Can choose between cpu and cuda. Default is cpu.
+"""
+
+
+def train_model():
+ # Argument Parsers
+ parser = argparse.ArgumentParser(description="Train the model from BGRXY to gxy.")
+ parser.add_argument(
+ "-b",
+ "--calib_dir",
+ type=str,
+ help="place where the calibration data is stored",
+ )
+ parser.add_argument(
+ "-ne", "--n_epochs", type=int, default=200, help="number of training epochs"
+ )
+ parser.add_argument("-lr", "--lr", type=float, default=0.002, help="learning rate")
+ parser.add_argument(
+ "-d",
+ "--device",
+ type=str,
+ choices=["cpu", "cuda"],
+ default="cpu",
+ help="the device to train NN",
+ )
+ args = parser.parse_args()
+
+ # Create the model directory
+ calib_dir = args.calib_dir
+ model_dir = os.path.join(calib_dir, "model")
+ if not os.path.isdir(model_dir):
+ os.makedirs(model_dir)
+
+ # Load the train and test split
+ data_path = os.path.join(calib_dir, "train_test_split.json")
+ with open(data_path, "r") as f:
+ data = json.load(f)
+ train_reldirs = data["train"]
+ test_reldirs = data["test"]
+
+ # Load the train and test data including the background data
+ train_data = {"all_bgrxys": [], "all_gxyangles": []}
+ for experiment_reldir in train_reldirs:
+ data_path = os.path.join(calib_dir, experiment_reldir, "data.npz")
+ if not os.path.isfile(data_path):
+ raise ValueError("Data file %s does not exist" % data_path)
+ data = np.load(data_path)
+ train_data["all_bgrxys"].append(data["bgrxys"][data["mask"]])
+ train_data["all_gxyangles"].append(data["gxyangles"][data["mask"]])
+ test_data = {"all_bgrxys": [], "all_gxyangles": []}
+ for experiment_reldir in test_reldirs:
+ data_path = os.path.join(calib_dir, experiment_reldir, "data.npz")
+ if not os.path.isfile(data_path):
+ raise ValueError("Data file %s does not exist" % data_path)
+ data = np.load(data_path)
+ test_data["all_bgrxys"].append(data["bgrxys"][data["mask"]])
+ test_data["all_gxyangles"].append(data["gxyangles"][data["mask"]])
+ #Load background data
+ bg_path = os.path.join(calib_dir, "background_data.npz")
+ bg_data = np.load(bg_path)
+ bgrxys = bg_data["bgrxys"][bg_data["mask"]]
+ gxyangles = bg_data["gxyangles"][bg_data["mask"]]
+ perm = np.random.permutation(len(bgrxys))
+ n_train = np.sum([len(bgrxys) for bgrxys in train_data["all_bgrxys"]]) // 5
+ n_test = np.sum([len(bgrxys) for bgrxys in test_data["all_bgrxys"]]) // 5
+ if n_train + n_test > len(bgrxys):
+ n_train = 4 * len(bgrxys) // 5
+ n_test = len(bgrxys) // 5
+ train_data["all_bgrxys"].append(bgrxys[perm[:n_train]])
+ train_data["all_gxyangles"].append(gxyangles[perm[:n_train]])
+ test_data["all_bgrxys"].append(bgrxys[perm[n_train : n_train + n_test]])
+ test_data["all_gxyangles"].append(gxyangles[perm[n_train : n_train + n_test]])
+ # Construct the train and test dataset
+ train_bgrxys = np.concatenate(train_data["all_bgrxys"])
+ train_gxyangles = np.concatenate(train_data["all_gxyangles"])
+ test_bgrxys = np.concatenate(test_data["all_bgrxys"])
+ test_gxyangles = np.concatenate(test_data["all_gxyangles"])
+
+ # Create train and test Dataloader
+ train_dataset = BGRXYDataset(train_bgrxys, train_gxyangles)
+ train_dataloader = DataLoader(train_dataset, batch_size=1024, shuffle=True)
+ test_dataset = BGRXYDataset(test_bgrxys, test_gxyangles)
+ test_dataloader = DataLoader(test_dataset, batch_size=1024, shuffle=False)
+
+ # Create the MLP Net for training
+ device = args.device
+ net = BGRXYMLPNet_().to(device)
+ criterion = nn.L1Loss()
+ optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=0.0)
+ scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
+
+ # Create MLP Net in the CNN format for saving
+ save_net = BGRXYMLPNet().to(device)
+
+ # Initial evaluation
+ train_mae = evaluate(net, train_dataloader, device)
+ test_mae = evaluate(net, test_dataloader, device)
+ naive_mae = np.mean(np.abs(test_gxyangles - np.mean(train_gxyangles, axis=0)))
+ traj = {"train_maes": [train_mae], "test_maes": [test_mae], "naive_mae": naive_mae}
+ print("Naive MAE (predict as mean): %.4f" % naive_mae)
+ print("without train, Train MAE: %.4f, Test MAE: %.4f" % (train_mae, test_mae))
+
+ # Train the model
+ for epoch_idx in range(args.n_epochs):
+ losses = []
+ net.train()
+ for bgrxys, gxyangles in train_dataloader:
+ bgrxys = bgrxys.to(device)
+ gxyangles = gxyangles.to(device)
+ optimizer.zero_grad()
+ outputs = net(bgrxys)
+ loss = criterion(outputs, gxyangles)
+ loss.backward()
+ optimizer.step()
+ diffs = outputs - gxyangles
+ losses.append(np.abs(diffs.cpu().detach().numpy()))
+ net.eval()
+ traj["train_maes"].append(np.mean(np.concatenate(losses)))
+ traj["test_maes"].append(evaluate(net, test_dataloader, device))
+ print(
+ "Epoch %i, Train MAE: %.4f, Test MAE: %.4f"
+ % (epoch_idx, traj["train_maes"][-1], traj["test_maes"][-1])
+ )
+ scheduler.step()
+
+ # Save model every 10 steps
+ if (epoch_idx + 1) % 10 == 0:
+ # Transfer weights to MLP Net and save
+ transfer_weights(net, save_net)
+ save_path = os.path.join(model_dir, "nnmodel.pth")
+ torch.save(save_net.state_dict(), save_path)
+
+ # Save the training curve
+ save_path = os.path.join(model_dir, "training_curve.png")
+ plt.plot(np.arange(len(traj["train_maes"])), traj["train_maes"], color="blue")
+ plt.plot(np.arange(len(traj["test_maes"])), traj["test_maes"], color="red")
+ plt.xlabel("Epochs")
+ plt.ylabel("MAE (rad)")
+ plt.title("MAE Curve")
+ plt.savefig(save_path)
+ plt.close()
+
+
+def evaluate(net, dataloader, device):
+ """
+ Evaluate the network loss on the dataset.
+
+ :param net: nn.Module; the network to evaluate.
+ :param dataloader: DataLoader; the dataloader for the dataset.
+ :param device: str; the device to evaluate the network.
+ """
+ losses = []
+ for bgrxys, gxyangles in dataloader:
+ bgrxys = bgrxys.to(device)
+ gxyangles = gxyangles.to(device)
+ outputs = net(bgrxys)
+ diffs = outputs - gxyangles
+ losses.append(np.abs(diffs.cpu().detach().numpy()))
+ mae = np.mean(np.concatenate(losses))
+ return mae
+
+
+if __name__ == "__main__":
+ train_model()
diff --git a/calibration/utils.py b/calibration/utils.py
new file mode 100644
index 0000000..ca0078b
--- /dev/null
+++ b/calibration/utils.py
@@ -0,0 +1,64 @@
+import csv
+import numpy as np
+
+
+def load_csv_as_dict(csv_path):
+ """
+ Load the csv file entries as dictionaries.
+
+ :params csv_path: str; the path of the csv file.
+ :returns: dict; the dictionary of the csv file.
+ """
+ with open(csv_path, "r") as f:
+ reader = csv.DictReader(f, delimiter=",")
+ data = list(reader)
+ keys = reader.fieldnames
+ data_dict = {}
+ for key in keys:
+ data_dict[key] = []
+ for line in data:
+ for key in keys:
+ data_dict[key].append(line[key])
+ return data_dict
+
+
+def transfer_weights(mlp_model, fcn_model):
+ """
+ transfer weights between BGRXYMLPNet_ to BGRXYMLPNet.
+
+ :param mlp_model: BGRXYMLPNet_; the model to transfer from.
+ :param fcn_model: BGRXYMLPNet; the model to transfer to.
+ """
+ # Copy weights from fc1 to conv1
+ fcn_model.conv1.weight.data = mlp_model.fc1.weight.data.view(
+ fcn_model.conv1.weight.size()
+ )
+ fcn_model.conv1.bias.data = mlp_model.fc1.bias.data
+ fcn_model.bn1.weight.data = mlp_model.bn1.weight.data
+ fcn_model.bn1.bias.data = mlp_model.bn1.bias.data
+ fcn_model.bn1.running_mean = mlp_model.bn1.running_mean
+ fcn_model.bn1.running_var = mlp_model.bn1.running_var
+ # Copy weights from fc2 to conv2
+ fcn_model.conv2.weight.data = mlp_model.fc2.weight.data.view(
+ fcn_model.conv2.weight.size()
+ )
+ fcn_model.conv2.bias.data = mlp_model.fc2.bias.data
+ fcn_model.bn2.weight.data = mlp_model.bn2.weight.data
+ fcn_model.bn2.bias.data = mlp_model.bn2.bias.data
+ fcn_model.bn2.running_mean = mlp_model.bn2.running_mean
+ fcn_model.bn2.running_var = mlp_model.bn2.running_var
+ # Copy weights from fc3 to conv3
+ fcn_model.conv3.weight.data = mlp_model.fc3.weight.data.view(
+ fcn_model.conv3.weight.size()
+ )
+ fcn_model.conv3.bias.data = mlp_model.fc3.bias.data
+ fcn_model.bn3.weight.data = mlp_model.bn3.weight.data
+ fcn_model.bn3.bias.data = mlp_model.bn3.bias.data
+ fcn_model.bn3.running_mean = mlp_model.bn3.running_mean
+ fcn_model.bn3.running_var = mlp_model.bn3.running_var
+ # Copy weights from fc4 to conv4
+ fcn_model.conv4.weight.data = mlp_model.fc4.weight.data.view(
+ fcn_model.conv4.weight.size()
+ )
+ fcn_model.conv4.bias.data = mlp_model.fc4.bias.data
+ return fcn_model
diff --git a/examples/configs/digit.yaml b/examples/configs/digit.yaml
new file mode 100644
index 0000000..29f3e88
--- /dev/null
+++ b/examples/configs/digit.yaml
@@ -0,0 +1,11 @@
+# Device Name
+device_name: "DIGIT"
+# Pixel Per Millimeter
+ppmm: 0.0405
+# Desired Image Width and Height
+imgh: 240
+imgw: 320
+# Raw image width, height, and framerate
+raw_imgh: 480
+raw_imgw: 640
+framerate: 60
\ No newline at end of file
diff --git a/examples/configs/gsmini.yaml b/examples/configs/gsmini.yaml
new file mode 100644
index 0000000..0238a21
--- /dev/null
+++ b/examples/configs/gsmini.yaml
@@ -0,0 +1,11 @@
+# Device Name
+device_name: "GelSight Mini"
+# Pixel Per Millimeter
+ppmm: 0.0634
+# Desired Image Width and Height
+imgh: 240
+imgw: 320
+# Raw image width, height, and framerate
+raw_imgh: 2464
+raw_imgw: 3280
+framerate: 25
diff --git a/examples/data/background.png b/examples/data/background.png
new file mode 100644
index 0000000..68506df
Binary files /dev/null and b/examples/data/background.png differ
diff --git a/examples/data/bead.png b/examples/data/bead.png
new file mode 100644
index 0000000..67510da
Binary files /dev/null and b/examples/data/bead.png differ
diff --git a/examples/data/key.png b/examples/data/key.png
new file mode 100644
index 0000000..4a577ca
Binary files /dev/null and b/examples/data/key.png differ
diff --git a/examples/data/seed.png b/examples/data/seed.png
new file mode 100644
index 0000000..9f3ccb0
Binary files /dev/null and b/examples/data/seed.png differ
diff --git a/examples/fast_stream_device.py b/examples/fast_stream_device.py
new file mode 100644
index 0000000..80547aa
--- /dev/null
+++ b/examples/fast_stream_device.py
@@ -0,0 +1,50 @@
+import os
+
+import cv2
+import yaml
+
+from gs_sdk.gs_device import FastCamera
+
+"""
+This script demonstrates how to use the FastCamera class from the gs_sdk package.
+
+It loads a configuration file, initializes the FastCamera, and streaming images with low latency.
+This script is only for GelSight Mini so far as only GelSight Mini has the frame dropping issue.
+
+Usage:
+ python fast_stream_device.py
+
+Press any key to quit the streaming session.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "configs")
+
+
+def fast_stream_device():
+ # Load the device configuration
+ config_path = os.path.join(config_dir, "gsmini.yaml")
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ device_name = config["device_name"]
+ imgh = config["imgh"]
+ imgw = config["imgw"]
+ raw_imgh = config["raw_imgh"]
+ raw_imgw = config["raw_imgw"]
+ framerate = config["framerate"]
+
+ # Create device and stream the device
+ device = FastCamera(device_name, imgh, imgw, raw_imgh, raw_imgw, framerate)
+ device.connect()
+ print("\nPrss any key to quit.\n")
+ while True:
+ image = device.get_image()
+ cv2.imshow(device_name, image)
+ key = cv2.waitKey(1)
+ if key != -1:
+ break
+ device.release()
+ cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+ fast_stream_device()
diff --git a/examples/models/gsmini.pth b/examples/models/gsmini.pth
new file mode 100644
index 0000000..422a215
Binary files /dev/null and b/examples/models/gsmini.pth differ
diff --git a/examples/reconstruct.py b/examples/reconstruct.py
new file mode 100644
index 0000000..6b927b0
--- /dev/null
+++ b/examples/reconstruct.py
@@ -0,0 +1,78 @@
+import argparse
+import os
+
+import cv2
+import matplotlib.pyplot as plt
+import yaml
+
+from gs_sdk.gs_reconstruct import Reconstructor
+from gs_sdk.viz_utils import plot_gradients
+
+"""
+This script demonstrates how to use the Reconstructor class from the gs_sdk package.
+
+It loads a configuration file, initialize the Reconstructor, reconstruct surface information from images,
+and save them to files in the "data/" directory.
+
+Usage:
+ python reconstruct.py --device {cuda, cpu}
+
+Arguments:
+ --device: The device to load the neural network model. Options are 'cuda' or 'cpu'.
+"""
+
+model_path = os.path.join(os.path.dirname(__file__), "models", "gsmini.pth")
+config_path = os.path.join(os.path.dirname(__file__), "configs", "gsmini.yaml")
+data_dir = os.path.join(os.path.dirname(__file__), "data")
+
+
+def reconstruct():
+ # Argument Parser
+ parser = argparse.ArgumentParser(description="Reconstruct surface info from data.")
+ parser.add_argument(
+ "-d",
+ "--device",
+ type=str,
+ choices=["cuda", "cpu"],
+ default="cpu",
+ help="The device to load and run the neural network model.",
+ )
+ args = parser.parse_args()
+
+ # Load the device configuration
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ ppmm = config["ppmm"]
+
+ # Create reconstructor
+ recon = Reconstructor(model_path, device=args.device)
+ bg_image = cv2.imread(os.path.join(data_dir, "background.png"))
+ recon.load_bg(bg_image)
+
+ # Reconstruct the surface information from data and save them to files
+ filenames = ["bead.png", "key.png", "seed.png"]
+ for filename in filenames:
+ image = cv2.imread(os.path.join(data_dir, filename))
+ G, H, C = recon.get_surface_info(image, ppmm)
+
+ # Plot the surface information
+ fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(15, 10))
+ axes[0, 0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+ axes[0, 0].set_title("GelSight Image")
+ plot_gradients(fig, axes[0, 1], G[:, :, 0], G[:, :, 1], mask=C, mode="rgb")
+ axes[0, 1].set_title("Reconstructed Gradients")
+ axes[1, 0].imshow(H, cmap="jet")
+ axes[1, 0].set_title("Reconstructed Heights")
+ axes[1, 1].imshow(C)
+ axes[1, 1].set_title("Predicted Contact Mask")
+ for ax in axes.flatten():
+ ax.set_xticks([])
+ ax.set_yticks([])
+ save_path = os.path.join(data_dir, "reconstructed_" + filename)
+ plt.savefig(save_path)
+ plt.close()
+ print("Save results to %s" % save_path)
+
+
+if __name__ == "__main__":
+ reconstruct()
diff --git a/examples/stream_device.py b/examples/stream_device.py
new file mode 100644
index 0000000..d3b64a4
--- /dev/null
+++ b/examples/stream_device.py
@@ -0,0 +1,69 @@
+import argparse
+import os
+
+import cv2
+import yaml
+
+from gs_sdk.gs_device import Camera
+
+"""
+This script demonstrates how to use the Camera class from the gs_sdk package.
+
+It loads a configuration file, initializes the Camera, and streaming images.
+
+Usage:
+ python stream_device.py --device_name {gsmini, digit}
+
+Arguments:
+ --device_name: The name of the device to stream from. Options are 'gsmini' or 'digit'.
+ Default is 'gsmini'.
+
+Press any key to quit the streaming session.
+"""
+
+config_dir = os.path.join(os.path.dirname(__file__), "configs")
+
+
+def stream_device():
+ # Argument parser
+ parser = argparse.ArgumentParser(description="Read and show image from the device.")
+ parser.add_argument(
+ "-n",
+ "--device_name",
+ type=str,
+ choices=["gsmini", "digit"],
+ default="gsmini",
+ help="The name of the device",
+ )
+ args = parser.parse_args()
+
+ # Load the device configuration
+ if args.device_name == "gsmini":
+ config_file = "gsmini.yaml"
+ elif args.device_name == "digit":
+ config_file = "digit.yaml"
+ else:
+ raise ValueError("Unknown device name %s." % args.device_name)
+ config_path = os.path.join(config_dir, config_file)
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ device_name = config["device_name"]
+ imgh = config["imgh"]
+ imgw = config["imgw"]
+
+ # Create device and stream the device
+ device = Camera(device_name, imgh, imgw)
+ device.connect()
+ print("\nPrss any key to quit.\n")
+ while True:
+ image = device.get_image()
+ cv2.imshow(device_name, image)
+ key = cv2.waitKey(1)
+ if key != -1:
+ break
+ device.release()
+ cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+ stream_device()
diff --git a/gs_sdk/__init__.py b/gs_sdk/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/gs_sdk/gs_device.py b/gs_sdk/gs_device.py
new file mode 100644
index 0000000..6a64152
--- /dev/null
+++ b/gs_sdk/gs_device.py
@@ -0,0 +1,206 @@
+import os
+import re
+import subprocess
+
+import cv2
+import ffmpeg
+import numpy as np
+
+
+class Camera:
+ """
+ The GelSight Camera Class.
+
+ This class handles camera initialization, image acquisition, and camera release.
+ Some sensors (GelSight Mini) might experience frame dropping issues, use FastCamera class instead.
+ """
+
+ def __init__(self, dev_type, imgh, imgw):
+ """
+ Initialize the camera.
+
+ :param dev_type: str; The type of the camera.
+ :param imgh: int; The height of the image.
+ :param imgw: int; The width of the image.
+ """
+ self.dev_type = dev_type
+ self.dev_id = get_camera_id(self.dev_type)
+ self.imgh = imgh
+ self.imgw = imgw
+ self.cam = None
+ self.data = None
+
+ def connect(self):
+ """
+ Connect to the camera using cv2 streamer.
+ """
+ self.cam = cv2.VideoCapture(self.dev_id)
+ self.cam.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+ if self.cam is None or not self.cam.isOpened():
+ print("Warning: unable to open video source %d" % (self.dev_id))
+ else:
+ print("Connect to %s at video source %d" % (self.dev_type, self.dev_id))
+
+ def get_image(self, flush=False):
+ """
+ Get the image from the camera.
+
+ :param flush: bool; Whether to flush the first few frames.
+ :return: np.ndarray; The image from the camera.
+ """
+ if flush:
+ # flush out fist few frames to remove black frames
+ for i in range(10):
+ ret, f0 = self.cam.read()
+ ret, f0 = self.cam.read()
+ if ret:
+ f0 = resize_crop(f0, self.imgw, self.imgh)
+ self.data = f0
+ else:
+ print("ERROR! reading image from video source %d" % (self.dev_id))
+ return self.data
+
+ def release(self):
+ """
+ Release the camera resource.
+ """
+ if self.cam is not None:
+ self.cam.release()
+ print("Video source %d released." % (self.dev_id))
+ else:
+ print("No camera to release.")
+
+
+class FastCamera:
+ """
+ The GelSight Camera Class with low latency.
+
+ This class handles camera initialization, image acquisition, and camera release with low latency.
+ """
+
+ def __init__(self, dev_type, imgh, imgw, raw_imgh, raw_imgw, framerate):
+ """
+ Initialize the low latency camera. Raw camera parameters are required to stream with low latency.
+
+ :param dev_type: str; The type of the camera.
+ :param imgh: int; The desired height of the image.
+ :param imgw: int; The desired width of the image.
+ :param raw_imgh: int; The raw height of the image.
+ :param raw_imgw: int; The raw width of the image.
+ :param framerate: int; The frame rate of the camera.
+ """
+ # Raw image size
+ self.raw_imgh = raw_imgh
+ self.raw_imgw = raw_imgw
+ self.raw_size = self.raw_imgh * self.raw_imgw * 3
+ self.framerate = framerate
+ # desired image size
+ self.imgh = imgh
+ self.imgw = imgw
+ # Get camera ID
+ self.dev_type = dev_type
+ self.dev_id = get_camera_id(self.dev_type)
+ self.device = "/dev/video" + str(self.dev_id)
+
+ def connect(self):
+ """
+ Connect to the camera using FFMpeg streamer.
+ """
+ # Command to capture video using ffmpeg and high resolution
+ self.ffmpeg_command = (
+ ffmpeg.input(
+ self.device,
+ format="v4l2",
+ framerate=self.framerate,
+ video_size="%dx%d" % (self.raw_imgw, self.raw_imgh),
+ )
+ .output("pipe:", format="rawvideo", pix_fmt="bgr24")
+ .global_args("-fflags", "nobuffer")
+ .global_args("-flags", "low_delay")
+ .global_args("-fflags", "+genpts")
+ .global_args("-rtbufsize", "0")
+ .compile()
+ )
+ self.process = subprocess.Popen(
+ self.ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+ )
+ # Warm-up phase: discard the first few frames
+ print("Warming up the camera...")
+ warm_up_frames = 100
+ for _ in range(warm_up_frames):
+ self.process.stdout.read(self.raw_size)
+ print("Camera ready for use!")
+
+ def get_image(self):
+ """
+ Get the image from the camera from raw data stream.
+
+ :return: np.ndarray; The image from the camera.
+ """
+ raw_frame = self.process.stdout.read(self.raw_size)
+ frame = np.frombuffer(raw_frame, np.uint8).reshape(
+ (self.raw_imgh, self.raw_imgw, 3)
+ )
+ frame = resize_crop(frame, self.imgw, self.imgh)
+ return frame
+
+ def release(self):
+ """
+ Release the camera resource.
+ """
+ self.process.stdout.close()
+ self.process.wait()
+
+
+def get_camera_id(camera_name):
+ """
+ Find the camera ID that has the corresponding camera name.
+
+ :param camera_name: str; The name of the camera.
+ :return: int; The camera ID.
+ """
+ cam_num = None
+ for file in os.listdir("/sys/class/video4linux"):
+ real_file = os.path.realpath("/sys/class/video4linux/" + file + "/name")
+ with open(real_file, "rt") as name_file:
+ name = name_file.read().rstrip()
+ if camera_name in name:
+ cam_num = int(re.search("\d+$", file).group(0))
+ found = "FOUND!"
+ else:
+ found = " "
+ print("{} {} -> {}".format(found, file, name))
+
+ return cam_num
+
+
+def resize_crop(img, imgw, imgh):
+ """
+ Resize and crop the image to the desired size.
+
+ :param img: np.ndarray; The image to resize and crop.
+ :param imgw: int; The width of the desired image.
+ :param imgh: int; The height of the desired image.
+ :return: np.ndarray; The resized and cropped image.
+ """
+ # remove 1/7th of border from each size
+ border_size_x, border_size_y = int(img.shape[0] * (1 / 7)), int(
+ np.floor(img.shape[1] * (1 / 7))
+ )
+ cropped_imgh = img.shape[0] - 2 * border_size_x
+ cropped_imgw = img.shape[1] - 2 * border_size_y
+ # Extra cropping to maintain aspect ratio
+ extra_border_h = 0
+ extra_border_w = 0
+ if cropped_imgh * imgw / imgh > cropped_imgw + 1e-8:
+ extra_border_h = int(cropped_imgh - cropped_imgw * imgh / imgw)
+ elif cropped_imgh * imgw / imgh < cropped_imgw - 1e-8:
+ extra_border_w = int(cropped_imgw - cropped_imgh * imgw / imgh)
+ # keep the ratio the same as the original image size
+ img = img[
+ border_size_x + extra_border_h : img.shape[0] - border_size_x,
+ border_size_y + extra_border_w : img.shape[1] - border_size_y,
+ ]
+ # final resize for the desired image size
+ img = cv2.resize(img, (imgw, imgh))
+ return img
diff --git a/gs_sdk/gs_reconstruct.py b/gs_sdk/gs_reconstruct.py
new file mode 100644
index 0000000..3f91f03
--- /dev/null
+++ b/gs_sdk/gs_reconstruct.py
@@ -0,0 +1,214 @@
+import math
+import os
+
+import cv2
+import numpy as np
+from scipy import fftpack
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BGRXYMLPNet(nn.Module):
+ """
+ The Neural Network architecture for GelSight calibration.
+
+ This class uses 1-by-1 convolution, which is technically the same as using MLP.
+ """
+
+ def __init__(self):
+ super(BGRXYMLPNet, self).__init__()
+ input_channels = 5
+ self.conv1 = nn.Conv2d(input_channels, 128, kernel_size=1)
+ self.bn1 = nn.BatchNorm2d(128)
+ self.conv2 = nn.Conv2d(128, 32, kernel_size=1)
+ self.bn2 = nn.BatchNorm2d(32)
+ self.conv3 = nn.Conv2d(32, 32, kernel_size=1)
+ self.bn3 = nn.BatchNorm2d(32)
+ self.conv4 = nn.Conv2d(32, 2, kernel_size=1)
+
+ def forward(self, x):
+ x = F.relu(self.bn1(self.conv1(x)))
+ x = F.relu(self.bn2(self.conv2(x)))
+ x = F.relu(self.bn3(self.conv3(x)))
+ x = self.conv4(x)
+ return x
+
+
+class Reconstructor:
+ """
+ The GelSight reconstruction class.
+
+ This class handles 3D reconstruction from calibrated GelSight images.
+ """
+
+ def __init__(self, model_path, contact_mode="standard", device="cpu"):
+ """
+ Initialize the reconstruction model.
+ Contact mode "flat" means the object in contact is flat, so a different threshold
+ is used to determine contact mask.
+
+ :param model_path: str; the path of the calibrated neural network model.
+ :param contact_mode: str {"standard", "flat"}; the mode to get the contact mask.
+ :param device: str {"cuda", "cpu"}; the device to run the model.
+ """
+ self.model_path = model_path
+ self.contact_mode = contact_mode
+ self.device = device
+ self.bg_image = None
+ # Load the gxy model
+ if not os.path.isfile(model_path):
+ raise ValueError("Error opening %s, file does not exist" % model_path)
+ self.gxy_net = BGRXYMLPNet()
+ self.gxy_net.load_state_dict(torch.load(model_path), self.device)
+ self.gxy_net.eval()
+
+ def load_bg(self, bg_image):
+ """
+ Load the background image.
+
+ :param bg_image: np.array (H, W, 3); the background image.
+ """
+ self.bg_image = bg_image
+
+ # Calculate the gradients of the background
+ bgrxys = image2bgrxys(bg_image)
+ bgrxys = bgrxys.transpose(2, 0, 1)
+ features = torch.from_numpy(bgrxys[np.newaxis, :, :, :]).float().to(self.device)
+ with torch.no_grad():
+ gxyangles = self.gxy_net(features)
+ gxyangles = gxyangles[0].cpu().detach().numpy()
+ self.bg_G = np.tan(gxyangles.transpose(1, 2, 0))
+
+ def get_surface_info(self, image, ppmm):
+ """
+ Get the surface information including gradients (G), height map (H), and contact mask (C).
+
+ :param image: np.array (H, W, 3); the gelsight image.
+ :param ppmm: float; the pixel per mm.
+ :return G: np.array (H, W, 2); the gradients.
+ H: np.array (H, W); the height map.
+ C: np.array (H, W); the contact mask.
+ """
+ # Calculate the gradients
+ bgrxys = image2bgrxys(image)
+ bgrxys = bgrxys.transpose(2, 0, 1)
+ features = torch.from_numpy(bgrxys[np.newaxis, :, :, :]).float().to(self.device)
+ with torch.no_grad():
+ gxyangles = self.gxy_net(features)
+ gxyangles = gxyangles[0].cpu().detach().numpy()
+ G = np.tan(gxyangles.transpose(1, 2, 0))
+ if self.bg_image is not None:
+ G = G - self.bg_G
+ else:
+ raise ValueError("Background image is not loaded.")
+
+ # Calculate the height map
+ H = poisson_dct_neumaan(G[:, :, 0], G[:, :, 1]).astype(np.float32)
+
+ # Calculate the contact mask
+ if self.contact_mode == "standard":
+ # Find the contact mask based on color difference
+ diff_image = image.astype(np.float32) - self.bg_image.astype(np.float32)
+ color_mask = np.linalg.norm(diff_image, axis=-1) > 15
+ color_mask = cv2.dilate(
+ color_mask.astype(np.uint8), np.ones((7, 7), np.uint8)
+ )
+ color_mask = cv2.erode(
+ color_mask.astype(np.uint8), np.ones((15, 15), np.uint8)
+ )
+
+ # Filter by height
+ cutoff = np.percentile(H, 85) - 0.2 / ppmm
+ height_mask = H < cutoff
+ C = np.logical_and(color_mask, height_mask)
+ elif self.contact_mode == "flat":
+ # Find the contact mask based on color difference
+ diff_image = image.astype(np.float32) - self.bg_image.astype(np.float32)
+ color_mask = np.linalg.norm(diff_image, axis=-1) > 10
+ color_mask = cv2.dilate(
+ color_mask.astype(np.uint8), np.ones((15, 15), np.uint8)
+ )
+ C = cv2.erode(
+ color_mask.astype(np.uint8), np.ones((25, 25), np.uint8)
+ ).astype(np.bool_)
+
+ return G, H, C
+
+
+def image2bgrxys(image):
+ """
+ Convert a bgr image to bgrxy feature.
+
+ :param image: np.array (H, W, 3); the bgr image.
+ :return: np.array (H, W, 5); the bgrxy feature.
+ """
+ xys = np.dstack(
+ np.meshgrid(np.arange(image.shape[1]), np.arange(image.shape[0]), indexing="xy")
+ )
+ xys = xys.astype(np.float32) / np.array([image.shape[1], image.shape[0]])
+ bgrs = image.copy() / 255
+ bgrxys = np.concatenate([bgrs, xys], axis=2)
+ return bgrxys
+
+
+def poisson_dct_neumaan(gx, gy):
+ """
+ 2D integration of depth from gx, gy using Poisson solver.
+
+ :param gx: np.array (H, W); the x gradient.
+ :param gy: np.array (H, W); the y gradient.
+ :return: np.array (H, W); the depth map.
+ """
+ # Compute Laplacian
+ gxx = 1 * (
+ gx[:, (list(range(1, gx.shape[1])) + [gx.shape[1] - 1])]
+ - gx[:, ([0] + list(range(gx.shape[1] - 1)))]
+ )
+ gyy = 1 * (
+ gy[(list(range(1, gx.shape[0])) + [gx.shape[0] - 1]), :]
+ - gy[([0] + list(range(gx.shape[0] - 1))), :]
+ )
+ f = gxx + gyy
+
+ # Right hand side of the boundary condition
+ b = np.zeros(gx.shape)
+ b[0, 1:-2] = -gy[0, 1:-2]
+ b[-1, 1:-2] = gy[-1, 1:-2]
+ b[1:-2, 0] = -gx[1:-2, 0]
+ b[1:-2, -1] = gx[1:-2, -1]
+ b[0, 0] = (1 / np.sqrt(2)) * (-gy[0, 0] - gx[0, 0])
+ b[0, -1] = (1 / np.sqrt(2)) * (-gy[0, -1] + gx[0, -1])
+ b[-1, -1] = (1 / np.sqrt(2)) * (gy[-1, -1] + gx[-1, -1])
+ b[-1, 0] = (1 / np.sqrt(2)) * (gy[-1, 0] - gx[-1, 0])
+
+ # Modification near the boundaries to enforce the non-homogeneous Neumann BC (Eq. 53 in [1])
+ f[0, 1:-2] = f[0, 1:-2] - b[0, 1:-2]
+ f[-1, 1:-2] = f[-1, 1:-2] - b[-1, 1:-2]
+ f[1:-2, 0] = f[1:-2, 0] - b[1:-2, 0]
+ f[1:-2, -1] = f[1:-2, -1] - b[1:-2, -1]
+
+ # Modification near the corners (Eq. 54 in [1])
+ f[0, -1] = f[0, -1] - np.sqrt(2) * b[0, -1]
+ f[-1, -1] = f[-1, -1] - np.sqrt(2) * b[-1, -1]
+ f[-1, 0] = f[-1, 0] - np.sqrt(2) * b[-1, 0]
+ f[0, 0] = f[0, 0] - np.sqrt(2) * b[0, 0]
+
+ # Cosine transform of f
+ tt = fftpack.dct(f, norm="ortho")
+ fcos = fftpack.dct(tt.T, norm="ortho").T
+
+ # Cosine transform of z (Eq. 55 in [1])
+ (x, y) = np.meshgrid(range(1, f.shape[1] + 1), range(1, f.shape[0] + 1), copy=True)
+ denom = 4 * (
+ (np.sin(0.5 * math.pi * x / (f.shape[1]))) ** 2
+ + (np.sin(0.5 * math.pi * y / (f.shape[0]))) ** 2
+ )
+
+ # Inverse Discrete cosine Transform
+ f = -fcos / denom
+ tt = fftpack.idct(f, norm="ortho")
+ img_tt = fftpack.idct(tt.T, norm="ortho").T
+ img_tt = img_tt.mean() + img_tt
+
+ return img_tt
diff --git a/gs_sdk/viz_utils.py b/gs_sdk/viz_utils.py
new file mode 100644
index 0000000..95d04de
--- /dev/null
+++ b/gs_sdk/viz_utils.py
@@ -0,0 +1,44 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_gradients(fig, ax, gx, gy, mask=None, mode="rgb", **kwargs):
+ """
+ Plot the gradients.
+
+ :params fig: plt.figure; the figure to plot the gradients.
+ :params ax: plt.axis; the axis to plot the gradients.
+ :params gx: np.array (H, W); the x gradient.
+ :params gy: np.array (H, W); the y gradient.
+ :params mask: np.array (H, W); the mask for gradients to be plotted
+ :params mode: str {"rgb", "quiver"}; the mode to plot the gradients.
+ """
+ if mode == "rgb":
+ # Plot the gradient in red and blue
+ grad_range = kwargs.get("grad_range", 3.0)
+ red = gx * 255 / grad_range + 127
+ red = np.clip(red, 0, 255)
+ blue = gy * 255 / grad_range + 127
+ blue = np.clip(blue, 0, 255)
+ image = np.stack((red, np.zeros_like(red), blue), axis=-1).astype(np.uint8)
+ if mask is not None:
+ image[np.logical_not(mask)] = np.array([127, 0, 127])
+ ax.imshow(image)
+ elif mode == "quiver":
+ # Plot the gradient in quiver
+ n_skip = kwargs.get("n_skip", 5)
+ quiver_scale = kwargs.get("quiver_scale", 10.0)
+ imgh, imgw = gx.shape
+ X, Y = np.meshgrid(np.arange(imgw)[::n_skip], np.arange(imgh)[::n_skip])
+ U = gx[::n_skip, ::n_skip] * quiver_scale
+ V = -gy[::n_skip, ::n_skip] * quiver_scale
+ if mask is None:
+ mask = np.ones_like(gx)
+ else:
+ mask = np.copy(mask)
+ mask = mask[::n_skip, ::n_skip]
+ ax.quiver(X[mask], Y[mask], U[mask], V[mask], units="xy", scale=1, color="red")
+ ax.set_xlim(0, imgw)
+ ax.set_ylim(imgh, 0)
+ else:
+ raise ValueError("Unknown plot gradient mode %s" % mode)
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..2fc1fc0
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,34 @@
+from setuptools import setup, find_packages
+
+setup(
+ name="gs_sdk",
+ version="0.1.0",
+ description="SDK for GelSight sensors usage, reconstruction, and calibration.",
+ author="Hung-Jui Huang, Ruihan Gao",
+ author_email="hungjuih@andrew.cmu.edu, ruihang@andrew.cmu.edu",
+ packages=find_packages(),
+ install_requires=[
+ "pillow==10.0.0",
+ "numpy==1.26.4",
+ "opencv-python>=4.9.0",
+ "scipy>=1.13.1",
+ "torch>=2.1.0",
+ "PyYaml>=6.0.1",
+ "matplotlib>=3.9.0",
+ "ffmpeg-python",
+ "nanogui"
+ ],
+ python_requires=">=3.9",
+ entry_points={
+ 'console_scripts': [
+ 'collect_data=calibration.collect_data:collect_data',
+ 'label_data=calibration.label_data:label_data',
+ 'prepare_data=calibration.prepare_data:prepare_data',
+ 'train_model=calibration.train_model:train_model',
+ 'test_model=calibration.test_model:test_model',
+ ],
+ },
+ classifiers=[
+ "Programming Language :: Python :: 3",
+ ],
+)