commit 5118d07b0f2cac30276fdefd14a582194294c12d Author: zhangmozhe Date: Mon Nov 30 22:57:10 2020 +0800 init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0b77750 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +*~ +*.pth \ No newline at end of file diff --git a/.style.yapf b/.style.yapf new file mode 100644 index 0000000..40632b4 --- /dev/null +++ b/.style.yapf @@ -0,0 +1,4 @@ +[style] +based_on_style: google +column_limit: 160 +split_arguments_when_comma_terminated: true diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..45f085a --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.pythonPath": "/home/zhanbo/anaconda3/envs/ColorVid/bin/python" +} \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..dd6534f --- /dev/null +++ b/README.md @@ -0,0 +1,67 @@ +# Deep Exemplar-based Video Colorization (Pytorch Implementation) + + + +### [Paper](https://arxiv.org/abs/1906.09909) | [Pretrained Model]() | [Youtube video](https://youtu.be/HXWR5h5vVYI) + +**Deep Exemplar-based Video Colorization, CVPR2019** + +[Bo Zhang](https://www.microsoft.com/en-us/research/people/zhanbo/)1,3, +[Mingming He](http://mingminghe.com/)1,5, +[Jing Liao](https://liaojing.github.io/html/)2, +[Pedro V. Sander](https://www.cse.ust.hk/~psander/)1, +[Lu Yuan](https://www.microsoft.com/en-us/research/people/luyuan/)4, +[Amine Bermak](https://eebermak.home.ece.ust.hk/)1, +[Dong Chen](https://www.microsoft.com/en-us/research/people/doch/)3
+1Hong Kong University of Science and Technology,2City University of Hong Kong, +3Microsoft Research Asia, 4Microsoft Cloud&AI, 5USC Institute for Creative Technologies + +## Prerequisites + +- Python 3.6+ +- Nvidia GPU + CUDA, CuDNN + +## Installation + +```bash +conda create -n ColorVid python=3.6 +source activate ColorVid +pip install -r requirements.txt +``` + +## Data preparation + +checkpoint, pre-trained vgg + +- input video frames +- reference image +- output + +## Test + +```bash +python test.py +``` + +## TODO + +- [ ] Release the training code +- [ ] Prepare the Colab Demo +- [ ] Output the video +- [ ] GIF teaser + +## Citation + +If you use this code for your research, please cite our paper. + +> @inproceedings{zhang2019deep, +> title={Deep exemplar-based video colorization}, +> author={Zhang, Bo and He, Mingming and Liao, Jing and Sander, Pedro V and Yuan, Lu and Bermak, Amine and Chen, Dong}, +> booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, +> pages={8052--8061}, +> year={2019} +> } + +## License + +This project is licensed under the MIT license. diff --git a/assets/teaser.png b/assets/teaser.png new file mode 100644 index 0000000..5ce212b Binary files /dev/null and b/assets/teaser.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1415.png b/legacy_videos/legacy_clips/v32_180/1415.png new file mode 100644 index 0000000..487122b Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1415.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1416.png b/legacy_videos/legacy_clips/v32_180/1416.png new file mode 100644 index 0000000..be81b41 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1416.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1417.png b/legacy_videos/legacy_clips/v32_180/1417.png new file mode 100644 index 0000000..c784cb8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1417.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1418.png b/legacy_videos/legacy_clips/v32_180/1418.png new file mode 100644 index 0000000..ba0d846 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1418.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1419.png b/legacy_videos/legacy_clips/v32_180/1419.png new file mode 100644 index 0000000..a72190f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1419.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1420.png b/legacy_videos/legacy_clips/v32_180/1420.png new file mode 100644 index 0000000..6c64f24 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1420.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1421.png b/legacy_videos/legacy_clips/v32_180/1421.png new file mode 100644 index 0000000..b6d8c42 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1421.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1422.png b/legacy_videos/legacy_clips/v32_180/1422.png new file mode 100644 index 0000000..355f99f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1422.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1423.png b/legacy_videos/legacy_clips/v32_180/1423.png new file mode 100644 index 0000000..dc76663 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1423.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1424.png b/legacy_videos/legacy_clips/v32_180/1424.png new file mode 100644 index 0000000..0638a00 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1424.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1425.png b/legacy_videos/legacy_clips/v32_180/1425.png new file mode 100644 index 0000000..c989bc5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1425.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1426.png b/legacy_videos/legacy_clips/v32_180/1426.png new file mode 100644 index 0000000..a45c0eb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1426.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1427.png b/legacy_videos/legacy_clips/v32_180/1427.png new file mode 100644 index 0000000..d5ff057 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1427.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1428.png b/legacy_videos/legacy_clips/v32_180/1428.png new file mode 100644 index 0000000..592b5db Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1428.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1429.png b/legacy_videos/legacy_clips/v32_180/1429.png new file mode 100644 index 0000000..689c86f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1429.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1430.png b/legacy_videos/legacy_clips/v32_180/1430.png new file mode 100644 index 0000000..4fa54f3 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1430.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1431.png b/legacy_videos/legacy_clips/v32_180/1431.png new file mode 100644 index 0000000..950ad8c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1431.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1432.png b/legacy_videos/legacy_clips/v32_180/1432.png new file mode 100644 index 0000000..b74b4a8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1432.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1433.png b/legacy_videos/legacy_clips/v32_180/1433.png new file mode 100644 index 0000000..ac71b27 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1433.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1434.png b/legacy_videos/legacy_clips/v32_180/1434.png new file mode 100644 index 0000000..d1224bc Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1434.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1435.png b/legacy_videos/legacy_clips/v32_180/1435.png new file mode 100644 index 0000000..d92c9c3 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1435.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1436.png b/legacy_videos/legacy_clips/v32_180/1436.png new file mode 100644 index 0000000..a45760d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1436.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1437.png b/legacy_videos/legacy_clips/v32_180/1437.png new file mode 100644 index 0000000..393b47b Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1437.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1438.png b/legacy_videos/legacy_clips/v32_180/1438.png new file mode 100644 index 0000000..2a2c7e0 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1438.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1439.png b/legacy_videos/legacy_clips/v32_180/1439.png new file mode 100644 index 0000000..938049c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1439.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1440.png b/legacy_videos/legacy_clips/v32_180/1440.png new file mode 100644 index 0000000..1bc8a41 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1440.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1441.png b/legacy_videos/legacy_clips/v32_180/1441.png new file mode 100644 index 0000000..141827d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1441.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1442.png b/legacy_videos/legacy_clips/v32_180/1442.png new file mode 100644 index 0000000..f51624d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1442.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1443.png b/legacy_videos/legacy_clips/v32_180/1443.png new file mode 100644 index 0000000..37b6747 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1443.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1444.png b/legacy_videos/legacy_clips/v32_180/1444.png new file mode 100644 index 0000000..afd80fc Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1444.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1445.png b/legacy_videos/legacy_clips/v32_180/1445.png new file mode 100644 index 0000000..2dba673 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1445.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1446.png b/legacy_videos/legacy_clips/v32_180/1446.png new file mode 100644 index 0000000..55360ff Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1446.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1447.png b/legacy_videos/legacy_clips/v32_180/1447.png new file mode 100644 index 0000000..8c3bf61 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1447.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1448.png b/legacy_videos/legacy_clips/v32_180/1448.png new file mode 100644 index 0000000..9813ffa Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1448.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1449.png b/legacy_videos/legacy_clips/v32_180/1449.png new file mode 100644 index 0000000..cd82b29 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1449.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1450.png b/legacy_videos/legacy_clips/v32_180/1450.png new file mode 100644 index 0000000..0b2f3d9 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1450.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1451.png b/legacy_videos/legacy_clips/v32_180/1451.png new file mode 100644 index 0000000..ddc7d99 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1451.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1452.png b/legacy_videos/legacy_clips/v32_180/1452.png new file mode 100644 index 0000000..9a5577f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1452.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1453.png b/legacy_videos/legacy_clips/v32_180/1453.png new file mode 100644 index 0000000..70e7af7 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1453.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1454.png b/legacy_videos/legacy_clips/v32_180/1454.png new file mode 100644 index 0000000..874375e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1454.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1455.png b/legacy_videos/legacy_clips/v32_180/1455.png new file mode 100644 index 0000000..6a6c4bf Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1455.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1456.png b/legacy_videos/legacy_clips/v32_180/1456.png new file mode 100644 index 0000000..cc8adb1 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1456.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1457.png b/legacy_videos/legacy_clips/v32_180/1457.png new file mode 100644 index 0000000..e084043 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1457.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1458.png b/legacy_videos/legacy_clips/v32_180/1458.png new file mode 100644 index 0000000..e0e8416 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1458.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1459.png b/legacy_videos/legacy_clips/v32_180/1459.png new file mode 100644 index 0000000..43d93c7 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1459.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1460.png b/legacy_videos/legacy_clips/v32_180/1460.png new file mode 100644 index 0000000..4c3a29d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1460.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1461.png b/legacy_videos/legacy_clips/v32_180/1461.png new file mode 100644 index 0000000..6f24968 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1461.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1462.png b/legacy_videos/legacy_clips/v32_180/1462.png new file mode 100644 index 0000000..558a7f1 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1462.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1463.png b/legacy_videos/legacy_clips/v32_180/1463.png new file mode 100644 index 0000000..17822d1 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1463.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1464.png b/legacy_videos/legacy_clips/v32_180/1464.png new file mode 100644 index 0000000..04e0eeb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1464.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1465.png b/legacy_videos/legacy_clips/v32_180/1465.png new file mode 100644 index 0000000..23fa1e2 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1465.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1466.png b/legacy_videos/legacy_clips/v32_180/1466.png new file mode 100644 index 0000000..6dfbf3a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1466.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1467.png b/legacy_videos/legacy_clips/v32_180/1467.png new file mode 100644 index 0000000..81dd5d8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1467.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1468.png b/legacy_videos/legacy_clips/v32_180/1468.png new file mode 100644 index 0000000..1d9514a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1468.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1469.png b/legacy_videos/legacy_clips/v32_180/1469.png new file mode 100644 index 0000000..48447eb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1469.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1470.png b/legacy_videos/legacy_clips/v32_180/1470.png new file mode 100644 index 0000000..742b023 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1470.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1471.png b/legacy_videos/legacy_clips/v32_180/1471.png new file mode 100644 index 0000000..e8b7c84 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1471.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1472.png b/legacy_videos/legacy_clips/v32_180/1472.png new file mode 100644 index 0000000..0fed450 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1472.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1473.png b/legacy_videos/legacy_clips/v32_180/1473.png new file mode 100644 index 0000000..f79eb13 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1473.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1474.png b/legacy_videos/legacy_clips/v32_180/1474.png new file mode 100644 index 0000000..51ffc3c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1474.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1475.png b/legacy_videos/legacy_clips/v32_180/1475.png new file mode 100644 index 0000000..8760505 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1475.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1476.png b/legacy_videos/legacy_clips/v32_180/1476.png new file mode 100644 index 0000000..9fc713d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1476.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1477.png b/legacy_videos/legacy_clips/v32_180/1477.png new file mode 100644 index 0000000..2ad14d5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1477.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1478.png b/legacy_videos/legacy_clips/v32_180/1478.png new file mode 100644 index 0000000..dad198f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1478.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1479.png b/legacy_videos/legacy_clips/v32_180/1479.png new file mode 100644 index 0000000..a6a9259 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1479.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1480.png b/legacy_videos/legacy_clips/v32_180/1480.png new file mode 100644 index 0000000..c044fd0 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1480.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1481.png b/legacy_videos/legacy_clips/v32_180/1481.png new file mode 100644 index 0000000..6805239 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1481.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1482.png b/legacy_videos/legacy_clips/v32_180/1482.png new file mode 100644 index 0000000..5034dbc Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1482.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1483.png b/legacy_videos/legacy_clips/v32_180/1483.png new file mode 100644 index 0000000..dee60b9 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1483.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1484.png b/legacy_videos/legacy_clips/v32_180/1484.png new file mode 100644 index 0000000..45fdcd9 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1484.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1485.png b/legacy_videos/legacy_clips/v32_180/1485.png new file mode 100644 index 0000000..145741e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1485.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1486.png b/legacy_videos/legacy_clips/v32_180/1486.png new file mode 100644 index 0000000..72d02ab Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1486.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1487.png b/legacy_videos/legacy_clips/v32_180/1487.png new file mode 100644 index 0000000..11a8e2e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1487.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1488.png b/legacy_videos/legacy_clips/v32_180/1488.png new file mode 100644 index 0000000..b82a323 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1488.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1489.png b/legacy_videos/legacy_clips/v32_180/1489.png new file mode 100644 index 0000000..0145599 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1489.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1490.png b/legacy_videos/legacy_clips/v32_180/1490.png new file mode 100644 index 0000000..8d04174 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1490.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1491.png b/legacy_videos/legacy_clips/v32_180/1491.png new file mode 100644 index 0000000..9b96ee5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1491.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1492.png b/legacy_videos/legacy_clips/v32_180/1492.png new file mode 100644 index 0000000..1df3f3e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1492.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1493.png b/legacy_videos/legacy_clips/v32_180/1493.png new file mode 100644 index 0000000..3c5ca8c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1493.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1494.png b/legacy_videos/legacy_clips/v32_180/1494.png new file mode 100644 index 0000000..7d2ec3e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1494.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1495.png b/legacy_videos/legacy_clips/v32_180/1495.png new file mode 100644 index 0000000..5ac3e92 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1495.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1496.png b/legacy_videos/legacy_clips/v32_180/1496.png new file mode 100644 index 0000000..7ffce24 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1496.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1497.png b/legacy_videos/legacy_clips/v32_180/1497.png new file mode 100644 index 0000000..19d74e5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1497.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1498.png b/legacy_videos/legacy_clips/v32_180/1498.png new file mode 100644 index 0000000..ab8b320 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1498.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1499.png b/legacy_videos/legacy_clips/v32_180/1499.png new file mode 100644 index 0000000..2d3ab5c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1499.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1500.png b/legacy_videos/legacy_clips/v32_180/1500.png new file mode 100644 index 0000000..3e2c749 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1500.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1501.png b/legacy_videos/legacy_clips/v32_180/1501.png new file mode 100644 index 0000000..257e66e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1501.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1502.png b/legacy_videos/legacy_clips/v32_180/1502.png new file mode 100644 index 0000000..05d79e5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1502.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1503.png b/legacy_videos/legacy_clips/v32_180/1503.png new file mode 100644 index 0000000..305c214 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1503.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1504.png b/legacy_videos/legacy_clips/v32_180/1504.png new file mode 100644 index 0000000..53bd32e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1504.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1505.png b/legacy_videos/legacy_clips/v32_180/1505.png new file mode 100644 index 0000000..3780a78 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1505.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1506.png b/legacy_videos/legacy_clips/v32_180/1506.png new file mode 100644 index 0000000..9d3a14a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1506.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1507.png b/legacy_videos/legacy_clips/v32_180/1507.png new file mode 100644 index 0000000..4ed8f90 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1507.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1508.png b/legacy_videos/legacy_clips/v32_180/1508.png new file mode 100644 index 0000000..a2c9854 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1508.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1509.png b/legacy_videos/legacy_clips/v32_180/1509.png new file mode 100644 index 0000000..fd3033a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1509.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1510.png b/legacy_videos/legacy_clips/v32_180/1510.png new file mode 100644 index 0000000..e089603 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1510.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1511.png b/legacy_videos/legacy_clips/v32_180/1511.png new file mode 100644 index 0000000..d9435b2 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1511.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1512.png b/legacy_videos/legacy_clips/v32_180/1512.png new file mode 100644 index 0000000..750f4cf Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1512.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1513.png b/legacy_videos/legacy_clips/v32_180/1513.png new file mode 100644 index 0000000..5180131 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1513.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1514.png b/legacy_videos/legacy_clips/v32_180/1514.png new file mode 100644 index 0000000..ba386f8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1514.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1515.png b/legacy_videos/legacy_clips/v32_180/1515.png new file mode 100644 index 0000000..595bc91 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1515.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1516.png b/legacy_videos/legacy_clips/v32_180/1516.png new file mode 100644 index 0000000..f395a89 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1516.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1517.png b/legacy_videos/legacy_clips/v32_180/1517.png new file mode 100644 index 0000000..d3fa97b Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1517.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1518.png b/legacy_videos/legacy_clips/v32_180/1518.png new file mode 100644 index 0000000..ec4e36d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1518.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1519.png b/legacy_videos/legacy_clips/v32_180/1519.png new file mode 100644 index 0000000..6ac75ca Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1519.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1520.png b/legacy_videos/legacy_clips/v32_180/1520.png new file mode 100644 index 0000000..647f07e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1520.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1521.png b/legacy_videos/legacy_clips/v32_180/1521.png new file mode 100644 index 0000000..8d566af Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1521.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1522.png b/legacy_videos/legacy_clips/v32_180/1522.png new file mode 100644 index 0000000..c71daf5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1522.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1523.png b/legacy_videos/legacy_clips/v32_180/1523.png new file mode 100644 index 0000000..8c27432 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1523.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1524.png b/legacy_videos/legacy_clips/v32_180/1524.png new file mode 100644 index 0000000..3ec4636 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1524.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1525.png b/legacy_videos/legacy_clips/v32_180/1525.png new file mode 100644 index 0000000..ba6abaf Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1525.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1526.png b/legacy_videos/legacy_clips/v32_180/1526.png new file mode 100644 index 0000000..d52bee8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1526.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1527.png b/legacy_videos/legacy_clips/v32_180/1527.png new file mode 100644 index 0000000..04e8ccb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1527.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1528.png b/legacy_videos/legacy_clips/v32_180/1528.png new file mode 100644 index 0000000..32896b2 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1528.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1529.png b/legacy_videos/legacy_clips/v32_180/1529.png new file mode 100644 index 0000000..8c83bb6 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1529.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1530.png b/legacy_videos/legacy_clips/v32_180/1530.png new file mode 100644 index 0000000..9d2622f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1530.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1531.png b/legacy_videos/legacy_clips/v32_180/1531.png new file mode 100644 index 0000000..4790f2a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1531.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1532.png b/legacy_videos/legacy_clips/v32_180/1532.png new file mode 100644 index 0000000..49eb080 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1532.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1533.png b/legacy_videos/legacy_clips/v32_180/1533.png new file mode 100644 index 0000000..d97108d Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1533.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1534.png b/legacy_videos/legacy_clips/v32_180/1534.png new file mode 100644 index 0000000..db03cc1 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1534.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1535.png b/legacy_videos/legacy_clips/v32_180/1535.png new file mode 100644 index 0000000..c81ccce Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1535.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1536.png b/legacy_videos/legacy_clips/v32_180/1536.png new file mode 100644 index 0000000..92bf9bb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1536.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1537.png b/legacy_videos/legacy_clips/v32_180/1537.png new file mode 100644 index 0000000..e2fb6cc Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1537.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1538.png b/legacy_videos/legacy_clips/v32_180/1538.png new file mode 100644 index 0000000..87d1573 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1538.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1539.png b/legacy_videos/legacy_clips/v32_180/1539.png new file mode 100644 index 0000000..75af2af Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1539.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1540.png b/legacy_videos/legacy_clips/v32_180/1540.png new file mode 100644 index 0000000..267e4e8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1540.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1541.png b/legacy_videos/legacy_clips/v32_180/1541.png new file mode 100644 index 0000000..06d98da Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1541.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1542.png b/legacy_videos/legacy_clips/v32_180/1542.png new file mode 100644 index 0000000..de6f063 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1542.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1543.png b/legacy_videos/legacy_clips/v32_180/1543.png new file mode 100644 index 0000000..3687a0e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1543.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1544.png b/legacy_videos/legacy_clips/v32_180/1544.png new file mode 100644 index 0000000..857e9f1 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1544.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1545.png b/legacy_videos/legacy_clips/v32_180/1545.png new file mode 100644 index 0000000..7bdf849 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1545.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1546.png b/legacy_videos/legacy_clips/v32_180/1546.png new file mode 100644 index 0000000..44f2253 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1546.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1547.png b/legacy_videos/legacy_clips/v32_180/1547.png new file mode 100644 index 0000000..30d0bd8 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1547.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1548.png b/legacy_videos/legacy_clips/v32_180/1548.png new file mode 100644 index 0000000..2938868 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1548.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1549.png b/legacy_videos/legacy_clips/v32_180/1549.png new file mode 100644 index 0000000..c804c43 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1549.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1550.png b/legacy_videos/legacy_clips/v32_180/1550.png new file mode 100644 index 0000000..b716ae4 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1550.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1551.png b/legacy_videos/legacy_clips/v32_180/1551.png new file mode 100644 index 0000000..de5bd16 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1551.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1552.png b/legacy_videos/legacy_clips/v32_180/1552.png new file mode 100644 index 0000000..22358c0 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1552.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1553.png b/legacy_videos/legacy_clips/v32_180/1553.png new file mode 100644 index 0000000..9392271 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1553.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1554.png b/legacy_videos/legacy_clips/v32_180/1554.png new file mode 100644 index 0000000..0992530 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1554.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1555.png b/legacy_videos/legacy_clips/v32_180/1555.png new file mode 100644 index 0000000..a4edd57 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1555.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1556.png b/legacy_videos/legacy_clips/v32_180/1556.png new file mode 100644 index 0000000..bc061c6 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1556.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1557.png b/legacy_videos/legacy_clips/v32_180/1557.png new file mode 100644 index 0000000..4406973 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1557.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1558.png b/legacy_videos/legacy_clips/v32_180/1558.png new file mode 100644 index 0000000..8c5857f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1558.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1559.png b/legacy_videos/legacy_clips/v32_180/1559.png new file mode 100644 index 0000000..ab42054 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1559.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1560.png b/legacy_videos/legacy_clips/v32_180/1560.png new file mode 100644 index 0000000..a7a9763 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1560.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1561.png b/legacy_videos/legacy_clips/v32_180/1561.png new file mode 100644 index 0000000..a5e2498 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1561.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1562.png b/legacy_videos/legacy_clips/v32_180/1562.png new file mode 100644 index 0000000..ca2f403 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1562.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1563.png b/legacy_videos/legacy_clips/v32_180/1563.png new file mode 100644 index 0000000..d3b6725 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1563.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1564.png b/legacy_videos/legacy_clips/v32_180/1564.png new file mode 100644 index 0000000..afb4d8a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1564.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1565.png b/legacy_videos/legacy_clips/v32_180/1565.png new file mode 100644 index 0000000..82bc964 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1565.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1566.png b/legacy_videos/legacy_clips/v32_180/1566.png new file mode 100644 index 0000000..57b5133 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1566.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1567.png b/legacy_videos/legacy_clips/v32_180/1567.png new file mode 100644 index 0000000..89db3da Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1567.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1568.png b/legacy_videos/legacy_clips/v32_180/1568.png new file mode 100644 index 0000000..f6167d5 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1568.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1569.png b/legacy_videos/legacy_clips/v32_180/1569.png new file mode 100644 index 0000000..2333da2 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1569.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1570.png b/legacy_videos/legacy_clips/v32_180/1570.png new file mode 100644 index 0000000..e27979c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1570.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1571.png b/legacy_videos/legacy_clips/v32_180/1571.png new file mode 100644 index 0000000..76beec6 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1571.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1572.png b/legacy_videos/legacy_clips/v32_180/1572.png new file mode 100644 index 0000000..4d7182e Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1572.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1573.png b/legacy_videos/legacy_clips/v32_180/1573.png new file mode 100644 index 0000000..1ac8350 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1573.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1574.png b/legacy_videos/legacy_clips/v32_180/1574.png new file mode 100644 index 0000000..b5f522f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1574.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1575.png b/legacy_videos/legacy_clips/v32_180/1575.png new file mode 100644 index 0000000..c094ebd Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1575.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1576.png b/legacy_videos/legacy_clips/v32_180/1576.png new file mode 100644 index 0000000..d0a4538 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1576.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1577.png b/legacy_videos/legacy_clips/v32_180/1577.png new file mode 100644 index 0000000..56c9250 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1577.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1578.png b/legacy_videos/legacy_clips/v32_180/1578.png new file mode 100644 index 0000000..daca191 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1578.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1579.png b/legacy_videos/legacy_clips/v32_180/1579.png new file mode 100644 index 0000000..31f8eae Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1579.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1580.png b/legacy_videos/legacy_clips/v32_180/1580.png new file mode 100644 index 0000000..193ed6a Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1580.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1581.png b/legacy_videos/legacy_clips/v32_180/1581.png new file mode 100644 index 0000000..ad8aef7 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1581.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1582.png b/legacy_videos/legacy_clips/v32_180/1582.png new file mode 100644 index 0000000..18aaa5c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1582.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1583.png b/legacy_videos/legacy_clips/v32_180/1583.png new file mode 100644 index 0000000..4eae318 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1583.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1584.png b/legacy_videos/legacy_clips/v32_180/1584.png new file mode 100644 index 0000000..46f1eac Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1584.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1585.png b/legacy_videos/legacy_clips/v32_180/1585.png new file mode 100644 index 0000000..bfa2e2f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1585.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1586.png b/legacy_videos/legacy_clips/v32_180/1586.png new file mode 100644 index 0000000..872aede Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1586.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1587.png b/legacy_videos/legacy_clips/v32_180/1587.png new file mode 100644 index 0000000..288a6cb Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1587.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1588.png b/legacy_videos/legacy_clips/v32_180/1588.png new file mode 100644 index 0000000..71b1d9f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1588.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1589.png b/legacy_videos/legacy_clips/v32_180/1589.png new file mode 100644 index 0000000..3fa019f Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1589.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1590.png b/legacy_videos/legacy_clips/v32_180/1590.png new file mode 100644 index 0000000..a7fa9ec Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1590.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1591.png b/legacy_videos/legacy_clips/v32_180/1591.png new file mode 100644 index 0000000..ba06607 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1591.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1592.png b/legacy_videos/legacy_clips/v32_180/1592.png new file mode 100644 index 0000000..31059df Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1592.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1593.png b/legacy_videos/legacy_clips/v32_180/1593.png new file mode 100644 index 0000000..712430c Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1593.png differ diff --git a/legacy_videos/legacy_clips/v32_180/1594.png b/legacy_videos/legacy_clips/v32_180/1594.png new file mode 100644 index 0000000..b8b2f75 Binary files /dev/null and b/legacy_videos/legacy_clips/v32_180/1594.png differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00000.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00000.jpg new file mode 100644 index 0000000..25b085b Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00000.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00001.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00001.jpg new file mode 100644 index 0000000..8efe746 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00001.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00002.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00002.jpg new file mode 100644 index 0000000..c0f2422 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00002.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00003.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00003.jpg new file mode 100644 index 0000000..9e1f225 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00003.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00004.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00004.jpg new file mode 100644 index 0000000..c4b8035 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00004.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00005.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00005.jpg new file mode 100644 index 0000000..59af7d3 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00005.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00006.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00006.jpg new file mode 100644 index 0000000..16af058 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00006.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00007.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00007.jpg new file mode 100644 index 0000000..a4dfa89 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00007.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00008.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00008.jpg new file mode 100644 index 0000000..8837b37 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00008.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00009.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00009.jpg new file mode 100644 index 0000000..ebfdd00 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00009.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00010.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00010.jpg new file mode 100644 index 0000000..fd05f13 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00010.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00011.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00011.jpg new file mode 100644 index 0000000..3d8cdb6 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00011.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00012.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00012.jpg new file mode 100644 index 0000000..478a4cc Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00012.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00013.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00013.jpg new file mode 100644 index 0000000..b7b422b Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00013.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00014.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00014.jpg new file mode 100644 index 0000000..59c94a9 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00014.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00015.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00015.jpg new file mode 100644 index 0000000..c9cc428 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00015.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00016.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00016.jpg new file mode 100644 index 0000000..b2cc286 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00016.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00017.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00017.jpg new file mode 100644 index 0000000..300edfa Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00017.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00018.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00018.jpg new file mode 100644 index 0000000..c667ad0 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00018.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00019.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00019.jpg new file mode 100644 index 0000000..1cdb014 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00019.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00020.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00020.jpg new file mode 100644 index 0000000..d0acc88 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00020.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00021.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00021.jpg new file mode 100644 index 0000000..72f602a Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00021.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00022.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00022.jpg new file mode 100644 index 0000000..21e2458 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00022.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00023.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00023.jpg new file mode 100644 index 0000000..70511e0 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00023.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00024.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00024.jpg new file mode 100644 index 0000000..5d5731a Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00024.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00025.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00025.jpg new file mode 100644 index 0000000..15be6a9 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00025.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00026.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00026.jpg new file mode 100644 index 0000000..b8cdd64 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00026.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00027.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00027.jpg new file mode 100644 index 0000000..6a714f7 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00027.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00028.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00028.jpg new file mode 100644 index 0000000..466fa8f Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00028.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00029.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00029.jpg new file mode 100644 index 0000000..15115ea Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00029.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00030.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00030.jpg new file mode 100644 index 0000000..5223f66 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00030.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00031.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00031.jpg new file mode 100644 index 0000000..7cc6233 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00031.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00032.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00032.jpg new file mode 100644 index 0000000..8ddaad0 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00032.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00033.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00033.jpg new file mode 100644 index 0000000..845948b Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00033.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00034.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00034.jpg new file mode 100644 index 0000000..3434418 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00034.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00035.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00035.jpg new file mode 100644 index 0000000..f546ca9 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00035.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00036.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00036.jpg new file mode 100644 index 0000000..79b35b0 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00036.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00037.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00037.jpg new file mode 100644 index 0000000..9c8e303 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00037.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00038.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00038.jpg new file mode 100644 index 0000000..0cdc597 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00038.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00039.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00039.jpg new file mode 100644 index 0000000..a9652c4 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00039.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00040.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00040.jpg new file mode 100644 index 0000000..95bbd33 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00040.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00041.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00041.jpg new file mode 100644 index 0000000..6353f02 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00041.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00042.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00042.jpg new file mode 100644 index 0000000..bedb552 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00042.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00043.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00043.jpg new file mode 100644 index 0000000..3c50a8a Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00043.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00044.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00044.jpg new file mode 100644 index 0000000..d4e0f12 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00044.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00045.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00045.jpg new file mode 100644 index 0000000..d7049b9 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00045.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00046.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00046.jpg new file mode 100644 index 0000000..49f2c54 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00046.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00047.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00047.jpg new file mode 100644 index 0000000..d3e2106 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00047.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00048.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00048.jpg new file mode 100644 index 0000000..51631f3 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00048.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00049.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00049.jpg new file mode 100644 index 0000000..a7d7d98 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00049.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00050.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00050.jpg new file mode 100644 index 0000000..57bee92 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00050.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00051.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00051.jpg new file mode 100644 index 0000000..9fea587 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00051.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00052.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00052.jpg new file mode 100644 index 0000000..283c4fc Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00052.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00053.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00053.jpg new file mode 100644 index 0000000..8a88075 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00053.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00054.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00054.jpg new file mode 100644 index 0000000..b9c5470 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00054.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00055.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00055.jpg new file mode 100644 index 0000000..0f30360 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00055.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00056.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00056.jpg new file mode 100644 index 0000000..e70dbf1 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00056.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00057.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00057.jpg new file mode 100644 index 0000000..fdc1c14 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00057.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00058.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00058.jpg new file mode 100644 index 0000000..3fba6f3 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00058.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00059.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00059.jpg new file mode 100644 index 0000000..2a55092 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00059.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00060.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00060.jpg new file mode 100644 index 0000000..de24b82 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00060.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00061.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00061.jpg new file mode 100644 index 0000000..b44a9f9 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00061.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00062.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00062.jpg new file mode 100644 index 0000000..2e16148 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00062.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00063.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00063.jpg new file mode 100644 index 0000000..9842f8c Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00063.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00064.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00064.jpg new file mode 100644 index 0000000..8b5bf1a Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00064.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00065.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00065.jpg new file mode 100644 index 0000000..b0513d8 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00065.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00066.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00066.jpg new file mode 100644 index 0000000..ff27933 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00066.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00067.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00067.jpg new file mode 100644 index 0000000..7942fc6 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00067.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00068.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00068.jpg new file mode 100644 index 0000000..4789054 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00068.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00069.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00069.jpg new file mode 100644 index 0000000..b285707 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00069.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00070.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00070.jpg new file mode 100644 index 0000000..d1e1f42 Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00070.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00071.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00071.jpg new file mode 100644 index 0000000..ae7054a Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00071.jpg differ diff --git a/legacy_videos/legacy_clips_output/v32_v32_01/00072.jpg b/legacy_videos/legacy_clips_output/v32_v32_01/00072.jpg new file mode 100644 index 0000000..d43028f Binary files /dev/null and b/legacy_videos/legacy_clips_output/v32_v32_01/00072.jpg differ diff --git a/legacy_videos/legacy_clips_ref/v32_01.jpg b/legacy_videos/legacy_clips_ref/v32_01.jpg new file mode 100644 index 0000000..430f4cd Binary files /dev/null and b/legacy_videos/legacy_clips_ref/v32_01.jpg differ diff --git a/legacy_videos/legacy_clips_ref/v32_02.jpg b/legacy_videos/legacy_clips_ref/v32_02.jpg new file mode 100644 index 0000000..f764bb1 Binary files /dev/null and b/legacy_videos/legacy_clips_ref/v32_02.jpg differ diff --git a/legacy_videos/legacy_clips_ref/v32_03.jpg b/legacy_videos/legacy_clips_ref/v32_03.jpg new file mode 100644 index 0000000..88023a1 Binary files /dev/null and b/legacy_videos/legacy_clips_ref/v32_03.jpg differ diff --git a/legacy_videos/legacy_clips_ref/v32_04.jpg b/legacy_videos/legacy_clips_ref/v32_04.jpg new file mode 100644 index 0000000..9b4476c Binary files /dev/null and b/legacy_videos/legacy_clips_ref/v32_04.jpg differ diff --git a/lib/.vscode/settings.json b/lib/.vscode/settings.json new file mode 100644 index 0000000..2b7e46d --- /dev/null +++ b/lib/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.formatting.provider": "yapf" +} \ No newline at end of file diff --git a/lib/FeatVGG.py b/lib/FeatVGG.py new file mode 100755 index 0000000..bb06ae3 --- /dev/null +++ b/lib/FeatVGG.py @@ -0,0 +1,65 @@ +import torch.nn as nn +import utils.vgg_util as vgg_util +from torchvision import models + + +def conv_to_relu(layer_names): + out_layernames = [] + for name in layer_names: + if name.startswith("conv"): + out_layernames.append("relu" + name[4:]) + else: + out_layernames.append(name) + return out_layernames + + +class FeatVGG(nn.Module): + def __init__(self, content_layers=["relu3_1"]): + super(FeatVGG, self).__init__() + self.content_layers = conv_to_relu(content_layers) + # set up vgg19 model + self.vgg19 = vgg_util.get_renamed_vgg() + self.last_c_layer = self.content_layers[-1] + is_last_content = False + replace_layers, del_layers = [], [] + for name, mod in self.vgg19.named_children(): + if is_last_content: + del_layers.append(name) + else: + if name == self.last_c_layer: + is_last_content = True + + for name in del_layers: + delattr(self.vgg19, name) + + # no need for gradweight vgg19 + for param in self.vgg19.parameters(): + param.requires_grad = False + + def forward(self, input_img): + # input image is BGR image + # each channel ranges in [0,255] + # should be normalized with mean = [0.406*255, 0.456*255, 0.485*255] = [103,116,123] + # out = {} + return self.vgg19(input_img) + + +class VGGNet_multilayer(nn.Module): + def __init__(self): + """Select conv1_1 ~ conv5_1 activation maps.""" + super(VGGNet_multilayer, self).__init__() + self.select = ["0", "5", "10", "19", "28"] + self.vgg = models.vgg19(pretrained=True).features + + def forward(self, x): + # x: rgb image + # ranges in [0,1] + # should be normalzied with mean = [0.485, 0.456, 0.406] + # and variance = [0.229, 0.224, 0.225] + """Extract multiple convolutional feature maps.""" + features = [] + for name, layer in self.vgg._modules.items(): + x = layer(x) + if name in self.select: + features.append(x) + return features diff --git a/lib/HardSigmoid.py b/lib/HardSigmoid.py new file mode 100755 index 0000000..267382b --- /dev/null +++ b/lib/HardSigmoid.py @@ -0,0 +1,20 @@ +from torch.autograd import Function, Variable + + +class HardSigmoid(Function): + @staticmethod + def forward(ctx, input): + output = (0.2 * input) + 0.5 + output[output > 1] = 1 + output[output < 0] = 0 + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + input = ctx.saved_variables + grad_input = input.new(input.size()) + grad_input.data.fill_(0.2) + grad_input[input <= -2.5] = 0 + grad_input[input >= 2.5] = 0 + return grad_input diff --git a/lib/RefPairTransformsWithGlobalOrigRGB.py b/lib/RefPairTransformsWithGlobalOrigRGB.py new file mode 100644 index 0000000..9c9dfc4 --- /dev/null +++ b/lib/RefPairTransformsWithGlobalOrigRGB.py @@ -0,0 +1,359 @@ +from __future__ import division + +import math +import random + +import torch +from PIL import Image, ImageEnhance, ImageOps + +try: + import accimage +except ImportError: + accimage = None +import collections +import numbers +import types +import warnings + +import numpy as np +from skimage import color + +import lib.functional as F + +__all__ = [ + "Compose", + "Concatenate", + "ToTensor", + "Normalize", + "Resize", + "Scale", + "CenterCrop", + "Pad", + "RandomCrop", + "RandomHorizontalFlip", + "RandomVerticalFlip", + "RandomResizedCrop", + "RandomSizedCrop", + "FiveCrop", + "TenCrop", + "RGB2Lab", +] + + +def CustomFunc(inputs, func, *args, **kwargs): + im_l = func(inputs[0], *args, **kwargs) + im_ab = func(inputs[1], *args, **kwargs) + warp_ba = func(inputs[2], *args, **kwargs) + warp_aba = func(inputs[3], *args, **kwargs) + im_gbl_ab = func(inputs[4], *args, **kwargs) + bgr_mc_im = func(inputs[5], *args, **kwargs) + + layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + + for l in range(5): + layer = inputs[6 + l] + err_ba = func(layer[0], *args, **kwargs) + err_ab = func(layer[1], *args, **kwargs) + + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Compose(object): + """Composes several transforms together. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.ToTensor(), + >>> ]) + """ + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, inputs): + for t in self.transforms: + inputs = t(inputs) + return inputs + + +class Concatenate(object): + """ + Input: [im_l, im_ab, inputs] + inputs = [warp_ba_l, warp_ba_ab, warp_aba, err_pm, err_aba] + + Output:[im_l, err_pm, warp_ba, warp_aba, im_ab, err_aba] + """ + + def __call__(self, inputs): + im_l = inputs[0] + im_ab = inputs[1] + warp_ba = inputs[2] + warp_aba = inputs[3] + im_glb_ab = inputs[4] + bgr_mc_im = inputs[5] + bgr_mc_im = bgr_mc_im[[2, 1, 0], ...] + + err_ba = [] + err_ab = [] + + for l in range(5): + layer = inputs[6 + l] + err_ba.append(layer[0]) + err_ab.append(layer[1]) + + cerr_ba = torch.cat(err_ba, 0) + cerr_ab = torch.cat(err_ab, 0) + + return (im_l, cerr_ba, warp_ba, warp_aba, im_glb_ab, bgr_mc_im, im_ab, cerr_ab) + + +class ToTensor(object): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]. + """ + + def __call__(self, inputs): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + inputs = CustomFunc(inputs, F.to_mytensor) + return inputs + + +class Normalize(object): + """Normalize an tensor image with mean and standard deviation. + Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform + will normalize each channel of the input ``torch.*Tensor`` i.e. + ``input[channel] = (input[channel] - mean[channel]) / std[channel]`` + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + """ + + def __call__(self, inputs): + """ + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + + im_l = F.normalize(inputs[0], 50, 1) # [0, 100] + im_ab = F.normalize(inputs[1], (0, 0), (1, 1)) # [-100, 100] + + inputs[2][0:1, :, :] = F.normalize(inputs[2][0:1, :, :], 50, 1) + inputs[2][1:3, :, :] = F.normalize(inputs[2][1:3, :, :], (0, 0), (1, 1)) + warp_ba = inputs[2] + + inputs[3][0:1, :, :] = F.normalize(inputs[3][0:1, :, :], 50, 1) + inputs[3][1:3, :, :] = F.normalize(inputs[3][1:3, :, :], (0, 0), (1, 1)) + warp_aba = inputs[3] + + im_gbl_ab = F.normalize(inputs[4], (0, 0), (1, 1)) # [-100, 100] + + bgr_mc_im = F.normalize(inputs[5], (123.68, 116.78, 103.938), (1, 1, 1)) + + layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + + for l in range(5): + layer = inputs[6 + l] + err_ba = F.normalize(layer[0], 127, 2) # [0, 255] + err_ab = F.normalize(layer[1], 127, 2) # [0, 255] + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Resize(object): + """Resize the input PIL Image to the given size. + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (h, w), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + i.e, if height > width, then image will be rescaled to + (size * height / width, size) + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR`` + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2) + self.size = size + self.interpolation = interpolation + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be scaled. + + Returns: + PIL Image: Rescaled image. + """ + return CustomFunc(inputs, F.resize, self.size, self.interpolation) + + +class RandomCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = random.randint(0, h - th) + j = random.randint(0, w - tw) + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + i, j, h, w = self.get_params(inputs[0], self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class CenterCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = (h - th) // 2 + j = (w - tw) // 2 + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + i, j, h, w = self.get_params(inputs[0], self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class RandomHorizontalFlip(object): + """Horizontally flip the given PIL Image randomly with a probability of 0.5.""" + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + if random.random() < 0.5: + return CustomFunc(inputs, F.hflip) + return inputs + + +class RGB2Lab(object): + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs[0]) + warp_ba_lab = color.rgb2lab(inputs[2]) + warp_aba_lab = color.rgb2lab(inputs[3]) + im_gbl_lab = color.rgb2lab(inputs[4]) + + inputs[0] = image_lab[:, :, :1] # l channel + inputs[1] = image_lab[:, :, 1:] # ab channel + inputs[2] = warp_ba_lab # lab channel + inputs[3] = warp_aba_lab # lab channel + inputs[4] = im_gbl_lab[:, :, 1:] # ab channel + + return inputs diff --git a/lib/TestTransforms.py b/lib/TestTransforms.py new file mode 100644 index 0000000..f9ac6f0 --- /dev/null +++ b/lib/TestTransforms.py @@ -0,0 +1,353 @@ +from __future__ import division + +import collections +import numbers +import random + +import torch +from PIL import Image +from skimage import color + +import lib.functional as F + +__all__ = [ + "Compose", + "Concatenate", + "ToTensor", + "Normalize", + "Resize", + "Scale", + "CenterCrop", + "Pad", + "RandomCrop", + "RandomHorizontalFlip", + "RandomVerticalFlip", + "RandomResizedCrop", + "RandomSizedCrop", + "FiveCrop", + "TenCrop", + "RGB2Lab", +] + + +def CustomFunc(inputs, func, *args, **kwargs): + im_l = func(inputs[0], *args, **kwargs) + im_ab = func(inputs[1], *args, **kwargs) + warp_ba = func(inputs[2], *args, **kwargs) + warp_aba = func(inputs[3], *args, **kwargs) + # im_gbl_ab = func(inputs[4], *args, **kwargs) + # bgr_mc_im = func(inputs[5], *args, **kwargs) + layer_data = [im_l, im_ab, warp_ba, warp_aba] + # layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + for l in range(5): + layer = inputs[4 + l] + err_ba = func(layer[0], *args, **kwargs) + err_ab = func(layer[1], *args, **kwargs) + + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Compose(object): + """Composes several transforms together. + + Args: + transforms (list of ``Transform`` objects): list of transforms to compose. + + Example: + >>> transforms.Compose([ + >>> transforms.CenterCrop(10), + >>> transforms.ToTensor(), + >>> ]) + """ + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, inputs): + for t in self.transforms: + inputs = t(inputs) + return inputs + + +class Concatenate(object): + """ + Input: [im_l, im_ab, inputs] + inputs = [warp_ba_l, warp_ba_ab, warp_aba, err_pm, err_aba] + + Output:[im_l, err_pm, warp_ba, warp_aba, im_ab, err_aba] + """ + + def __call__(self, inputs): + im_l = inputs[0] + im_ab = inputs[1] + warp_ba = inputs[2] + warp_aba = inputs[3] + # im_glb_ab = inputs[4] + # bgr_mc_im = inputs[5] + # bgr_mc_im = bgr_mc_im[[2, 1, 0], ...] + + err_ba = [] + err_ab = [] + + for l in range(5): + layer = inputs[4 + l] + err_ba.append(layer[0]) + err_ab.append(layer[1]) + + cerr_ba = torch.cat(err_ba, 0) + cerr_ab = torch.cat(err_ab, 0) + + return (im_l, cerr_ba, warp_ba, warp_aba, im_ab, cerr_ab) + # return (im_l, cerr_ba, warp_ba, warp_aba, im_glb_ab, bgr_mc_im, im_ab, cerr_ab) + + +class ToTensor(object): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + Converts a PIL Image or numpy.ndarray (H x W x C) in the range + [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]. + """ + + def __call__(self, inputs): + """ + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + inputs = CustomFunc(inputs, F.to_mytensor) + return inputs + + +class Normalize(object): + """Normalize an tensor image with mean and standard deviation. + Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform + will normalize each channel of the input ``torch.*Tensor`` i.e. + ``input[channel] = (input[channel] - mean[channel]) / std[channel]`` + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + """ + + def __call__(self, inputs): + """ + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + + im_l = F.normalize(inputs[0], 50, 1) # [0, 100] + im_ab = F.normalize(inputs[1], (0, 0), (1, 1)) # [-100, 100] + + inputs[2][0:1, :, :] = F.normalize(inputs[2][0:1, :, :], 50, 1) + inputs[2][1:3, :, :] = F.normalize(inputs[2][1:3, :, :], (0, 0), (1, 1)) + warp_ba = inputs[2] + + inputs[3][0:1, :, :] = F.normalize(inputs[3][0:1, :, :], 50, 1) + inputs[3][1:3, :, :] = F.normalize(inputs[3][1:3, :, :], (0, 0), (1, 1)) + warp_aba = inputs[3] + + # im_gbl_ab = F.normalize(inputs[4], (0, 0), (1, 1)) # [-100, 100] + # + # bgr_mc_im = F.normalize(inputs[5], (123.68, 116.78, 103.938), (1, 1, 1)) + + # layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im] + layer_data = [im_l, im_ab, warp_ba, warp_aba] + + for l in range(5): + layer = inputs[4 + l] + err_ba = F.normalize(layer[0], 127, 2) # [0, 255] + err_ab = F.normalize(layer[1], 127, 2) # [0, 255] + layer_data.append([err_ba, err_ab]) + + return layer_data + + +class Resize(object): + """Resize the input PIL Image to the given size. + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (h, w), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + i.e, if height > width, then image will be rescaled to + (size * height / width, size) + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR`` + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2) + self.size = size + self.interpolation = interpolation + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be scaled. + + Returns: + PIL Image: Rescaled image. + """ + return CustomFunc(inputs, F.resize, self.size, self.interpolation) + + +class RandomCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = random.randint(0, h - th) + j = random.randint(0, w - tw) + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + i, j, h, w = self.get_params(inputs[0], self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class CenterCrop(object): + """Crop the given PIL Image at a random location. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + padding (int or sequence, optional): Optional padding on each border + of the image. Default is 0, i.e no padding. If a sequence of length + 4 is provided, it is used to pad left, top, right, bottom borders + respectively. + """ + + def __init__(self, size, padding=0): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + self.padding = padding + + @staticmethod + def get_params(img, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + w, h = img.size + th, tw = output_size + if w == tw and h == th: + return 0, 0, h, w + + i = (h - th) // 2 + j = (w - tw) // 2 + return i, j, th, tw + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be cropped. + + Returns: + PIL Image: Cropped image. + """ + if self.padding > 0: + inputs = CustomFunc(inputs, F.pad, self.padding) + + if type(inputs) is list: + i, j, h, w = self.get_params(inputs[0], self.size) + else: + i, j, h, w = self.get_params(inputs, self.size) + return CustomFunc(inputs, F.crop, i, j, h, w) + + +class RandomHorizontalFlip(object): + """Horizontally flip the given PIL Image randomly with a probability of 0.5.""" + + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + if random.random() < 0.5: + return CustomFunc(inputs, F.hflip) + return inputs + + +class RGB2Lab(object): + def __call__(self, inputs): + """ + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Randomly flipped image. + """ + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs[0]) + warp_ba_lab = color.rgb2lab(inputs[2]) + warp_aba_lab = color.rgb2lab(inputs[3]) + # im_gbl_lab = color.rgb2lab(inputs[4]) + + inputs[0] = image_lab[:, :, :1] # l channel + inputs[1] = image_lab[:, :, 1:] # ab channel + inputs[2] = warp_ba_lab # lab channel + inputs[3] = warp_aba_lab # lab channel + # inputs[4] = im_gbl_lab[:, :, 1:] # ab channel + + return inputs diff --git a/lib/VGGFeatureLoss.py b/lib/VGGFeatureLoss.py new file mode 100755 index 0000000..d9c782e --- /dev/null +++ b/lib/VGGFeatureLoss.py @@ -0,0 +1,90 @@ +import torch.nn as nn +import utils.vgg_util as vgg_util + + +def conv_to_relu(layer_names): + out_layernames = [] + for name in layer_names: + if name.startswith("conv"): + out_layernames.append("relu" + name[4:]) + else: + out_layernames.append(name) + return out_layernames + + +class VGGFeatureLoss(object): + def __init__(self, content_layers=["relu3_1"], content_weights=[1]): + self.content_layers = conv_to_relu(content_layers) + self.content_weights = content_weights + # set up vgg19 model + self.vgg19 = vgg_util.get_renamed_vgg() + self.last_c_layer = self.content_layers[-1] + is_last_content = False + replace_layers, del_layers = [], [] + for name, mod in self.vgg19.named_children(): + if is_last_content: + del_layers.append(name) + else: + if name == self.last_c_layer: + is_last_content = True + print("del_layers", del_layers) + for name in del_layers: + delattr(self.vgg19, name) + + # set up loss + self.loss_layers = self.content_layers + self.contentloss_fns = [nn.MSELoss()] * len(self.content_layers) + self.content_loss = None + + # target + self.content_targets = None + + # no need for gradweight vgg19 + for param in self.vgg19.parameters(): + param.requires_grad = False + + def vgg_forward(self, input_img, layers, last_layer): + # out = {} + x = input_img + out = None + # print(last_layer) + assert len(layers) == 1 and layers[-1] == last_layer + for name, mod in self.vgg19.named_children(): + # print('name is : ', name) + x = mod(x) + if name == last_layer: + out = x + break + # out[name] = mod(prev_input) + # prev_input = out[name] + # if name == last_layer: + # break + return [out] + # return [out[key] for key in layers] + + def cuda(self, gpu_ids): + if len(gpu_ids) > 1: + self.vgg19 = nn.DataParallel(self.vgg19.cuda(), gpu_ids) + else: + self.vgg19.cuda() + self.contentloss_fns = [loss_fn.cuda() for loss_fn in self.contentloss_fns] + + def set_content_targets(self, content_img): + self.content_targets = [ + A.detach() for A in self.vgg_forward(content_img, self.content_layers, self.last_c_layer) + ] + + def forward(self, img): + assert self.content_targets is not None, "set up content targets first" + out = self.vgg_forward(img, self.content_layers, self.last_c_layer) + self.content_loss = sum( + [ + self.content_weights[idx] * self.contentloss_fns[idx](out[idx], self.content_targets[idx]) + for idx in range(len(self.content_layers)) + ] + ) + return self.content_loss + + # def backward(self): + # self.loss.backward() + diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/lib/flowlib.py b/lib/flowlib.py new file mode 100644 index 0000000..f982829 --- /dev/null +++ b/lib/flowlib.py @@ -0,0 +1,522 @@ +#!/usr/bin/python +""" +# ============================== +# flowlib.py +# library for optical flow processing +# Author: Ruoteng Li +# Date: 6th Aug 2016 +# ============================== +""" +import matplotlib.colors as cl +import matplotlib.pyplot as plt +import numpy as np +import png +from PIL import Image + +UNKNOWN_FLOW_THRESH = 1e7 +SMALLFLOW = 0.0 +LARGEFLOW = 1e8 + +""" +============= +Flow Section +============= +""" + + +def show_flow(filename): + """ + visualize optical flow map using matplotlib + :param filename: optical flow file + :return: None + """ + flow = read_flow(filename) + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + + +def visualize_flow(flow, mode="Y"): + """ + this function visualize the input flow + :param flow: input flow in array + :param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color) + :return: None + """ + if mode == "Y": + # Ccbcr color wheel + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + elif mode == "RGB": + (h, w) = flow.shape[0:2] + du = flow[:, :, 0] + dv = flow[:, :, 1] + valid = flow[:, :, 2] + max_flow = max(np.max(du), np.max(dv)) + img = np.zeros((h, w, 3), dtype=np.float64) + # angle layer + img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi) + # magnitude layer, normalized to 1 + img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow + # phase layer + img[:, :, 2] = 8 - img[:, :, 1] + # clip to [0,1] + small_idx = img[:, :, 0:3] < 0 + large_idx = img[:, :, 0:3] > 1 + img[small_idx] = 0 + img[large_idx] = 1 + # convert to rgb + img = cl.hsv_to_rgb(img) + # remove invalid point + img[:, :, 0] = img[:, :, 0] * valid + img[:, :, 1] = img[:, :, 1] * valid + img[:, :, 2] = img[:, :, 2] * valid + # show + plt.imshow(img) + plt.show() + + return None + + +def read_flow(filename): + """ + read optical flow from Middlebury .flo file + :param filename: name of the flow file + :return: optical flow data in matrix + """ + f = open(filename, "rb") + try: + magic = np.fromfile(f, np.float32, count=1)[0] # For Python3.x + except: + magic = np.fromfile(f, np.float32, count=1) # For Python2.x + data2d = None + + if 202021.25 != magic: + print("Magic number incorrect. Invalid .flo file") + else: + w = np.fromfile(f, np.int32, count=1)[0] + h = np.fromfile(f, np.int32, count=1)[0] + # print("Reading %d x %d flo file" % (h, w)) + data2d = np.fromfile(f, np.float32, count=2 * w * h) + # reshape data into 3D array (columns, rows, channels) + data2d = np.resize(data2d, (h, w, 2)) + f.close() + return data2d + + +def read_flow_png(flow_file): + """ + Read optical flow from KITTI .png file + :param flow_file: name of the flow file + :return: optical flow data in matrix + """ + flow_object = png.Reader(filename=flow_file) + flow_direct = flow_object.asDirect() + flow_data = list(flow_direct[2]) + (w, h) = flow_direct[3]["size"] + flow = np.zeros((h, w, 3), dtype=np.float64) + for i in range(len(flow_data)): + flow[i, :, 0] = flow_data[i][0::3] + flow[i, :, 1] = flow_data[i][1::3] + flow[i, :, 2] = flow_data[i][2::3] + + invalid_idx = flow[:, :, 2] == 0 + flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0 + flow[invalid_idx, 0] = 0 + flow[invalid_idx, 1] = 0 + return flow + + +def write_flow(flow, filename): + """ + write optical flow in Middlebury .flo format + :param flow: optical flow map + :param filename: optical flow file path to be saved + :return: None + """ + f = open(filename, "wb") + magic = np.array([202021.25], dtype=np.float32) + (height, width) = flow.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + flow.tofile(f) + f.close() + + +def segment_flow(flow): + h = flow.shape[0] + w = flow.shape[1] + u = flow[:, :, 0] + v = flow[:, :, 1] + + idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW) + idx2 = abs(u) == SMALLFLOW + class0 = (v == 0) & (u == 0) + u[idx2] = 0.00001 + tan_value = v / u + + class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0) + class2 = (tan_value >= 1) & (u >= 0) & (v >= 0) + class3 = (tan_value < -1) & (u <= 0) & (v >= 0) + class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0) + class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0) + class7 = (tan_value < -1) & (u >= 0) & (v <= 0) + class6 = (tan_value >= 1) & (u <= 0) & (v <= 0) + class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0) + + seg = np.zeros((h, w)) + + seg[class1] = 1 + seg[class2] = 2 + seg[class3] = 3 + seg[class4] = 4 + seg[class5] = 5 + seg[class6] = 6 + seg[class7] = 7 + seg[class8] = 8 + seg[class0] = 0 + seg[idx] = 0 + + return seg + + +def flow_error(tu, tv, u, v): + """ + Calculate average end point error + :param tu: ground-truth horizontal flow map + :param tv: ground-truth vertical flow map + :param u: estimated horizontal flow map + :param v: estimated vertical flow map + :return: End point error of the estimated flow + """ + smallflow = 0.0 + """ + stu = tu[bord+1:end-bord,bord+1:end-bord] + stv = tv[bord+1:end-bord,bord+1:end-bord] + su = u[bord+1:end-bord,bord+1:end-bord] + sv = v[bord+1:end-bord,bord+1:end-bord] + """ + stu = tu[:] + stv = tv[:] + su = u[:] + sv = v[:] + + idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH) + stu[idxUnknow] = 0 + stv[idxUnknow] = 0 + su[idxUnknow] = 0 + sv[idxUnknow] = 0 + + ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)] + index_su = su[ind2] + index_sv = sv[ind2] + an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1) + un = index_su * an + vn = index_sv * an + + index_stu = stu[ind2] + index_stv = stv[ind2] + tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1) + tun = index_stu * tn + tvn = index_stv * tn + + """ + angle = un * tun + vn * tvn + (an * tn) + index = [angle == 1.0] + angle[index] = 0.999 + ang = np.arccos(angle) + mang = np.mean(ang) + mang = mang * 180 / np.pi + """ + + epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2) + epe = epe[ind2] + mepe = np.mean(epe) + return mepe + + +def flow_to_image(flow, display=False): + """ + Convert flow into middlebury color code image + :param flow: optical flow map + :return: optical flow image in middlebury color + """ + u = flow[:, :, 0] + v = flow[:, :, 1] + + maxu = -999.0 + maxv = -999.0 + minu = 999.0 + minv = 999.0 + + idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH) + u[idxUnknow] = 0 + v[idxUnknow] = 0 + + maxu = max(maxu, np.max(u)) + minu = min(minu, np.min(u)) + + maxv = max(maxv, np.max(v)) + minv = min(minv, np.min(v)) + + rad = np.sqrt(u ** 2 + v ** 2) + maxrad = max(-1, np.max(rad)) + + if display: + print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv)) + + u = u / (maxrad + np.finfo(float).eps) + v = v / (maxrad + np.finfo(float).eps) + + img = compute_color(u, v) + + idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2) + img[idx] = 0 + + return np.uint8(img) + + +def evaluate_flow_file(gt, pred): + """ + evaluate the estimated optical flow end point error according to ground truth provided + :param gt: ground truth file path + :param pred: estimated optical flow file path + :return: end point error, float32 + """ + # Read flow files and calculate the errors + gt_flow = read_flow(gt) # ground truth flow + eva_flow = read_flow(pred) # predicted flow + # Calculate errors + average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1]) + return average_pe + + +def evaluate_flow(gt_flow, pred_flow): + """ + gt: ground-truth flow + pred: estimated flow + """ + average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], pred_flow[:, :, 0], pred_flow[:, :, 1]) + return average_pe + + +""" +============== +Disparity Section +============== +""" + + +def read_disp_png(file_name): + """ + Read optical flow from KITTI .png file + :param file_name: name of the flow file + :return: optical flow data in matrix + """ + image_object = png.Reader(filename=file_name) + image_direct = image_object.asDirect() + image_data = list(image_direct[2]) + (w, h) = image_direct[3]["size"] + channel = len(image_data[0]) / w + flow = np.zeros((h, w, channel), dtype=np.uint16) + for i in range(len(image_data)): + for j in range(channel): + flow[i, :, j] = image_data[i][j::channel] + return flow[:, :, 0] / 256 + + +def disp_to_flowfile(disp, filename): + """ + Read KITTI disparity file in png format + :param disp: disparity matrix + :param filename: the flow file name to save + :return: None + """ + f = open(filename, "wb") + magic = np.array([202021.25], dtype=np.float32) + (height, width) = disp.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + empty_map = np.zeros((height, width), dtype=np.float32) + data = np.dstack((disp, empty_map)) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + data.tofile(f) + f.close() + + +""" +============== +Image Section +============== +""" + + +def read_image(filename): + """ + Read normal image of any format + :param filename: name of the image file + :return: image data in matrix uint8 type + """ + img = Image.open(filename) + im = np.array(img) + return im + + +def warp_image(im, flow): + """ + Use optical flow to warp image to the next + :param im: image to warp + :param flow: optical flow + :return: warped image + """ + from scipy import interpolate + + image_height = im.shape[0] + image_width = im.shape[1] + flow_height = flow.shape[0] + flow_width = flow.shape[1] + n = image_height * image_width + (iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width]) + (fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width]) + fx += flow[:, :, 0].astype(np.float) + fy += flow[:, :, 1] + mask = np.logical_or(fx < 0, fx > flow_width) + mask = np.logical_or(mask, fy < 0) + mask = np.logical_or(mask, fy > flow_height) + fx = np.minimum(np.maximum(fx, 0), flow_width) + fy = np.minimum(np.maximum(fy, 0), flow_height) + points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1) + xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1) + warp = np.zeros((image_height, image_width, im.shape[2])) + for i in range(im.shape[2]): + channel = im[:, :, i] + # plt.imshow(channel, cmap='gray') + values = channel.reshape(n, 1) + new_channel = interpolate.griddata(points, values, xi, method="cubic") + new_channel = np.reshape(new_channel, [flow_height, flow_width]) + new_channel[mask] = 1 + warp[:, :, i] = new_channel.astype(np.uint8) + + return warp.astype(np.uint8) + + +""" +============== +Others +============== +""" + + +def scale_image(image, new_range): + """ + Linearly scale the image into desired range + :param image: input image + :param new_range: the new range to be aligned + :return: image normalized in new range + """ + min_val = np.min(image).astype(np.float32) + max_val = np.max(image).astype(np.float32) + min_val_new = np.array(min(new_range), dtype=np.float32) + max_val_new = np.array(max(new_range), dtype=np.float32) + scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new + return scaled_image.astype(np.uint8) + + +def compute_color(u, v): + """ + compute optical flow color map + :param u: optical flow horizontal map + :param v: optical flow vertical map + :return: optical flow in color code + """ + [h, w] = u.shape + img = np.zeros([h, w, 3]) + nanIdx = np.isnan(u) | np.isnan(v) + u[nanIdx] = 0 + v[nanIdx] = 0 + + colorwheel = make_color_wheel() + ncols = np.size(colorwheel, 0) + + rad = np.sqrt(u ** 2 + v ** 2) + + a = np.arctan2(-v, -u) / np.pi + + fk = (a + 1) / 2 * (ncols - 1) + 1 + + k0 = np.floor(fk).astype(int) + + k1 = k0 + 1 + k1[k1 == ncols + 1] = 1 + f = fk - k0 + + for i in range(0, np.size(colorwheel, 1)): + tmp = colorwheel[:, i] + col0 = tmp[k0 - 1] / 255 + col1 = tmp[k1 - 1] / 255 + col = (1 - f) * col0 + f * col1 + + idx = rad <= 1 + col[idx] = 1 - rad[idx] * (1 - col[idx]) + notidx = np.logical_not(idx) + + col[notidx] *= 0.75 + img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx))) + + return img + + +def make_color_wheel(): + """ + Generate color wheel according Middlebury color code + :return: Color wheel + """ + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + + colorwheel = np.zeros([ncols, 3]) + + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY)) + col += RY + + # YG + colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG)) + colorwheel[col : col + YG, 1] = 255 + col += YG + + # GC + colorwheel[col : col + GC, 1] = 255 + colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC)) + col += GC + + # CB + colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB)) + colorwheel[col : col + CB, 2] = 255 + col += CB + + # BM + colorwheel[col : col + BM, 2] = 255 + colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM)) + col += +BM + + # MR + colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) + colorwheel[col : col + MR, 0] = 255 + + return colorwheel diff --git a/lib/functional.py b/lib/functional.py new file mode 100755 index 0000000..f36181d --- /dev/null +++ b/lib/functional.py @@ -0,0 +1,606 @@ +from __future__ import division + +import math +import random + +import torch +from PIL import Image, ImageEnhance, ImageOps + +try: + import accimage +except ImportError: + accimage = None +import collections +import numbers +import types +import warnings + +import numpy as np + + +def _is_pil_image(img): + if accimage is not None: + return isinstance(img, (Image.Image, accimage.Image)) + else: + return isinstance(img, Image.Image) + + +def _is_tensor_image(img): + return torch.is_tensor(img) and img.ndimension() == 3 + + +def _is_numpy_image(img): + return isinstance(img, np.ndarray) and (img.ndim in {2, 3}) + + +def to_tensor(pic): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + See ``ToTensor`` for more details. + + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + if not (_is_pil_image(pic) or _is_numpy_image(pic)): + raise TypeError("pic should be PIL Image or ndarray. Got {}".format(type(pic))) + + if isinstance(pic, np.ndarray): + # handle numpy array + img = torch.from_numpy(pic.transpose((2, 0, 1))) + # backward compatibility + return img.float().div(255) + + if accimage is not None and isinstance(pic, accimage.Image): + nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32) + pic.copyto(nppic) + return torch.from_numpy(nppic) + + # handle PIL Image + if pic.mode == "I": + img = torch.from_numpy(np.array(pic, np.int32, copy=False)) + elif pic.mode == "I;16": + img = torch.from_numpy(np.array(pic, np.int16, copy=False)) + else: + img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) + # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK + if pic.mode == "YCbCr": + nchannel = 3 + elif pic.mode == "I;16": + nchannel = 1 + else: + nchannel = len(pic.mode) + img = img.view(pic.size[1], pic.size[0], nchannel) + # put it from HWC to CHW format + # yikes, this transpose takes 80% of the loading time/CPU + img = img.transpose(0, 1).transpose(0, 2).contiguous() + if isinstance(img, torch.ByteTensor): + return img.float().div(255) + else: + return img + + +def to_mytensor(pic): + """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. + + See ``ToTensor`` for more details. + + Args: + pic (PIL Image or numpy.ndarray): Image to be converted to tensor. + + Returns: + Tensor: Converted image. + """ + pic_arr = np.array(pic) + if pic_arr.ndim == 2: + pic_arr = pic_arr[..., np.newaxis] + img = torch.from_numpy(pic_arr.transpose((2, 0, 1))) + if not isinstance(img, torch.FloatTensor): + return img.float() # no normalize .div(255) + else: + return img + + +def to_pil_image(pic, mode=None): + """Convert a tensor or an ndarray to PIL Image. + + See :class:`~torchvision.transforms.ToPIlImage` for more details. + + Args: + pic (Tensor or numpy.ndarray): Image to be converted to PIL Image. + mode (`PIL.Image mode`_): color space and pixel depth of input data (optional). + + .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes + + Returns: + PIL Image: Image converted to PIL Image. + """ + if not (_is_numpy_image(pic) or _is_tensor_image(pic)): + raise TypeError("pic should be Tensor or ndarray. Got {}.".format(type(pic))) + + npimg = pic + if isinstance(pic, torch.FloatTensor): + pic = pic.mul(255).byte() + if torch.is_tensor(pic): + npimg = np.transpose(pic.numpy(), (1, 2, 0)) + + if not isinstance(npimg, np.ndarray): + raise TypeError("Input pic must be a torch.Tensor or NumPy ndarray, " + "not {}".format(type(npimg))) + + if npimg.shape[2] == 1: + expected_mode = None + npimg = npimg[:, :, 0] + if npimg.dtype == np.uint8: + expected_mode = "L" + if npimg.dtype == np.int16: + expected_mode = "I;16" + if npimg.dtype == np.int32: + expected_mode = "I" + elif npimg.dtype == np.float32: + expected_mode = "F" + if mode is not None and mode != expected_mode: + raise ValueError( + "Incorrect mode ({}) supplied for input type {}. Should be {}".format(mode, np.dtype, expected_mode) + ) + mode = expected_mode + + elif npimg.shape[2] == 4: + permitted_4_channel_modes = ["RGBA", "CMYK"] + if mode is not None and mode not in permitted_4_channel_modes: + raise ValueError("Only modes {} are supported for 4D inputs".format(permitted_4_channel_modes)) + + if mode is None and npimg.dtype == np.uint8: + mode = "RGBA" + else: + permitted_3_channel_modes = ["RGB", "YCbCr", "HSV"] + if mode is not None and mode not in permitted_3_channel_modes: + raise ValueError("Only modes {} are supported for 3D inputs".format(permitted_3_channel_modes)) + if mode is None and npimg.dtype == np.uint8: + mode = "RGB" + + if mode is None: + raise TypeError("Input type {} is not supported".format(npimg.dtype)) + + return Image.fromarray(npimg, mode=mode) + + +def normalize(tensor, mean, std): + """Normalize a tensor image with mean and standard deviation. + + See ``Normalize`` for more details. + + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channely. + + Returns: + Tensor: Normalized Tensor image. + """ + if not _is_tensor_image(tensor): + raise TypeError("tensor is not a torch image.") + # TODO: make efficient + if tensor.size(0) == 1: + tensor.sub_(mean).div_(std) + else: + for t, m, s in zip(tensor, mean, std): + t.sub_(m).div_(s) + return tensor + + +def resize(img, size, interpolation=Image.BILINEAR): + """Resize the input PIL Image to the given size. + + Args: + img (PIL Image): Image to be resized. + size (sequence or int): Desired output size. If size is a sequence like + (h, w), the output size will be matched to this. If size is an int, + the smaller edge of the image will be matched to this number maintaing + the aspect ratio. i.e, if height > width, then image will be rescaled to + (size * height / width, size) + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR`` + + Returns: + PIL Image: Resized image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + if not (isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)): + raise TypeError("Got inappropriate size arg: {}".format(size)) + + if isinstance(size, int): + w, h = img.size + if (w <= h and w == size) or (h <= w and h == size): + return img + if w < h: + ow = size + oh = int(round(size * h / w)) + return img.resize((ow, oh), interpolation) + else: + oh = size + ow = int(round(size * w / h)) + return img.resize((ow, oh), interpolation) + else: + return img.resize(size[::-1], interpolation) + + +def scale(*args, **kwargs): + warnings.warn("The use of the transforms.Scale transform is deprecated, " + "please use transforms.Resize instead.") + return resize(*args, **kwargs) + + +def pad(img, padding, fill=0): + """Pad the given PIL Image on all sides with the given "pad" value. + + Args: + img (PIL Image): Image to be padded. + padding (int or tuple): Padding on each border. If a single int is provided this + is used to pad all borders. If tuple of length 2 is provided this is the padding + on left/right and top/bottom respectively. If a tuple of length 4 is provided + this is the padding for the left, top, right and bottom borders + respectively. + fill: Pixel fill value. Default is 0. If a tuple of + length 3, it is used to fill R, G, B channels respectively. + + Returns: + PIL Image: Padded image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + if not isinstance(padding, (numbers.Number, tuple)): + raise TypeError("Got inappropriate padding arg") + if not isinstance(fill, (numbers.Number, str, tuple)): + raise TypeError("Got inappropriate fill arg") + + if isinstance(padding, collections.Sequence) and len(padding) not in [2, 4]: + raise ValueError( + "Padding must be an int or a 2, or 4 element tuple, not a " + "{} element tuple".format(len(padding)) + ) + + return ImageOps.expand(img, border=padding, fill=fill) + + +def crop(img, i, j, h, w): + """Crop the given PIL Image. + + Args: + img (PIL Image): Image to be cropped. + i: Upper pixel coordinate. + j: Left pixel coordinate. + h: Height of the cropped image. + w: Width of the cropped image. + + Returns: + PIL Image: Cropped image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + return img.crop((j, i, j + w, i + h)) + + +def center_crop(img, output_size): + if isinstance(output_size, numbers.Number): + output_size = (int(output_size), int(output_size)) + w, h = img.size + th, tw = output_size + i = int(round((h - th) / 2.0)) + j = int(round((w - tw) / 2.0)) + return crop(img, i, j, th, tw) + + +def resized_crop(img, i, j, h, w, size, interpolation=Image.BILINEAR): + """Crop the given PIL Image and resize it to desired size. + + Notably used in RandomResizedCrop. + + Args: + img (PIL Image): Image to be cropped. + i: Upper pixel coordinate. + j: Left pixel coordinate. + h: Height of the cropped image. + w: Width of the cropped image. + size (sequence or int): Desired output size. Same semantics as ``scale``. + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR``. + Returns: + PIL Image: Cropped image. + """ + assert _is_pil_image(img), "img should be PIL Image" + img = crop(img, i, j, h, w) + img = resize(img, size, interpolation) + return img + + +def hflip(img): + """Horizontally flip the given PIL Image. + + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Horizontall flipped image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + return img.transpose(Image.FLIP_LEFT_RIGHT) + + +def vflip(img): + """Vertically flip the given PIL Image. + + Args: + img (PIL Image): Image to be flipped. + + Returns: + PIL Image: Vertically flipped image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + return img.transpose(Image.FLIP_TOP_BOTTOM) + + +def five_crop(img, size): + """Crop the given PIL Image into four corners and the central crop. + + .. Note:: + This transform returns a tuple of images and there may be a + mismatch in the number of inputs and targets your ``Dataset`` returns. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + Returns: + tuple: tuple (tl, tr, bl, br, center) corresponding top left, + top right, bottom left, bottom right and center crop. + """ + if isinstance(size, numbers.Number): + size = (int(size), int(size)) + else: + assert len(size) == 2, "Please provide only two dimensions (h, w) for size." + + w, h = img.size + crop_h, crop_w = size + if crop_w > w or crop_h > h: + raise ValueError("Requested crop size {} is bigger than input size {}".format(size, (h, w))) + tl = img.crop((0, 0, crop_w, crop_h)) + tr = img.crop((w - crop_w, 0, w, crop_h)) + bl = img.crop((0, h - crop_h, crop_w, h)) + br = img.crop((w - crop_w, h - crop_h, w, h)) + center = center_crop(img, (crop_h, crop_w)) + return (tl, tr, bl, br, center) + + +def ten_crop(img, size, vertical_flip=False): + """Crop the given PIL Image into four corners and the central crop plus the + flipped version of these (horizontal flipping is used by default). + + .. Note:: + This transform returns a tuple of images and there may be a + mismatch in the number of inputs and targets your ``Dataset`` returns. + + Args: + size (sequence or int): Desired output size of the crop. If size is an + int instead of sequence like (h, w), a square crop (size, size) is + made. + vertical_flip (bool): Use vertical flipping instead of horizontal + + Returns: + tuple: tuple (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip, + br_flip, center_flip) corresponding top left, top right, + bottom left, bottom right and center crop and same for the + flipped image. + """ + if isinstance(size, numbers.Number): + size = (int(size), int(size)) + else: + assert len(size) == 2, "Please provide only two dimensions (h, w) for size." + + first_five = five_crop(img, size) + + if vertical_flip: + img = vflip(img) + else: + img = hflip(img) + + second_five = five_crop(img, size) + return first_five + second_five + + +def adjust_brightness(img, brightness_factor): + """Adjust brightness of an Image. + + Args: + img (PIL Image): PIL Image to be adjusted. + brightness_factor (float): How much to adjust the brightness. Can be + any non negative number. 0 gives a black image, 1 gives the + original image while 2 increases the brightness by a factor of 2. + + Returns: + PIL Image: Brightness adjusted image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + enhancer = ImageEnhance.Brightness(img) + img = enhancer.enhance(brightness_factor) + return img + + +def adjust_contrast(img, contrast_factor): + """Adjust contrast of an Image. + + Args: + img (PIL Image): PIL Image to be adjusted. + contrast_factor (float): How much to adjust the contrast. Can be any + non negative number. 0 gives a solid gray image, 1 gives the + original image while 2 increases the contrast by a factor of 2. + + Returns: + PIL Image: Contrast adjusted image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + enhancer = ImageEnhance.Contrast(img) + img = enhancer.enhance(contrast_factor) + return img + + +def adjust_saturation(img, saturation_factor): + """Adjust color saturation of an image. + + Args: + img (PIL Image): PIL Image to be adjusted. + saturation_factor (float): How much to adjust the saturation. 0 will + give a black and white image, 1 will give the original image while + 2 will enhance the saturation by a factor of 2. + + Returns: + PIL Image: Saturation adjusted image. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + enhancer = ImageEnhance.Color(img) + img = enhancer.enhance(saturation_factor) + return img + + +def adjust_hue(img, hue_factor): + """Adjust hue of an image. + + The image hue is adjusted by converting the image to HSV and + cyclically shifting the intensities in the hue channel (H). + The image is then converted back to original image mode. + + `hue_factor` is the amount of shift in H channel and must be in the + interval `[-0.5, 0.5]`. + + See https://en.wikipedia.org/wiki/Hue for more details on Hue. + + Args: + img (PIL Image): PIL Image to be adjusted. + hue_factor (float): How much to shift the hue channel. Should be in + [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in + HSV space in positive and negative direction respectively. + 0 means no shift. Therefore, both -0.5 and 0.5 will give an image + with complementary colors while 0 gives the original image. + + Returns: + PIL Image: Hue adjusted image. + """ + if not (-0.5 <= hue_factor <= 0.5): + raise ValueError("hue_factor is not in [-0.5, 0.5].".format(hue_factor)) + + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + input_mode = img.mode + if input_mode in {"L", "1", "I", "F"}: + return img + + h, s, v = img.convert("HSV").split() + + np_h = np.array(h, dtype=np.uint8) + # uint8 addition take cares of rotation across boundaries + with np.errstate(over="ignore"): + np_h += np.uint8(hue_factor * 255) + h = Image.fromarray(np_h, "L") + + img = Image.merge("HSV", (h, s, v)).convert(input_mode) + return img + + +def adjust_gamma(img, gamma, gain=1): + """Perform gamma correction on an image. + + Also known as Power Law Transform. Intensities in RGB mode are adjusted + based on the following equation: + + I_out = 255 * gain * ((I_in / 255) ** gamma) + + See https://en.wikipedia.org/wiki/Gamma_correction for more details. + + Args: + img (PIL Image): PIL Image to be adjusted. + gamma (float): Non negative real number. gamma larger than 1 make the + shadows darker, while gamma smaller than 1 make dark regions + lighter. + gain (float): The constant multiplier. + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + if gamma < 0: + raise ValueError("Gamma should be a non-negative real number") + + input_mode = img.mode + img = img.convert("RGB") + + np_img = np.array(img, dtype=np.float32) + np_img = 255 * gain * ((np_img / 255) ** gamma) + np_img = np.uint8(np.clip(np_img, 0, 255)) + + img = Image.fromarray(np_img, "RGB").convert(input_mode) + return img + + +def rotate(img, angle, resample=False, expand=False, center=None): + """Rotate the image by angle and then (optionally) translate it by (n_columns, n_rows) + + + Args: + img (PIL Image): PIL Image to be rotated. + angle ({float, int}): In degrees degrees counter clockwise order. + resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional): + An optional resampling filter. + See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters + If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST. + expand (bool, optional): Optional expansion flag. + If true, expands the output image to make it large enough to hold the entire rotated image. + If false or omitted, make the output image the same size as the input image. + Note that the expand flag assumes rotation around the center and no translation. + center (2-tuple, optional): Optional center of rotation. + Origin is the upper left corner. + Default is the center of the image. + """ + + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + return img.rotate(angle, resample, expand, center) + + +def to_grayscale(img, num_output_channels=1): + """Convert image to grayscale version of image. + + Args: + img (PIL Image): Image to be converted to grayscale. + + Returns: + PIL Image: Grayscale version of the image. + if num_output_channels == 1 : returned image is single channel + if num_output_channels == 3 : returned image is 3 channel with r == g == b + """ + if not _is_pil_image(img): + raise TypeError("img should be PIL Image. Got {}".format(type(img))) + + if num_output_channels == 1: + img = img.convert("L") + elif num_output_channels == 3: + img = img.convert("L") + np_img = np.array(img, dtype=np.uint8) + np_img = np.dstack([np_img, np_img, np_img]) + img = Image.fromarray(np_img, "RGB") + else: + raise ValueError("num_output_channels should be either 1 or 3") + + return img diff --git a/lib/imageloader_imagenet.py b/lib/imageloader_imagenet.py new file mode 100644 index 0000000..c4bef26 --- /dev/null +++ b/lib/imageloader_imagenet.py @@ -0,0 +1,179 @@ +import os.path as osp + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.utils.data as data +import torchvision.transforms as transforms +from PIL import Image +from scipy.ndimage.filters import gaussian_filter +from utils.util_distortion import CenterPadCrop_numpy, Distortion_with_flow, RGB2Lab, ToTensor, random_mask + +import lib.functional as F +from lib.RefPairDatasetWithGlobalOrigRGB_optimize import ToTensor, combo5_loader, image_loader, parse_images + +cv2.setNumThreads(0) + + +class ImageNetDataset(data.Dataset): + def __init__( + self, + data_root, + epoch, + image_size, + with_bad=False, + with_mid=False, + transforms_imagenet=None, + distortion_level=3, + brightnessjitter=5, + extra_reference_transform=None, + real_reference_probability=1, + mask_reference=False, + no_reference_probability=0, + ): + image_pairs = [] + curr_image_pairs = parse_images(data_root, with_bad, with_mid) + image_pairs += curr_image_pairs + print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs))) + if len(image_pairs) == 0: + raise RuntimeError("Found 0 image_a pairs in all the data_roots") + + self.image_pairs = image_pairs + self.epoch = epoch + self.image_size = image_size + self.real_len = len(self.image_pairs) + self.image_pairs *= epoch + + self.real_reference_probability = real_reference_probability + self.no_reference_probability = no_reference_probability + self.mask_reference = mask_reference + + self.transforms_imagenet_raw = transforms_imagenet + self.transforms_imagenet = transforms.Compose(transforms_imagenet) + self.extra_reference_transform = transforms.Compose(extra_reference_transform) + self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()]) + + self.distortion_level = distortion_level + self.distortion_transform = Distortion_with_flow() + self.brightnessjitter = brightnessjitter + + def __getitem__(self, index): + try: + # print('imagenet loader') + + image_id = 0 + pair_id = index + + combo_path = None + image_a_path = None + image_b_path = None + + image_names = ["", ""] + dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id] + sub_dir = osp.join(dir_root, cls_dir) + if is_good >= 1: + image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id]) + elif is_good == 0: + image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id]) + else: + image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id]) + + if np.random.random() > 0.5: + image_a_path, image_b_path = image_b_path, image_a_path + + I1 = image_loader(image_a_path) + if np.random.random() < self.no_reference_probability: + no_reference_flag = 1 + I_reference_fake = Image.fromarray(np.array(I1) * 0) + I_reference_real = I_reference_fake + else: + no_reference_flag = 0 + I_reference_fake = I1 + I_reference_real = image_loader(image_b_path) + + ## generate the flow + height, width = np.array(I1).shape[0], np.array(I1).shape[1] + alpha = np.random.rand() * self.distortion_level + distortion_range = 50 + random_state = np.random.RandomState(None) + shape = self.image_size[0], self.image_size[1] + # dx: flow on the vertical direction; dy: flow on the horizontal direction + forward_dx = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + forward_dy = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + + # transform the input image for colorization + I1 = self.transforms_imagenet(I1) + + # DEBUG + # from utils.util import batch_lab2rgb_transpose_mc + # plt.imshow(image_loader(image_a_path)) + # plt.figure() + # I1 = I1.unsqueeze(0) + # plt.imshow(batch_lab2rgb_transpose_mc(I1[0:32, 0:1, :, :], I1[0:32, 1:3, :, :])) + + # transform the fake reference + I_reference_fake = self.extra_reference_transform(I_reference_fake) + for transform in self.transforms_imagenet_raw: + if type(transform) is RGB2Lab and self.extra_reference_transform is not None: + I_reference_fake = self.distortion_transform(I_reference_fake, forward_dx, forward_dy) + I_reference_fake = transform(I_reference_fake) + I_reference_fake[0:1, :, :] = I_reference_fake[0:1, :, :] + torch.randn(1) * self.brightnessjitter + + # transform the real reference + I_reference_real = self.transforms_imagenet(I_reference_real) + + # generate the masked image + if self.mask_reference: + mask = torch.Tensor(random_mask(self.image_size[0], self.image_size[1])) + I_reference_fake[1:2, :, :] = ( + I_reference_fake[1:2, :, :] * (1 - mask) + I_reference_real[1:2, :, :] * mask + ) + I_reference_fake[2:3, :, :] = ( + I_reference_fake[2:3, :, :] * (1 - mask) + I_reference_real[2:3, :, :] * mask + ) + + # choose the reference + if no_reference_flag: + I_reference_output, fake_refernce_flag, no_reference_flag = ( + I_reference_real, + torch.Tensor([0]), + torch.Tensor([1]).unsqueeze(-1).unsqueeze(-1), + ) + else: + no_reference_flag = torch.Tensor([0]).unsqueeze(-1).unsqueeze(-1) + I_reference_output, fake_refernce_flag = ( + (I_reference_real, torch.Tensor([0])) + if np.random.random() < self.real_reference_probability + else (I_reference_fake, torch.Tensor([1])) + ) + + outputs = [] + outputs.append(I1) + outputs.append(I_reference_output) + outputs.append(fake_refernce_flag) + outputs.append(no_reference_flag) + + except Exception as e: + if combo_path is not None: + print("problem in ", combo_path) + print("problem in, ", image_a_path) + print(e) + return self.__getitem__(np.random.randint(0, len(self.image_pairs))) + + return outputs + + def __len__(self): + return len(self.image_pairs) + # return len(self.image_pairs) * 2 diff --git a/lib/imports.py b/lib/imports.py new file mode 100755 index 0000000..2c49d6d --- /dev/null +++ b/lib/imports.py @@ -0,0 +1,82 @@ +import collections +import datetime +import itertools +import json +import math +import os +import pickle +import random +import re +import string +import sys +import threading +import time +import warnings +from abc import abstractmethod +from collections import Counter, Iterable, OrderedDict +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from functools import partial +from glob import glob, iglob +from itertools import chain +from operator import attrgetter, itemgetter + +import bcolz +import cv2 +import graphviz +import IPython +import matplotlib +import numpy as np +import pandas as pd +import PIL +import scipy +import seaborn as sns +import sklearn +import sklearn_pandas +from IPython.lib.deepreload import reload as dreload +from IPython.lib.display import FileLink +from ipywidgets import fixed, interact, interactive, widgets +from isoweek import Week +from matplotlib import animation +from matplotlib import pyplot as plt +from matplotlib import rcParams +from pandas_summary import DataFrameSummary +from PIL import Image, ImageEnhance, ImageOps +from sklearn import ensemble, metrics, preprocessing + +matplotlib.rc("animation", html="html5") +np.set_printoptions(precision=5, linewidth=110, suppress=True) + + +def in_notebook(): + return "ipykernel" in sys.modules + + +import tqdm as tq +from tqdm import tnrange, tqdm_notebook + + +def clear_tqdm(): + inst = getattr(tq.tqdm, "_instances", None) + if not inst: + return + for i in range(len(inst)): + inst.pop().close() + + +if in_notebook(): + + def tqdm(*args, **kwargs): + clear_tqdm() + return tq.tqdm(*args, file=sys.stdout, **kwargs) + + def trange(*args, **kwargs): + clear_tqdm() + return tq.trange(*args, file=sys.stdout, **kwargs) + + +else: + from tqdm import tqdm, trange + + tnrange = trange + tqdm_notebook = tqdm + diff --git a/lib/videoloader.py b/lib/videoloader.py new file mode 100644 index 0000000..0d281d0 --- /dev/null +++ b/lib/videoloader.py @@ -0,0 +1,390 @@ +import sys + +sys.path.insert(0, "..") +import os +import struct + +import cv2 +import numpy as np +import torch.utils.data as data +from PIL import Image + +cv2.setNumThreads(0) +import pdb +import random + +import matplotlib.pyplot as plt +import torch +import torchvision +import torchvision.transforms as transforms +from PIL import Image +from skimage import color +from utils.flowlib import read_flow +from utils.util_distortion import CenterPad, CenterPadCrop_numpy +from utils.warping import WarpingLayer + +import lib.functional as F + + +class RGB2Lab(object): + def __init__(self): + pass + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs) + return image_lab + + +class Normalize(object): + def __init__(self): + pass + + def __call__(self, inputs): + inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1) + inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1)) + return inputs + + +class ToTensor(object): + def __init__(self): + pass + + def __call__(self, inputs): + outputs = F.to_mytensor(inputs) # permute channel and transform to tensor + return outputs + + +class RandomErasing(object): + def __init__(self, probability=0.6, sl=0.05, sh=0.6): + self.probability = probability + self.sl = sl + self.sh = sh + + def __call__(self, img): + img = np.array(img) + if random.uniform(0, 1) > self.probability: + return Image.fromarray(img) + + area = img.shape[0] * img.shape[1] + h0 = img.shape[0] + w0 = img.shape[1] + channel = img.shape[2] + + h = int(round(random.uniform(self.sl, self.sh) * h0)) + w = int(round(random.uniform(self.sl, self.sh) * w0)) + + if w < img.shape[1] and h < img.shape[0]: + x1 = random.randint(0, img.shape[0] - h) + y1 = random.randint(0, img.shape[1] - w) + img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255 + return Image.fromarray(img) + + return Image.fromarray(img) + + +class CenterCrop(object): + """ + center crop the numpy array + """ + + def __init__(self, image_size): + # target size + self.h0, self.w0 = image_size + + def __call__(self, input_numpy): + if input_numpy.ndim == 3: + h, w, channel = input_numpy.shape + output_numpy = np.zeros((self.h0, self.w0, channel)) + output_numpy = input_numpy[ + (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, : + ] + else: + h, w = input_numpy.shape + output_numpy = np.zeros((self.h0, self.w0)) + output_numpy = input_numpy[ + (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0 + ] + return output_numpy + + +def parse_images(data_root): + image_pairs = [] + subdirs = sorted(os.listdir(data_root)) + for subdir in subdirs: + path = os.path.join(data_root, subdir) + if not os.path.isdir(path): + continue + + # parse_file = os.path.join(path, 'pairs_output.txt') + parse_file = os.path.join(path, "pairs_output_new.txt") + if os.path.exists(parse_file): + f = open(parse_file, "r") + lines = f.readlines() + for line in lines: + line = line.replace("\n", "") + ( + image1_name, + image2_name, + reference_video_name, + reference_video_name1, + reference_name1, + reference_name2, + reference_name3, + reference_name4, + reference_name5, + reference_gt1, + reference_gt2, + reference_gt3, + ) = line.split() + image1_name = image1_name.split(".")[0] + image2_name = image2_name.split(".")[0] + reference_video_name = reference_video_name.split(".")[0] + reference_video_name1 = reference_video_name1.split(".")[0] + reference_name1 = reference_name1.split(".")[0] + reference_name2 = reference_name2.split(".")[0] + reference_name3 = reference_name3.split(".")[0] + reference_name4 = reference_name4.split(".")[0] + reference_name5 = reference_name5.split(".")[0] + + reference_gt1 = reference_gt1.split(".")[0] + reference_gt2 = reference_gt2.split(".")[0] + reference_gt3 = reference_gt3.split(".")[0] + + flow_forward_name = image1_name + "_forward" + flow_backward_name = image1_name + "_backward" + mask_name = image1_name + "_mask" + + item = ( + image1_name + ".jpg", + image2_name + ".jpg", + reference_video_name + ".jpg", + reference_name1 + ".JPEG", + reference_name2 + ".JPEG", + reference_name3 + ".JPEG", + reference_name4 + ".JPEG", + reference_name5 + ".JPEG", + flow_forward_name + ".flo", + flow_backward_name + ".flo", + mask_name + ".pgm", + reference_gt1 + ".jpg", + reference_gt2 + ".jpg", + reference_gt3 + ".jpg", + path, + ) + image_pairs.append(item) + + f.close() + + else: + raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n")) + + return image_pairs + + +class VideosDataset(torch.utils.data.Dataset): + def __init__( + self, + data_root, + epoch, + image_size, + image_transform=None, + use_google_reference=False, + real_reference_probability=1, + nonzero_placeholder_probability=0.5, + ): + self.data_root = data_root + self.image_transform = image_transform + self.CenterPad = CenterPad(image_size) + self.ToTensor = ToTensor() + self.CenterCrop = CenterCrop(image_size) + + assert len(self.data_root) > 0, "find no dataroot" + self.epoch = epoch + self.image_pairs = parse_images(self.data_root) + self.real_len = len(self.image_pairs) + print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len)) + self.image_pairs *= epoch + self.use_google_reference = use_google_reference + self.real_reference_probability = real_reference_probability + self.nonzero_placeholder_probability = nonzero_placeholder_probability + + def __getitem__(self, index): + ( + image1_name, + image2_name, + reference_video_name, + reference_name1, + reference_name2, + reference_name3, + reference_name4, + reference_name5, + flow_forward_name, + flow_backward_name, + mask_name, + reference_gt1, + reference_gt2, + reference_gt3, + path, + ) = self.image_pairs[index] + try: + # print('video loader') + I1 = Image.open(os.path.join(path, "input_pad", image1_name)) + I2 = Image.open(os.path.join(path, "input_pad", image2_name)) + + # randomly choose a reference frame from the video + I_reference_video = Image.open( + os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3])) + ) + I_reference_video_real = Image.open( + os.path.join( + path, + "reference", + random.choice( + [reference_name1, reference_name2, reference_name3, reference_name4, reference_name5] + ), + ) + ) + + flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name)) # numpy + flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name)) # numpy + mask = Image.open(os.path.join(path, "mask", mask_name)) + + # binary mask + mask = np.array(mask) + mask[mask < 240] = 0 + mask[mask >= 240] = 1 + + # transform + I1 = self.image_transform(I1) + I2 = self.image_transform(I2) + I_reference_video = self.image_transform(self.CenterPad(I_reference_video)) + I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real)) + flow_forward = self.ToTensor(self.CenterCrop(flow_forward)) + flow_backward = self.ToTensor(self.CenterCrop(flow_backward)) + mask = self.ToTensor(self.CenterCrop(mask)) + + if np.random.random() < self.real_reference_probability: + I_reference_output = I_reference_video_real + placeholder = torch.zeros_like(I1) + self_ref_flag = torch.zeros_like(I1) + else: + I_reference_output = I_reference_video + placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1) + self_ref_flag = torch.ones_like(I1) + + outputs = [] + outputs.append(I1) + outputs.append(I2) + outputs.append(I_reference_output) + # outputs.append(I_reference_video_real) + outputs.append(flow_forward) + outputs.append(flow_backward) + outputs.append(mask) + outputs.append(placeholder) + outputs.append(self_ref_flag) + + except Exception as e: + print("problem in, ", path) + print(e) + return self.__getitem__(np.random.randint(0, len(self.image_pairs))) + + return outputs + + def __len__(self): + return len(self.image_pairs) + + +import torchvision.utils as vutils +from torch.autograd import Variable + + +def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8): + if isinstance(img_l_mc, Variable): + img_l_mc = img_l_mc.data.cpu() + if isinstance(img_ab_mc, Variable): + img_ab_mc = img_ab_mc.data.cpu() + + if img_l_mc.is_cuda: + img_l_mc = img_l_mc.cpu() + if img_ab_mc.is_cuda: + img_ab_mc = img_ab_mc.cpu() + + assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input" + + l_norm, ab_norm = 1.0, 1.0 + l_mean, ab_mean = 50.0, 0 + img_l = img_l_mc * l_norm + l_mean + img_ab = img_ab_mc * ab_norm + ab_mean + pred_lab = torch.cat((img_l, img_ab), dim=1) + grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64") + grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8") + return grid_rgb + + +# if __name__ == "__main__": +def test_videoloader(): + # dataset initI1lization + # image size = 256 + train_dataset = VideosDataset(data_root="../../Dataset/video/video_pair", epoch=10) + + data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0) + + iter_instance = iter(data_loader) + sample = iter_instance.next() + I1_lab, I2_lab, I_reference_lab, flow_forward, flow_backward, mask = sample + + # l: [-50,50], ab: [-128, 128] + l_norm, ab_norm = 1.0, 1.0 + l_mean, ab_mean = 50.0, 0 + I1_l = I1_lab[:, 0:1, :, :] + I1_ab = I1_lab[:, 1:3, :, :] + I2_l = I2_lab[:, 0:1, :, :] + I2_ab = I2_lab[:, 1:3, :, :] + I_reference_l = I_reference_lab[:, 0:1, :, :] + I_reference_ab = I_reference_lab[:, 1:3, :, :] + + # test image samples + image_A_grid = batch_lab2rgb_transpose_mc(I1_l, I1_ab, nrow=4).astype(np.uint8) + image_B_grid = batch_lab2rgb_transpose_mc(I2_l, I2_ab, nrow=4).astype(np.uint8) + image_ref_grid = batch_lab2rgb_transpose_mc(I_reference_l, I_reference_ab, nrow=4).astype(np.uint8) + plt.imsave("A.png", image_A_grid) + plt.imsave("B.png", image_B_grid) + plt.figure() + plt.imshow(image_A_grid) + plt.title("image A") + plt.figure() + plt.imshow(image_B_grid) + plt.title("image B") + plt.figure() + plt.imshow(image_ref_grid) + plt.title("image ref") + # plt.show() + + # test warping + print(I2_lab.shape) + print(flow_forward.shape) + warping_layer = WarpingLayer("gpu") + I21_lab = warping_layer(I2_lab.cuda(), flow_forward.cuda()) + I21_lab = I21_lab * mask.cuda() + + I21_l = I21_lab[:, 0:1, :, :] + I21_ab = I21_lab[:, 1:3, :, :] + image_warp_grid = batch_lab2rgb_transpose_mc(I21_l, I21_ab, nrow=4).astype(np.uint8) + plt.imsave("BA.png", image_warp_grid) + plt.figure() + plt.imshow(image_warp_grid) + plt.title("image warp") + # plt.show() + + # image difference + plt.figure() + plt.imshow(np.abs(image_warp_grid - image_A_grid).astype(np.uint8)) + plt.title("warp-A") + plt.figure() + plt.imshow(np.abs(image_B_grid - image_A_grid).astype(np.uint8)) + plt.title("B-A") + # plt.show() + plt.imsave("difference_warp.png", np.abs(image_warp_grid - image_A_grid).astype(np.uint8)) + plt.imsave("difference_BA.png", np.abs(image_B_grid - image_A_grid).astype(np.uint8)) + + return None diff --git a/lib/videoloader_imagenet.py b/lib/videoloader_imagenet.py new file mode 100644 index 0000000..12af761 --- /dev/null +++ b/lib/videoloader_imagenet.py @@ -0,0 +1,203 @@ +import os +import os.path as osp +import pdb +import struct +import time + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.utils.data as data +import torchvision.transforms as transforms +from PIL import Image +from scipy.misc import imread, imresize, imsave +from scipy.ndimage import geometric_transform, map_coordinates +from scipy.ndimage.filters import gaussian_filter +from skimage import color +from utils.util_distortion import ( + CenterPad, + CenterPadCrop_numpy, + Distortion, + Distortion_with_flow, + Normalize, + RGB2Lab, + ToTensor, +) + +import lib.functional as F +from lib.RefPairDatasetWithGlobalOrigRGB_optimize import ToTensor, combo5_loader, image_loader, parse_images + +cv2.setNumThreads(0) + + +class VideosDataset_ImageNet(data.Dataset): + def __init__( + self, + data_root, + epoch, + image_size, + with_bad=False, + with_mid=False, + transforms_imagenet=None, + distortion_level=3, + brightnessjitter=0, + nonzero_placeholder_probability=0.5, + extra_reference_transform=None, + real_reference_probability=1, + ): + image_pairs = [] + curr_image_pairs = parse_images(data_root, with_bad, with_mid) + image_pairs += curr_image_pairs + print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs))) + if len(image_pairs) == 0: + raise RuntimeError("Found 0 image_a pairs in all the data_roots") + + self.image_pairs = image_pairs + self.transforms_imagenet_raw = transforms_imagenet + self.extra_reference_transform = transforms.Compose(extra_reference_transform) + self.real_reference_probability = real_reference_probability + self.transforms_imagenet = transforms.Compose(transforms_imagenet) + self.epoch = epoch + self.image_size = image_size + self.real_len = len(self.image_pairs) + self.image_pairs *= epoch + self.distortion_level = distortion_level + self.distortion_transform = Distortion_with_flow() + self.brightnessjitter = brightnessjitter + self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()]) + self.nonzero_placeholder_probability = nonzero_placeholder_probability + self.ToTensor = ToTensor() + self.Normalize = Normalize() + + def __getitem__(self, index): + try: + # print('imagenet loader') + + image_id = 0 + pair_id = index + + combo_path = None + image_a_path = None + image_b_path = None + + image_names = ["", ""] + dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id] + sub_dir = osp.join(dir_root, cls_dir) + if is_good >= 1: + image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id]) + elif is_good == 0: + image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id]) + else: + image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id]) + image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id]) + + if np.random.random() > 0.5: + image_a_path, image_b_path = image_b_path, image_a_path + + I1 = image_loader(image_a_path) + I2 = I1 + I_reference_video = I1 + I_reference_video_real = image_loader(image_b_path) + + ## generate the flow + height, width = np.array(I2).shape[0], np.array(I2).shape[1] + alpha = np.random.rand() * self.distortion_level + distortion_range = 50 + random_state = np.random.RandomState(None) + shape = self.image_size[0], self.image_size[1] + # dx: flow on the vertical direction; dy: flow on the horizontal direction + forward_dx = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + forward_dy = ( + gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0) + * alpha + * 1000 + ) + + for transform in self.transforms_imagenet_raw: + if type(transform) is RGB2Lab: + I1_raw = I1 + I1 = transform(I1) + else: + I1 = transform(I1) + + for transform in self.transforms_imagenet_raw: + if type(transform) is RGB2Lab: + I2 = self.distortion_transform(I2, forward_dx, forward_dy) + I2_raw = I2 + I2 = transform(I2) + I2[0:1, :, :] = I2[0:1, :, :] + torch.randn(1) * self.brightnessjitter + + I_reference_video = self.extra_reference_transform(I_reference_video) + for transform in self.transforms_imagenet_raw: + # if type(transform) is RGB2Lab and self.extra_reference_transform is not None: + # I_reference_video = self.extra_reference_transform(I_reference_video) + I_reference_video = transform(I_reference_video) + + I_reference_video_real = self.transforms_imagenet(I_reference_video_real) + + flow_forward_raw = np.stack((forward_dy, forward_dx), axis=-1) + flow_backward_raw = np.zeros_like(flow_forward_raw) + flow_forward = self.flow_transform(flow_forward_raw) + flow_backward = self.flow_transform(flow_backward_raw) + + # update the mask for the pixels on the border + grid_x, grid_y = np.meshgrid(np.arange(self.image_size[0]), np.arange(self.image_size[1]), indexing="ij") + grid = np.stack((grid_y, grid_x), axis=-1) + grid_warp = grid + flow_forward_raw + location_y = grid_warp[:, :, 0].flatten() + location_x = grid_warp[:, :, 1].flatten() + I2_raw = np.array(I2_raw).astype(float) + I21_r = map_coordinates(I2_raw[:, :, 0], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_g = map_coordinates(I2_raw[:, :, 1], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_b = map_coordinates(I2_raw[:, :, 2], np.stack((location_x, location_y)), cval=-1).reshape( + (self.image_size[0], self.image_size[1]) + ) + I21_raw = np.stack((I21_r, I21_g, I21_b), axis=2) + mask = np.ones((self.image_size[0], self.image_size[1])) + mask[(I21_raw[:, :, 0] == -1) & (I21_raw[:, :, 1] == -1) & (I21_raw[:, :, 2] == -1)] = 0 + mask[abs(I21_raw - I1_raw).sum(axis=-1) > 50] = 0 + mask = self.ToTensor(mask) + + if np.random.random() < self.real_reference_probability: + I_reference_output = I_reference_video_real + placeholder = torch.zeros_like(I1) + self_ref_flag = torch.zeros_like(I1) + else: + I_reference_output = I_reference_video + placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1) + self_ref_flag = torch.ones_like(I1) + + outputs = [] + outputs.append(I1) + outputs.append(I2) + outputs.append(I_reference_output) + # outputs.append(I_reference_video_real) + outputs.append(flow_forward) + outputs.append(flow_backward) + outputs.append(mask) + outputs.append(placeholder) + outputs.append(self_ref_flag) + + except Exception as e: + if combo_path is not None: + print("problem in ", combo_path) + print("problem in, ", image_a_path) + print(e) + return self.__getitem__(np.random.randint(0, len(self.image_pairs))) + + return outputs + + def __len__(self): + return len(self.image_pairs) + # return len(self.image_pairs) * 2 diff --git a/models/ColorVidNet.py b/models/ColorVidNet.py new file mode 100644 index 0000000..b27792e --- /dev/null +++ b/models/ColorVidNet.py @@ -0,0 +1,265 @@ +from os import path as osp + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.parallel + + +def type_name(m): + name = type(m).__name__ + if name == "TorchObject": + name = m._typename.replace("cudnn.", "") + return name + + +def model_tm2pm(model): + layer_map = {} + for m in model.modules: + layer_map[m.name] = tm2pm(m) + return layer_map + + +"""' +removed layers: [[ 'ab_conv1_1', 'bw_conv1_1' ]] +""" + +ref_color_layernames = [ + "relu1_1", + "conv1_2", + "relu1_2", + "conv1_2norm", + "conv1_2norm_ss", + "conv2_1", + "relu2_1", + "conv2_2", + "relu2_2", + "conv2_2norm", + "conv2_2norm_ss", + "conv3_1", + "relu3_1", + "conv3_2", + "relu3_2", + "conv3_3", + "relu3_3", + "conv3_3norm", + "conv3_3norm_ss", + "conv4_1", + "relu4_1", + "conv4_2", + "relu4_2", + "conv4_3", + "relu4_3", + "conv4_3norm", + "conv5_1", + "relu5_1", + "conv5_2", + "relu5_2", + "conv5_3", + "relu5_3", + "conv5_3norm", + "conv6_1", + "relu6_1", + "conv6_2", + "relu6_2", + "conv6_3", + "relu6_3", + "conv6_3norm", + "conv7_1", + "relu7_1", + "conv7_2", + "relu7_2", + "conv7_3", + "relu7_3", + "conv7_3norm", + "conv8_1", + "conv3_3_short", + "relu8_1_comb", + "conv8_2", + "relu8_2", + "conv8_3", + "relu8_3", + "conv8_3norm", + "conv9_1", + "conv2_2_short", + "relu9_1_comb", + "conv9_2", + "relu9_2", + "conv9_2norm", + "conv10_1", + "conv1_2_short", + "relu10_1_comb", + "conv10_2", + "relu10_2", + "conv10_ab", +] + + +class ColorVidNet(nn.Module): + def __init__( + self, ic, ref_color_torch_model="convert_caffe_model/torch_model/reference_model/color_nodist_model.t7" + ): + super(ColorVidNet, self).__init__() + """torch_model = load_lua(ref_color_torch_model, unknown_classes=True) + if type(torch_model).__name__ == 'hashable_uniq_dict': + torch_model = torch_model.model + + # fix the ugly typo error + for m in torch_model.modules: + if m.name == 'conv8_1_comb' and type_name(m) == 'ReLU': + m.name = 'relu8_1_comb' + + pytorch_model = model_tm2pm(torch_model) + + self.add_module('conv1_1', nn.Sequential(nn.Conv2d(ic, 32, 3, 1, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 1, 1), nn.ReLU())) + + for l_name in ref_color_layernames: + self.add_module(l_name, pytorch_model[l_name]) + + ### replace all deconvolution with neareast upsample + convolution + assert type(self.conv8_1).__name__=='ConvTranspose2d' and type(self.conv9_1).__name__=='ConvTranspose2d' and \ + type(self.conv10_1).__name__=='ConvTranspose2d', 'conv8_1, conv9_1, conv10_1 should be all ConvTranspose2d' + """ + + ##### Written by Ak. ##### + self.conv1_1 = nn.Sequential(nn.Conv2d(ic, 32, 3, 1, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 1, 1)) + self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 1) + self.conv1_2norm = nn.BatchNorm2d(64, affine=False) + self.conv1_2norm_ss = nn.Conv2d(64, 64, 1, 2, bias=False, groups=64) + self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 1) + self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 1) + self.conv2_2norm = nn.BatchNorm2d(128, affine=False) + self.conv2_2norm_ss = nn.Conv2d(128, 128, 1, 2, bias=False, groups=128) + self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 1) + self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 1) + self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 1) + self.conv3_3norm = nn.BatchNorm2d(256, affine=False) + self.conv3_3norm_ss = nn.Conv2d(256, 256, 1, 2, bias=False, groups=256) + self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 1) + self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 1) + self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 1) + self.conv4_3norm = nn.BatchNorm2d(512, affine=False) + self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv5_2 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv5_3 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv5_3norm = nn.BatchNorm2d(512, affine=False) + self.conv6_1 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv6_2 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv6_3 = nn.Conv2d(512, 512, 3, 1, 2, 2) + self.conv6_3norm = nn.BatchNorm2d(512, affine=False) + self.conv7_1 = nn.Conv2d(512, 512, 3, 1, 1) + self.conv7_2 = nn.Conv2d(512, 512, 3, 1, 1) + self.conv7_3 = nn.Conv2d(512, 512, 3, 1, 1) + self.conv7_3norm = nn.BatchNorm2d(512, affine=False) + self.conv8_1 = nn.ConvTranspose2d(512, 256, 4, 2, 1) + self.conv3_3_short = nn.Conv2d(256, 256, 3, 1, 1) + self.conv8_2 = nn.Conv2d(256, 256, 3, 1, 1) + self.conv8_3 = nn.Conv2d(256, 256, 3, 1, 1) + self.conv8_3norm = nn.BatchNorm2d(256, affine=False) + self.conv9_1 = nn.ConvTranspose2d(256, 128, 4, 2, 1) + self.conv2_2_short = nn.Conv2d(128, 128, 3, 1, 1) + self.conv9_2 = nn.Conv2d(128, 128, 3, 1, 1) + self.conv9_2norm = nn.BatchNorm2d(128, affine=False) + self.conv10_1 = nn.ConvTranspose2d(128, 128, 4, 2, 1) + self.conv1_2_short = nn.Conv2d(64, 128, 3, 1, 1) + self.conv10_2 = nn.Conv2d(128, 128, 3, 1, 1) + self.conv10_ab = nn.Conv2d(128, 2, 1, 1) + + # add self.relux_x + self.relu1_1 = nn.ReLU() + self.relu1_2 = nn.ReLU() + self.relu2_1 = nn.ReLU() + self.relu2_2 = nn.ReLU() + self.relu3_1 = nn.ReLU() + self.relu3_2 = nn.ReLU() + self.relu3_3 = nn.ReLU() + self.relu4_1 = nn.ReLU() + self.relu4_2 = nn.ReLU() + self.relu4_3 = nn.ReLU() + self.relu5_1 = nn.ReLU() + self.relu5_2 = nn.ReLU() + self.relu5_3 = nn.ReLU() + self.relu6_1 = nn.ReLU() + self.relu6_2 = nn.ReLU() + self.relu6_3 = nn.ReLU() + self.relu7_1 = nn.ReLU() + self.relu7_2 = nn.ReLU() + self.relu7_3 = nn.ReLU() + self.relu8_1_comb = nn.ReLU() + self.relu8_2 = nn.ReLU() + self.relu8_3 = nn.ReLU() + self.relu9_1_comb = nn.ReLU() + self.relu9_2 = nn.ReLU() + self.relu10_1_comb = nn.ReLU() + # self.relu10_2 = nn.ReLU() + self.relu10_2 = nn.LeakyReLU(0.2, True) + + ##### Above ##### + + print("replace all deconv with [nearest + conv]") + self.conv8_1 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(512, 256, 3, 1, 1)) + self.conv9_1 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(256, 128, 3, 1, 1)) + self.conv10_1 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(128, 128, 3, 1, 1)) + + print("replace all batchnorm with instancenorm") + self.conv1_2norm = nn.InstanceNorm2d(64) + self.conv2_2norm = nn.InstanceNorm2d(128) + self.conv3_3norm = nn.InstanceNorm2d(256) + self.conv4_3norm = nn.InstanceNorm2d(512) + self.conv5_3norm = nn.InstanceNorm2d(512) + self.conv6_3norm = nn.InstanceNorm2d(512) + self.conv7_3norm = nn.InstanceNorm2d(512) + self.conv8_3norm = nn.InstanceNorm2d(256) + self.conv9_2norm = nn.InstanceNorm2d(128) + # print(self) + + def forward(self, x): + """ x: gray image (1 channel), ab(2 channel), ab_err, ba_err""" + conv1_1 = self.relu1_1(self.conv1_1(x)) + conv1_2 = self.relu1_2(self.conv1_2(conv1_1)) + conv1_2norm = self.conv1_2norm(conv1_2) + conv1_2norm_ss = self.conv1_2norm_ss(conv1_2norm) + conv2_1 = self.relu2_1(self.conv2_1(conv1_2norm_ss)) + conv2_2 = self.relu2_2(self.conv2_2(conv2_1)) + conv2_2norm = self.conv2_2norm(conv2_2) + conv2_2norm_ss = self.conv2_2norm_ss(conv2_2norm) + conv3_1 = self.relu3_1(self.conv3_1(conv2_2norm_ss)) + conv3_2 = self.relu3_2(self.conv3_2(conv3_1)) + conv3_3 = self.relu3_3(self.conv3_3(conv3_2)) + conv3_3norm = self.conv3_3norm(conv3_3) + conv3_3norm_ss = self.conv3_3norm_ss(conv3_3norm) + conv4_1 = self.relu4_1(self.conv4_1(conv3_3norm_ss)) + conv4_2 = self.relu4_2(self.conv4_2(conv4_1)) + conv4_3 = self.relu4_3(self.conv4_3(conv4_2)) + conv4_3norm = self.conv4_3norm(conv4_3) + conv5_1 = self.relu5_1(self.conv5_1(conv4_3norm)) + conv5_2 = self.relu5_2(self.conv5_2(conv5_1)) + conv5_3 = self.relu5_3(self.conv5_3(conv5_2)) + conv5_3norm = self.conv5_3norm(conv5_3) + conv6_1 = self.relu6_1(self.conv6_1(conv5_3norm)) + conv6_2 = self.relu6_2(self.conv6_2(conv6_1)) + conv6_3 = self.relu6_3(self.conv6_3(conv6_2)) + conv6_3norm = self.conv6_3norm(conv6_3) + conv7_1 = self.relu7_1(self.conv7_1(conv6_3norm)) + conv7_2 = self.relu7_2(self.conv7_2(conv7_1)) + conv7_3 = self.relu7_3(self.conv7_3(conv7_2)) + conv7_3norm = self.conv7_3norm(conv7_3) + conv8_1 = self.conv8_1(conv7_3norm) + conv3_3_short = self.conv3_3_short(conv3_3norm) + conv8_1_comb = self.relu8_1_comb(conv8_1 + conv3_3_short) + conv8_2 = self.relu8_2(self.conv8_2(conv8_1_comb)) + conv8_3 = self.relu8_3(self.conv8_3(conv8_2)) + conv8_3norm = self.conv8_3norm(conv8_3) + conv9_1 = self.conv9_1(conv8_3norm) + conv2_2_short = self.conv2_2_short(conv2_2norm) + conv9_1_comb = self.relu9_1_comb(conv9_1 + conv2_2_short) + conv9_2 = self.relu9_2(self.conv9_2(conv9_1_comb)) + conv9_2norm = self.conv9_2norm(conv9_2) + conv10_1 = self.conv10_1(conv9_2norm) + conv1_2_short = self.conv1_2_short(conv1_2norm) + conv10_1_comb = self.relu10_1_comb(conv10_1 + conv1_2_short) + conv10_2 = self.relu10_2(self.conv10_2(conv10_1_comb)) + conv10_ab = self.conv10_ab(conv10_2) + pred_ab = torch.tanh(conv10_ab) * 128 + + return pred_ab diff --git a/models/FrameColor.py b/models/FrameColor.py new file mode 100644 index 0000000..ba2acbf --- /dev/null +++ b/models/FrameColor.py @@ -0,0 +1,80 @@ +# from imports import * +import torch +from utils.tensor_lab2rgb import tensor_lab2rgb +from utils.util import * + +if torch.cuda.is_available(): + device = torch.device("cuda") +else: + device = troch.device("cpu") + + +def warp_color(IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise=0, temperature=0.01): + """ + joint_train=True, enable grad; otherwise disable grad + """ + # change to rgb for feature extraction + IA_rgb_from_gray = gray2rgb_batch(IA_l) + + with torch.no_grad(): + A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet( + IA_rgb_from_gray, ["r12", "r22", "r32", "r42", "r52"], preprocess=True + ) + B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B + + # NOTE: output the feature before normalization + features_A = [A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1] + + A_relu2_1 = feature_normalize(A_relu2_1) + A_relu3_1 = feature_normalize(A_relu3_1) + A_relu4_1 = feature_normalize(A_relu4_1) + A_relu5_1 = feature_normalize(A_relu5_1) + B_relu2_1 = feature_normalize(B_relu2_1) + B_relu3_1 = feature_normalize(B_relu3_1) + B_relu4_1 = feature_normalize(B_relu4_1) + B_relu5_1 = feature_normalize(B_relu5_1) + + # if feature_noise > 0: + # A_relu2_1 = A_relu2_1 + torch.randn_like(A_relu2_1, requires_grad=False) * feature_noise + # B_relu2_1 = B_relu2_1 + torch.randn_like(B_relu2_1, requires_grad=False) * feature_noise + # A_relu3_1 = A_relu3_1 + torch.randn_like(A_relu3_1, requires_grad=False) * feature_noise + # B_relu3_1 = B_relu3_1 + torch.randn_like(B_relu3_1, requires_grad=False) * feature_noise + # A_relu4_1 = A_relu4_1 + torch.randn_like(A_relu4_1, requires_grad=False) * feature_noise + # B_relu4_1 = B_relu4_1 + torch.randn_like(B_relu4_1, requires_grad=False) * feature_noise + # A_relu5_1 = A_relu5_1 + torch.randn_like(A_relu5_1, requires_grad=False) * feature_noise + # B_relu5_1 = B_relu5_1 + torch.randn_like(B_relu5_1, requires_grad=False) * feature_noise + + nonlocal_BA_lab, similarity_map = nonlocal_net( + IB_lab, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1, temperature=0.01 + ) + + return nonlocal_BA_lab, similarity_map, features_A + + +def frame_colorization( + IA_lab, + IB_lab, + IA_last_lab, + features_B, + vggnet, + nonlocal_net, + colornet, + joint_training=True, + feature_noise=0, + luminance_noise=0, + temperature=0.01, +): + # change to rgb for feature extraction + IA_l = IA_lab[:, 0:1, :, :] + if luminance_noise: + IA_l = IA_l + torch.randn_like(IA_l, requires_grad=False) * luminance_noise + + with torch.autograd.set_grad_enabled(joint_training): + nonlocal_BA_lab, similarity_map, features_A_gray = warp_color( + IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise, temperature=temperature + ) + nonlocal_BA_ab = nonlocal_BA_lab[:, 1:3, :, :] + color_input = torch.cat((IA_l, nonlocal_BA_ab, similarity_map, IA_last_lab), dim=1) + IA_ab_predict = colornet(color_input) + + return IA_ab_predict, nonlocal_BA_lab, features_A_gray diff --git a/models/GAN_models.py b/models/GAN_models.py new file mode 100644 index 0000000..024b3a3 --- /dev/null +++ b/models/GAN_models.py @@ -0,0 +1,413 @@ +# DCGAN-like generator and discriminator +import torch +import torch.nn.functional as F +from torch import nn + +from models.spectral_normalization import SpectralNorm + + +class Generator(nn.Module): + def __init__(self, z_dim): + super(Generator, self).__init__() + self.z_dim = z_dim + + self.model = nn.Sequential( + nn.ConvTranspose2d(z_dim, 512, 4, stride=1), + nn.InstanceNorm2d(512), + nn.ReLU(), + nn.ConvTranspose2d(512, 256, 4, stride=2, padding=(1, 1)), + nn.InstanceNorm2d(256), + nn.ReLU(), + nn.ConvTranspose2d(256, 128, 4, stride=2, padding=(1, 1)), + nn.InstanceNorm2d(128), + nn.ReLU(), + nn.ConvTranspose2d(128, 64, 4, stride=2, padding=(1, 1)), + nn.InstanceNorm2d(64), + nn.ReLU(), + nn.ConvTranspose2d(64, channels, 3, stride=1, padding=(1, 1)), + nn.Tanh(), + ) + + def forward(self, z): + return self.model(z.view(-1, self.z_dim, 1, 1)) + + +channels = 3 +leak = 0.1 +w_g = 4 + + +class Discriminator(nn.Module): + def __init__(self): + super(Discriminator, self).__init__() + + self.conv1 = SpectralNorm(nn.Conv2d(channels, 64, 3, stride=1, padding=(1, 1))) + self.conv2 = SpectralNorm(nn.Conv2d(64, 64, 4, stride=2, padding=(1, 1))) + self.conv3 = SpectralNorm(nn.Conv2d(64, 128, 3, stride=1, padding=(1, 1))) + self.conv4 = SpectralNorm(nn.Conv2d(128, 128, 4, stride=2, padding=(1, 1))) + self.conv5 = SpectralNorm(nn.Conv2d(128, 256, 3, stride=1, padding=(1, 1))) + self.conv6 = SpectralNorm(nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1))) + self.conv7 = SpectralNorm(nn.Conv2d(256, 256, 3, stride=1, padding=(1, 1))) + self.conv8 = SpectralNorm(nn.Conv2d(256, 512, 4, stride=2, padding=(1, 1))) + self.fc = SpectralNorm(nn.Linear(w_g * w_g * 512, 1)) + + def forward(self, x): + m = x + m = nn.LeakyReLU(leak)(self.conv1(m)) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(64)(self.conv2(m))) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(128)(self.conv3(m))) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(128)(self.conv4(m))) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(256)(self.conv5(m))) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(256)(self.conv6(m))) + m = nn.LeakyReLU(leak)(nn.InstanceNorm2d(256)(self.conv7(m))) + m = nn.LeakyReLU(leak)(self.conv8(m)) + + return self.fc(m.view(-1, w_g * w_g * 512)) + + +class Self_Attention(nn.Module): + """ Self attention Layer""" + + def __init__(self, in_dim): + super(Self_Attention, self).__init__() + self.chanel_in = in_dim + + self.query_conv = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 1, kernel_size=1)) + self.key_conv = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 1, kernel_size=1)) + self.value_conv = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)) + self.gamma = nn.Parameter(torch.zeros(1)) + + self.softmax = nn.Softmax(dim=-1) # + + def forward(self, x): + """ + inputs : + x : input feature maps( B X C X W X H) + returns : + out : self attention value + input feature + attention: B X N X N (N is Width*Height) + """ + m_batchsize, C, width, height = x.size() + proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1) # B X CX(N) + proj_key = self.key_conv(x).view(m_batchsize, -1, width * height) # B X C x (*W*H) + energy = torch.bmm(proj_query, proj_key) # transpose check + attention = self.softmax(energy) # BX (N) X (N) + proj_value = self.value_conv(x).view(m_batchsize, -1, width * height) # B X C X N + + out = torch.bmm(proj_value, attention.permute(0, 2, 1)) + out = out.view(m_batchsize, C, width, height) + + out = self.gamma * out + x + # return out, attention + return out + + +class Discriminator_WITHOUT_FC_x64_video(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=6, ndf=64): + super(Discriminator_WITHOUT_FC_x64_video, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.layer1 = nn.Sequential( + # input size is in_size x 64 x 64 + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer2 = nn.Sequential( + # state size: ndf x 32 x 32 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer3 = nn.Sequential( + # state size: (ndf * 2) x 16 x 16 + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer4 = nn.Sequential( + # state size: (ndf * 4) x 8 x 8 + # Self_Attention(self.ndf * 4, 'relu'), + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + ) + + # sate size: (ndf * 8) x 4 x 4 + # Self_Attention(self.ndf * 8, 'relu') + + self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0)) + # nn.Sigmoid() + # state size: 1 x 1 x 1 + + # for m in self.modules(): + # if isinstance(m, nn.Conv2d): + # m.weight.data.normal_(0.0, 0.02) + # if m.bias is not None: + # m.bias.data.zero_() + + def forward(self, input): + feature1 = self.layer1(input) + feature2 = self.layer2(feature1) + feature3 = self.layer3(feature2) + feature4 = self.layer4(feature3) + output = self.last(feature4) + output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) + + return output, feature4 + + +class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=6, ndf=64): + super(Discriminator_WITHOUT_FC_x64_video_BIG, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.layer1 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True) + ) + self.layer2 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf), + nn.LeakyReLU(0.2, inplace=True), + ) + self.attention = Self_Attention(self.ndf) + self.layer3 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer4 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer5 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + ) + self.layer6 = nn.Sequential( + SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 16), + nn.LeakyReLU(0.2, inplace=True), + ) + + self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, [3, 6], 1, 0)) + + def forward(self, input): + feature1 = self.layer1(input) + feature2 = self.layer2(feature1) + feature_attention = self.attention(feature2) + feature3 = self.layer3(feature_attention) + feature4 = self.layer4(feature3) + feature5 = self.layer5(feature4) + feature6 = self.layer6(feature5) + output = self.last(feature6) + output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) + + return output, feature4 + + +class Discriminator_WITHOUT_FC_x64(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=3, ndf=64): + super(Discriminator_WITHOUT_FC_x64, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.main = nn.Sequential( + # input size is in_size x 64 x 64 + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), + nn.LeakyReLU(0.2, inplace=True), + # state size: ndf x 32 x 32 + Self_Attention(self.ndf), + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + # state size: (ndf * 2) x 16 x 16 + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + # state size: (ndf * 4) x 8 x 8 + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + # sate size: (ndf * 8) x 4 x 4 + ) + + self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) + # nn.Sigmoid() + # state size: 1 x 1 x 1 + + # for m in self.modules(): + # if isinstance(m, nn.Conv2d): + # m.weight.data.normal_(0.0, 0.02) + # if m.bias is not None: + # m.bias.data.zero_() + + def forward(self, input): + + feature = self.main(input) + output = self.last(feature) + output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) + + return output, feature + + +class Discriminator_WITHOUT_FC_x64_BIG(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=3, ndf=64): + super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.main = nn.Sequential( + # 256 X 256 + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), + nn.LeakyReLU(0.2, inplace=True), + # 128 X 128 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf), + nn.LeakyReLU(0.2, inplace=True), + # 64 X 64 + Self_Attention(self.ndf), + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + # 32 X 32 + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + # 16 X 16 + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + # 8 X 8 + SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 16), + nn.LeakyReLU(0.2, inplace=True), + ) + # 4 X 4 + self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0)) + + def forward(self, input): + + feature = self.main(input) + output = self.last(feature) + output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1) + + return output, feature + + +class Discriminator_WITHOUT_FC_x128(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=3, ndf=64): + super(Discriminator_WITHOUT_FC_x128, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.main = nn.Sequential( + # input size is in_size x 128 x 128 + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), + nn.LeakyReLU(0.2, inplace=True), + # state size: ndf x 64 x 64 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf), + nn.LeakyReLU(0.2, inplace=True), + # state size: ndf x 32 x 32 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + # state size: (ndf * 24) x 16 x 16 + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + # sate size: (ndf * 4) x 8 x 8 + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + # sate size: (ndf * 8) x 4 x 4 + ) + + self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) + + def forward(self, input): + + feature = self.main(input) + output = self.last(feature) + + return output, feature + + +class Discriminator_WITHOUT_FC_x256(nn.Module): + """ + Discriminative Network + """ + + def __init__(self, in_size=3, ndf=64): + super(Discriminator_WITHOUT_FC_x256, self).__init__() + self.in_size = in_size + self.ndf = ndf + + self.main = nn.Sequential( + # input size is in_size x 256 x 256 + SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)), + nn.LeakyReLU(0.2, inplace=True), + # state size: ndf x 128 x 128 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf), + nn.LeakyReLU(0.2, inplace=True), + # state size: ndf x 64 x 64 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf), + nn.LeakyReLU(0.2, inplace=True), + # state size: (ndf * 1) x 32 x 32 + SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + # state size: (ndf * 2) x 16 x 16 + SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + # sate size: (ndf * 4) x 8 x 8 + SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)), + nn.InstanceNorm2d(self.ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + # sate size: (ndf * 8) x 4 x 4 + ) + + self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0)) + + def forward(self, input): + + feature = self.main(input) + output = self.last(feature) + + return output, feature + + +if __name__ == "__main__": + discriminator1 = Discriminator_WITHOUT_FC_x64() + discriminator2 = Discriminator_WITHOUT_FC_x128() + discriminator3 = Discriminator_WITHOUT_FC_x256() + print(discriminator1) + print(discriminator2) + print(discriminator3) diff --git a/models/NonlocalNet.py b/models/NonlocalNet.py new file mode 100644 index 0000000..d15cdb1 --- /dev/null +++ b/models/NonlocalNet.py @@ -0,0 +1,571 @@ +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import models as torch_models +from utils.util import uncenter_l, vgg_preprocess + +from models.vgg19_gray import vgg19_gray, vgg19_gray_new + + +def find_local_patch(x, patch_size): + N, C, H, W = x.shape + x_unfold = F.unfold( + x, kernel_size=(patch_size, patch_size), padding=(patch_size // 2, patch_size // 2), stride=(1, 1) + ) + out = x_unfold.view(N, x_unfold.shape[1], H, W) + return out + + +class WeightedAverage(nn.Module): + def __init__(self,): + super(WeightedAverage, self).__init__() + + def forward(self, x_lab, patch_size=3, alpha=1, scale_factor=1): + # alpha=0: less smooth; alpha=inf: smoother + x_lab = F.interpolate(x_lab, scale_factor=scale_factor) + l = x_lab[:, 0:1, :, :] + a = x_lab[:, 1:2, :, :] + b = x_lab[:, 2:3, :, :] + local_l = find_local_patch(l, patch_size) + local_a = find_local_patch(a, patch_size) + local_b = find_local_patch(b, patch_size) + local_difference_l = (local_l - l) ** 2 + correlation = nn.functional.softmax(-1 * local_difference_l / alpha, dim=1) + + weighted_ab = torch.cat( + ( + torch.sum(correlation * local_a, dim=1, keepdim=True), + torch.sum(correlation * local_b, dim=1, keepdim=True), + ), + 1, + ) + return weighted_ab + + +class WeightedAverage_color(nn.Module): + """ + smooth the image according to the color distance in the LAB space + """ + + def __init__(self,): + super(WeightedAverage_color, self).__init__() + + def forward(self, x_lab, x_lab_predict, patch_size=3, alpha=1, scale_factor=1): + # alpha=0: less smooth; alpha=inf: smoother + x_lab = F.interpolate(x_lab, scale_factor=scale_factor) + l = uncenter_l(x_lab[:, 0:1, :, :]) + a = x_lab[:, 1:2, :, :] + b = x_lab[:, 2:3, :, :] + a_predict = x_lab_predict[:, 1:2, :, :] + b_predict = x_lab_predict[:, 2:3, :, :] + local_l = find_local_patch(l, patch_size) + local_a = find_local_patch(a, patch_size) + local_b = find_local_patch(b, patch_size) + local_a_predict = find_local_patch(a_predict, patch_size) + local_b_predict = find_local_patch(b_predict, patch_size) + + local_color_difference = (local_l - l) ** 2 + (local_a - a) ** 2 + (local_b - b) ** 2 + correlation = nn.functional.softmax( + -1 * local_color_difference / alpha, dim=1 + ) # so that sum of weights equal to 1 + + weighted_ab = torch.cat( + ( + torch.sum(correlation * local_a_predict, dim=1, keepdim=True), + torch.sum(correlation * local_b_predict, dim=1, keepdim=True), + ), + 1, + ) + return weighted_ab + + +class NonlocalWeightedAverage(nn.Module): + def __init__(self,): + super(NonlocalWeightedAverage, self).__init__() + + def forward(self, x_lab, feature, patch_size=3, alpha=0.1, scale_factor=1): + # alpha=0: less smooth; alpha=inf: smoother + # input feature is normalized feature + x_lab = F.interpolate(x_lab, scale_factor=scale_factor) + batch_size, channel, height, width = x_lab.shape + feature = F.interpolate(feature, size=(height, width)) + batch_size = x_lab.shape[0] + x_ab = x_lab[:, 1:3, :, :].view(batch_size, 2, -1) + x_ab = x_ab.permute(0, 2, 1) + + local_feature = find_local_patch(feature, patch_size) + local_feature = local_feature.view(batch_size, local_feature.shape[1], -1) + + correlation_matrix = torch.matmul(local_feature.permute(0, 2, 1), local_feature) + correlation_matrix = nn.functional.softmax(correlation_matrix / alpha, dim=-1) + + weighted_ab = torch.matmul(correlation_matrix, x_ab) + weighted_ab = weighted_ab.permute(0, 2, 1).contiguous() + weighted_ab = weighted_ab.view(batch_size, 2, height, width) + return weighted_ab + + +class CorrelationLayer(nn.Module): + def __init__(self, search_range): + super(CorrelationLayer, self).__init__() + self.search_range = search_range + + def forward(self, x1, x2, alpha=1, raw_output=False, metric="similarity"): + # args = self.args + + shape = list(x1.size()) + shape[1] = (self.search_range * 2 + 1) ** 2 + cv = torch.zeros(shape).to(torch.device("cuda")) + + if metric == "similarity": + for i in range(-self.search_range, self.search_range + 1): + for j in range(-self.search_range, self.search_range + 1): + if i < 0: + slice_h, slice_h_r = slice(None, i), slice(-i, None) + elif i > 0: + slice_h, slice_h_r = slice(i, None), slice(None, -i) + else: + slice_h, slice_h_r = slice(None), slice(None) + + if j < 0: + slice_w, slice_w_r = slice(None, j), slice(-j, None) + elif j > 0: + slice_w, slice_w_r = slice(j, None), slice(None, -j) + else: + slice_w, slice_w_r = slice(None), slice(None) + + # storage sequence (eg. search_range=3): -24, -23, .., -1, 0, 1, ..., 23, 24 + cv[:, (self.search_range * 2 + 1) * i + j, slice_h, slice_w] = ( + x1[:, :, slice_h, slice_w] * x2[:, :, slice_h_r, slice_w_r] + ).sum(1) + else: # patchwise subtraction + for i in range(-self.search_range, self.search_range + 1): + for j in range(-self.search_range, self.search_range + 1): + if i < 0: + slice_h, slice_h_r = slice(None, i), slice(-i, None) + elif i > 0: + slice_h, slice_h_r = slice(i, None), slice(None, -i) + else: + slice_h, slice_h_r = slice(None), slice(None) + + if j < 0: + slice_w, slice_w_r = slice(None, j), slice(-j, None) + elif j > 0: + slice_w, slice_w_r = slice(j, None), slice(None, -j) + else: + slice_w, slice_w_r = slice(None), slice(None) + + # storage sequence (eg. search_range=3): -24, -23, .., -1, 0, 1, ..., 23, 24 + cv[:, (self.search_range * 2 + 1) * i + j, slice_h, slice_w] = -( + (x1[:, :, slice_h, slice_w] - x2[:, :, slice_h_r, slice_w_r]) ** 2 + ).sum(1) + + ## TODO sigmoid? + if raw_output: + return cv + else: + return nn.functional.softmax(cv / alpha, dim=1) + + +class Self_Attn(nn.Module): + """ Self attention Layer""" + + def __init__(self, in_dim, activation): + super(Self_Attn, self).__init__() + self.chanel_in = in_dim + self.activation = activation + + self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 4, kernel_size=1) + self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 4, kernel_size=1) + self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1) + self.gamma = nn.Parameter(torch.zeros(1)) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x): + """ + inputs : + x : input feature maps(B X C X W X H) + returns : + out : self attention value + input feature + attention: B X N X N (N is Width*Height) + """ + m_batchsize, C, width, height = x.size() + proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1) # B X N X C + proj_key = self.key_conv(x).view(m_batchsize, -1, width * height) # B X C x N + energy = torch.bmm(proj_query, proj_key) # transpose check + attention = self.softmax(energy) # B X (N) X (N) + proj_value = self.value_conv(x).view(m_batchsize, -1, width * height) # B X C X N + + out = torch.bmm(proj_value, attention.permute(0, 1, 2)) + out = out.view(m_batchsize, C, width, height) + + out = self.gamma * out + x + return out + + +class VGG19_feature_color_torchversion(nn.Module): + """ + NOTE: no need to pre-process the input; input tensor should range in [0,1] + """ + + def __init__(self, pool="max"): + super(VGG19_feature_color_torchversion, self).__init__() + self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) + self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) + self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1) + self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1) + self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.conv3_4 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1) + self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv4_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + self.conv5_4 = nn.Conv2d(512, 512, kernel_size=3, padding=1) + if pool == "max": + self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) + self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) + self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) + self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) + self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2) + elif pool == "avg": + self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2) + self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2) + self.pool3 = nn.AvgPool2d(kernel_size=2, stride=2) + self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2) + self.pool5 = nn.AvgPool2d(kernel_size=2, stride=2) + + def forward(self, x, out_keys, preprocess=True): + """ + NOTE: input tensor should range in [0,1] + """ + out = {} + if preprocess: + x = vgg_preprocess(x) + out["r11"] = F.relu(self.conv1_1(x)) + out["r12"] = F.relu(self.conv1_2(out["r11"])) + out["p1"] = self.pool1(out["r12"]) + out["r21"] = F.relu(self.conv2_1(out["p1"])) + out["r22"] = F.relu(self.conv2_2(out["r21"])) + out["p2"] = self.pool2(out["r22"]) + out["r31"] = F.relu(self.conv3_1(out["p2"])) + out["r32"] = F.relu(self.conv3_2(out["r31"])) + out["r33"] = F.relu(self.conv3_3(out["r32"])) + out["r34"] = F.relu(self.conv3_4(out["r33"])) + out["p3"] = self.pool3(out["r34"]) + out["r41"] = F.relu(self.conv4_1(out["p3"])) + out["r42"] = F.relu(self.conv4_2(out["r41"])) + out["r43"] = F.relu(self.conv4_3(out["r42"])) + out["r44"] = F.relu(self.conv4_4(out["r43"])) + out["p4"] = self.pool4(out["r44"]) + out["r51"] = F.relu(self.conv5_1(out["p4"])) + out["r52"] = F.relu(self.conv5_2(out["r51"])) + out["r53"] = F.relu(self.conv5_3(out["r52"])) + out["r54"] = F.relu(self.conv5_4(out["r53"])) + out["p5"] = self.pool5(out["r54"]) + return [out[key] for key in out_keys] + + +class VGG19_feature_color(nn.Module): + def __init__(self): + super(VGG19_feature_color, self).__init__() + # self.select = ['0', '5', '10', '19', '28'] # Select conv1_1 ~ conv5_1 activation maps. + self.select = ["1", "6", "11", "20", "29"] # Select relu1_1 ~ relu5_1 activation maps. + self.vgg = torch_models.vgg19(pretrained=True).features + + def forward(self, x): + """Extract multiple convolutional feature maps.""" + features = [] + for name, layer in self.vgg._modules.items(): + x = layer(x) + if name in self.select: + features.append(x) + return features + + +class VGG19_feature(nn.Module): + # input: [LLL] channels, range=[0,255] + def __init__(self, gpu_ids): + super(VGG19_feature, self).__init__() + self.vgg19_gray = vgg19_gray().cuda() + + def forward(self, A_l, B_l): + A_relu3_1, A_relu4_1, A_relu5_1 = self.vgg19_gray(A_l) + B_relu3_1, B_relu4_1, B_relu5_1 = self.vgg19_gray(B_l) + return A_relu3_1, A_relu4_1, A_relu5_1, B_relu3_1, B_relu4_1, B_relu5_1 + + +class VGG19_feature_new(nn.Module): + # input: [LLL] channels, range=[0,255] + def __init__(self, gpu_ids): + super(VGG19_feature_new, self).__init__() + self.vgg19_gray_new = vgg19_gray_new().cuda() + + def forward(self, A_l, B_l): + A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = self.vgg19_gray_new(A_l) + B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = self.vgg19_gray_new(B_l) + return A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 + + +class WTA(torch.autograd.Function): + """ + We can implement our own custom autograd Functions by subclassing + torch.autograd.Function and implementing the forward and backward passes + which operate on Tensors. + """ + + @staticmethod + def forward(ctx, input): + """ + In the forward pass we receive a Tensor containing the input and return a + Tensor containing the output. You can cache arbitrary Tensors for use in the + backward pass using the save_for_backward method. + """ + height = input.shape[2] + width = input.shape[3] + activation_max, index_max = torch.topk(input, -1, keepdim=True) + input_zeros = input * 0 + output_max_only = torch.where(input == activation_max, input, input_zeros) + + mask = (output_max_only > 0).type(torch.float) + ctx.save_for_backward(input, mask) + return output_max_only + + @staticmethod + def backward(ctx, grad_output): + """ + In the backward pass we receive a Tensor containing the gradient of the loss + with respect to the output, and we need to compute the gradient of the loss + with respect to the input. + """ + # import pdb + # pdb.set_trace() + input, mask = ctx.saved_tensors + grad_input = grad_output.clone() * mask + return grad_input + + +class WTA_scale(torch.autograd.Function): + """ + We can implement our own custom autograd Functions by subclassing + torch.autograd.Function and implementing the forward and backward passes + which operate on Tensors. + """ + + @staticmethod + def forward(ctx, input, scale=1e-4): + """ + In the forward pass we receive a Tensor containing the input and return a + Tensor containing the output. You can cache arbitrary Tensors for use in the + backward pass using the save_for_backward method. + """ + activation_max, index_max = torch.max(input, -1, keepdim=True) + input_scale = input * scale # default: 1e-4 + # input_scale = input * scale # default: 1e-4 + output_max_scale = torch.where(input == activation_max, input, input_scale) + + mask = (input == activation_max).type(torch.float) + ctx.save_for_backward(input, mask) + return output_max_scale + + @staticmethod + def backward(ctx, grad_output): + """ + In the backward pass we receive a Tensor containing the gradient of the loss + with respect to the output, and we need to compute the gradient of the loss + with respect to the input. + """ + # import pdb + # pdb.set_trace() + input, mask = ctx.saved_tensors + mask_ones = torch.ones_like(mask) + mask_small_ones = torch.ones_like(mask) * 1e-4 + # mask_small_ones = torch.ones_like(mask) * 1e-4 + + grad_scale = torch.where(mask == 1, mask_ones, mask_small_ones) + grad_input = grad_output.clone() * grad_scale + return grad_input, None + + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, stride=1): + super(ResidualBlock, self).__init__() + self.padding1 = nn.ReflectionPad2d(padding) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=0, stride=stride) + self.bn1 = nn.InstanceNorm2d(out_channels) + self.prelu = nn.PReLU() + self.padding2 = nn.ReflectionPad2d(padding) + self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=0, stride=stride) + self.bn2 = nn.InstanceNorm2d(out_channels) + + def forward(self, x): + residual = x + out = self.padding1(x) + out = self.conv1(out) + out = self.bn1(out) + out = self.prelu(out) + out = self.padding2(out) + out = self.conv2(out) + out = self.bn2(out) + out += residual + out = self.prelu(out) + return out + + +## NOTE: THIS IS THE FINAL VERSION +class WarpNet(nn.Module): + # input is Al, Bl, channel = 1, range~[0,255] + def __init__(self, batch_size): + super(WarpNet, self).__init__() + self.feature_channel = 64 + self.in_channels = self.feature_channel * 4 + self.inter_channels = 256 + # 44*44 + self.layer2_1 = nn.Sequential( + nn.ReflectionPad2d(1), + nn.Conv2d(128, 128, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(128), + nn.PReLU(), + nn.ReflectionPad2d(1), + nn.Conv2d(128, self.feature_channel, kernel_size=3, padding=0, stride=2), + nn.InstanceNorm2d(self.feature_channel), + nn.PReLU(), + ) + self.layer3_1 = nn.Sequential( + nn.ReflectionPad2d(1), + nn.Conv2d(256, 128, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(128), + nn.PReLU(), + nn.ReflectionPad2d(1), + nn.Conv2d(128, self.feature_channel, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(self.feature_channel), + nn.PReLU(), + ) + + # 22*22->44*44 + self.layer4_1 = nn.Sequential( + nn.ReflectionPad2d(1), + nn.Conv2d(512, 256, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(256), + nn.PReLU(), + nn.ReflectionPad2d(1), + nn.Conv2d(256, self.feature_channel, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(self.feature_channel), + nn.PReLU(), + nn.Upsample(scale_factor=2), + ) + + # 11*11->44*44 + self.layer5_1 = nn.Sequential( + nn.ReflectionPad2d(1), + nn.Conv2d(512, 256, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(256), + nn.PReLU(), + nn.Upsample(scale_factor=2), + nn.ReflectionPad2d(1), + nn.Conv2d(256, self.feature_channel, kernel_size=3, padding=0, stride=1), + nn.InstanceNorm2d(self.feature_channel), + nn.PReLU(), + nn.Upsample(scale_factor=2), + ) + + self.layer = nn.Sequential( + ResidualBlock(self.feature_channel * 4, self.feature_channel * 4, kernel_size=3, padding=1, stride=1), + ResidualBlock(self.feature_channel * 4, self.feature_channel * 4, kernel_size=3, padding=1, stride=1), + ResidualBlock(self.feature_channel * 4, self.feature_channel * 4, kernel_size=3, padding=1, stride=1), + ) + + self.theta = nn.Conv2d( + in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0 + ) + self.phi = nn.Conv2d( + in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0 + ) + + self.upsampling = nn.Upsample(scale_factor=4) + + def forward( + self, + B_lab_map, + A_relu2_1, + A_relu3_1, + A_relu4_1, + A_relu5_1, + B_relu2_1, + B_relu3_1, + B_relu4_1, + B_relu5_1, + temperature=0.001 * 5, + detach_flag=False, + WTA_scale_weight=1, + feature_noise=0, + ): + batch_size = B_lab_map.shape[0] + channel = B_lab_map.shape[1] + image_height = B_lab_map.shape[2] + image_width = B_lab_map.shape[3] + feature_height = int(image_height / 4) + feature_width = int(image_width / 4) + + # scale feature size to 44*44 + A_feature2_1 = self.layer2_1(A_relu2_1) + B_feature2_1 = self.layer2_1(B_relu2_1) + A_feature3_1 = self.layer3_1(A_relu3_1) + B_feature3_1 = self.layer3_1(B_relu3_1) + A_feature4_1 = self.layer4_1(A_relu4_1) + B_feature4_1 = self.layer4_1(B_relu4_1) + A_feature5_1 = self.layer5_1(A_relu5_1) + B_feature5_1 = self.layer5_1(B_relu5_1) + + # concatenate features + if A_feature5_1.shape[2] != A_feature2_1.shape[2] or A_feature5_1.shape[3] != A_feature2_1.shape[3]: + A_feature5_1 = F.pad(A_feature5_1, (0, 0, 1, 1), "replicate") + B_feature5_1 = F.pad(B_feature5_1, (0, 0, 1, 1), "replicate") + A_features = self.layer(torch.cat((A_feature2_1, A_feature3_1, A_feature4_1, A_feature5_1), 1)) + B_features = self.layer(torch.cat((B_feature2_1, B_feature3_1, B_feature4_1, B_feature5_1), 1)) + + # pairwise cosine similarity + theta = self.theta(A_features).view(batch_size, self.inter_channels, -1) # 2*256*(feature_height*feature_width) + theta = theta - theta.mean(dim=-1, keepdim=True) # center the feature + theta_norm = torch.norm(theta, 2, 1, keepdim=True) + sys.float_info.epsilon + theta = torch.div(theta, theta_norm) + theta_permute = theta.permute(0, 2, 1) # 2*(feature_height*feature_width)*256 + phi = self.phi(B_features).view(batch_size, self.inter_channels, -1) # 2*256*(feature_height*feature_width) + phi = phi - phi.mean(dim=-1, keepdim=True) # center the feature + phi_norm = torch.norm(phi, 2, 1, keepdim=True) + sys.float_info.epsilon + phi = torch.div(phi, phi_norm) + f = torch.matmul(theta_permute, phi) # 2*(feature_height*feature_width)*(feature_height*feature_width) + if detach_flag: + f = f.detach() + + f_similarity = f.unsqueeze_(dim=1) + similarity_map = torch.max(f_similarity, -1, keepdim=True)[0] + similarity_map = similarity_map.view(batch_size, 1, feature_height, feature_width) + + # f can be negative + if WTA_scale_weight == 1: + f_WTA = f + else: + f_WTA = WTA_scale.apply(f, WTA_scale_weight) + f_WTA = f_WTA / temperature + f_div_C = F.softmax(f_WTA.squeeze_(), dim=-1) # 2*1936*1936; softmax along the horizontal line (dim=-1) + + # downsample the reference color + B_lab = F.avg_pool2d(B_lab_map, 4) + B_lab = B_lab.view(batch_size, channel, -1) + B_lab = B_lab.permute(0, 2, 1) # 2*1936*channel + + # multiply the corr map with color + y = torch.matmul(f_div_C, B_lab) # 2*1936*channel + y = y.permute(0, 2, 1).contiguous() + y = y.view(batch_size, channel, feature_height, feature_width) # 2*3*44*44 + y = self.upsampling(y) + similarity_map = self.upsampling(similarity_map) + + return y, similarity_map diff --git a/models/guided_filter.py b/models/guided_filter.py new file mode 100644 index 0000000..b2cab00 --- /dev/null +++ b/models/guided_filter.py @@ -0,0 +1,125 @@ +import torch +from torch import nn +from torch.autograd import Variable +from torch.nn import functional as F + +# from .box_filter import BoxFilter + + +def diff_x(input, r): + assert input.dim() == 4 + + left = input[:, :, r : 2 * r + 1] + middle = input[:, :, 2 * r + 1 :] - input[:, :, : -2 * r - 1] + right = input[:, :, -1:] - input[:, :, -2 * r - 1 : -r - 1] + + output = torch.cat([left, middle, right], dim=2) + + return output + + +def diff_y(input, r): + assert input.dim() == 4 + + left = input[:, :, :, r : 2 * r + 1] + middle = input[:, :, :, 2 * r + 1 :] - input[:, :, :, : -2 * r - 1] + right = input[:, :, :, -1:] - input[:, :, :, -2 * r - 1 : -r - 1] + + output = torch.cat([left, middle, right], dim=3) + + return output + + +class BoxFilter(nn.Module): + def __init__(self, r): + super(BoxFilter, self).__init__() + + self.r = r + + def forward(self, x): + assert x.dim() == 4 + + return diff_y(diff_x(x.cumsum(dim=2), self.r).cumsum(dim=3), self.r) + + +class FastGuidedFilter(nn.Module): + def __init__(self, r, eps=1e-8): + super(FastGuidedFilter, self).__init__() + + self.r = r + self.eps = eps + self.boxfilter = BoxFilter(r) + + def forward(self, lr_x, lr_y, hr_x): + n_lrx, c_lrx, h_lrx, w_lrx = lr_x.size() + n_lry, c_lry, h_lry, w_lry = lr_y.size() + n_hrx, c_hrx, h_hrx, w_hrx = hr_x.size() + + assert n_lrx == n_lry and n_lry == n_hrx + assert c_lrx == c_hrx and (c_lrx == 1 or c_lrx == c_lry) + assert h_lrx == h_lry and w_lrx == w_lry + assert h_lrx > 2 * self.r + 1 and w_lrx > 2 * self.r + 1 + + ## N + N = self.boxfilter(Variable(lr_x.data.new().resize_((1, 1, h_lrx, w_lrx)).fill_(1.0))) + + ## mean_x + mean_x = self.boxfilter(lr_x) / N + ## mean_y + mean_y = self.boxfilter(lr_y) / N + ## cov_xy + cov_xy = self.boxfilter(lr_x * lr_y) / N - mean_x * mean_y + ## var_x + var_x = self.boxfilter(lr_x * lr_x) / N - mean_x * mean_x + + ## A + A = cov_xy / (var_x + self.eps) + ## b + b = mean_y - A * mean_x + + ## mean_A; mean_b + mean_A = F.upsample(A, (h_hrx, w_hrx), mode="bilinear") + mean_b = F.upsample(b, (h_hrx, w_hrx), mode="bilinear") + + return mean_A * hr_x + mean_b + + +class GuidedFilter(nn.Module): + def __init__(self, r, eps=1e-8): + super(GuidedFilter, self).__init__() + + self.r = r + self.eps = eps + self.boxfilter = BoxFilter(r) + + def forward(self, x, y): + n_x, c_x, h_x, w_x = x.size() + n_y, c_y, h_y, w_y = y.size() + + assert n_x == n_y + assert c_x == 1 or c_x == c_y + assert h_x == h_y and w_x == w_y + assert h_x > 2 * self.r + 1 and w_x > 2 * self.r + 1 + + # N + N = self.boxfilter(Variable(x.data.new().resize_((1, 1, h_x, w_x)).fill_(1.0))) + + # mean_x + mean_x = self.boxfilter(x) / N + # mean_y + mean_y = self.boxfilter(y) / N + # cov_xy + cov_xy = self.boxfilter(x * y) / N - mean_x * mean_y + # var_x + var_x = self.boxfilter(x * x) / N - mean_x * mean_x + + # A + A = cov_xy / (var_x + self.eps) + # b + b = mean_y - A * mean_x + + # mean_A; mean_b + mean_A = self.boxfilter(A) / N + mean_b = self.boxfilter(b) / N + + return mean_A * x + mean_b diff --git a/models/spectral_normalization.py b/models/spectral_normalization.py new file mode 100644 index 0000000..30279fe --- /dev/null +++ b/models/spectral_normalization.py @@ -0,0 +1,65 @@ +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.autograd import Variable +from torch.nn import Parameter +from torch.optim.optimizer import Optimizer, required + + +def l2normalize(v, eps=1e-12): + return v / (v.norm() + eps) + + +class SpectralNorm(nn.Module): + def __init__(self, module, name="weight", power_iterations=1): + super(SpectralNorm, self).__init__() + self.module = module + self.name = name + self.power_iterations = power_iterations + if not self._made_params(): + self._make_params() + + def _update_u_v(self): + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + + height = w.data.shape[0] + for _ in range(self.power_iterations): + v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data)) + u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data)) + + # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data)) + sigma = u.dot(w.view(height, -1).mv(v)) + setattr(self.module, self.name, w / sigma.expand_as(w)) + + def _made_params(self): + try: + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + return True + except AttributeError: + return False + + def _make_params(self): + w = getattr(self.module, self.name) + + height = w.data.shape[0] + width = w.view(height, -1).data.shape[1] + + u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) + v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) + u.data = l2normalize(u.data) + v.data = l2normalize(v.data) + w_bar = Parameter(w.data) + + del self.module._parameters[self.name] + + self.module.register_parameter(self.name + "_u", u) + self.module.register_parameter(self.name + "_v", v) + self.module.register_parameter(self.name + "_bar", w_bar) + + def forward(self, *args): + self._update_u_v() + return self.module.forward(*args) diff --git a/models/vgg19_gray.py b/models/vgg19_gray.py new file mode 100644 index 0000000..a6b8315 --- /dev/null +++ b/models/vgg19_gray.py @@ -0,0 +1,236 @@ +from collections import OrderedDict, namedtuple +from functools import reduce + +import torch +import torch.nn as nn +from torch.autograd import Variable + + +class LambdaBase(nn.Sequential): + def __init__(self, fn, *args): + super(LambdaBase, self).__init__(*args) + self.lambda_func = fn + + def forward_prepare(self, input): + output = [] + for module in self._modules.values(): + output.append(module(input)) + return output if output else input + + +class Lambda(LambdaBase): + def forward(self, input): + return self.lambda_func(self.forward_prepare(input)) + + +class LambdaMap(LambdaBase): + def forward(self, input): + return list(map(self.lambda_func, self.forward_prepare(input))) + + +class LambdaReduce(LambdaBase): + def forward(self, input): + return reduce(self.lambda_func, self.forward_prepare(input)) + + +layer_names = [ + "conv1_1", + "relu1_1", + "conv1_2", + "relu1_2", + "pool1", + "conv2_1", + "relu2_1", + "conv2_2", + "relu2_2", + "pool2", + "conv3_1", + "relu3_1", + "conv3_2", + "relu3_2", + "conv3_3", + "relu3_3", + "conv3_4", + "relu3_4", + "pool3", + "conv4_1", + "relu4_1", + "conv4_2", + "relu4_2", + "conv4_3", + "relu4_3", + "conv4_4", + "relu4_4", + "pool4", + "conv5_1", + "relu5_1", + "conv5_2", + "relu5_2", + "conv5_3", + "relu5_3", + "conv5_4", + "relu5_4", + "pool5", + "view1", + "fc6", + "fc6_relu", + "fc7", + "fc7_relu", + "fc8", +] + +model = nn.Sequential( # Sequential, + nn.Conv2d(3, 64, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.Conv2d(128, 256, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.Conv2d(256, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)), + nn.ReLU(), + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + Lambda(lambda x: x.view(x.size(0), -1)), # View, + nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(25088, 4096)), # Linear, + nn.ReLU(), + nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(4096, 4096)), # Linear, + nn.ReLU(), + nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(4096, 1000)), # Linear, +) + + +model.load_state_dict(torch.load("models/vgg19_gray.pth")) +# print(model) +vgg19_gray_net = torch.nn.Sequential() +for (name, layer) in model._modules.items(): + vgg19_gray_net.add_module(layer_names[int(name)], model[int(name)]) +# print(vgg19_gray_net) + + +for param in vgg19_gray_net.parameters(): + param.requires_grad = False +vgg19_gray_net.eval() + + +class vgg19_gray(torch.nn.Module): + def __init__(self, requires_grad=False): + super(vgg19_gray, self).__init__() + vgg_pretrained_features = vgg19_gray_net + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + for x in range(12): + self.slice1.add_module(layer_names[x], vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice2.add_module(layer_names[x], vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice3.add_module(layer_names[x], vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + h = self.slice1(X) + h_relu3_1 = h + h = self.slice2(h) + h_relu4_1 = h + h = self.slice3(h) + h_relu5_1 = h + # vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1']) + # out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1) + return h_relu3_1, h_relu4_1, h_relu5_1 + + +class vgg19_gray_new(torch.nn.Module): + def __init__(self, requires_grad=False): + super(vgg19_gray_new, self).__init__() + vgg_pretrained_features = vgg19_gray_net + self.slice0 = torch.nn.Sequential() + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + for x in range(7): + self.slice0.add_module(layer_names[x], vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice1.add_module(layer_names[x], vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice2.add_module(layer_names[x], vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice3.add_module(layer_names[x], vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + h = self.slice0(X) + h_relu2_1 = h + h = self.slice1(h) + h_relu3_1 = h + h = self.slice2(h) + h_relu4_1 = h + h = self.slice3(h) + h_relu5_1 = h + # vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1']) + # out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1) + return h_relu2_1, h_relu3_1, h_relu4_1, h_relu5_1 + + +# class vgg19_gray_conv(torch.nn.Module): +# def __init__(self, requires_grad=False): +# super(vgg19_gray_conv, self).__init__() +# vgg_pretrained_features = vgg19_gray_net +# self.slice0 = torch.nn.Sequential() +# self.slice1 = torch.nn.Sequential() +# self.slice2 = torch.nn.Sequential() +# self.slice3 = torch.nn.Sequential() +# for x in range(6): +# self.slice0.add_module(layer_names[x], vgg_pretrained_features[x]) +# for x in range(6, 11): +# self.slice1.add_module(layer_names[x], vgg_pretrained_features[x]) +# for x in range(11, 20): +# self.slice2.add_module(layer_names[x], vgg_pretrained_features[x]) +# for x in range(20, 29): +# self.slice3.add_module(layer_names[x], vgg_pretrained_features[x]) +# if not requires_grad: +# for param in self.parameters(): +# param.requires_grad = False + +# def forward(self, X): +# h = self.slice0(X) +# h_relu2_1 = h +# h = self.slice1(h) +# h_relu3_1 = h +# h = self.slice2(h) +# h_relu4_1 = h +# h = self.slice3(h) +# h_relu5_1 = h +# # vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1']) +# # out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1) +# return h_relu2_1, h_relu3_1, h_relu4_1, h_relu5_1 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7539545 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,17 @@ +numpy +tqdm>=4.23.0 +scipy==1.2 +torchsummary>=1.3 +matplotlib>=2.2.2 +opencv_contrib_python>=4.1.0.25 +torchvision>=0.3.0 +scikit_image>=0.15.0 +torchviz>=0.0.1 +Pillow>=6.1.0 +torch +easydict +prefetch_generator +PyYAML +tensorboard +future +numba \ No newline at end of file diff --git a/tensorboardX/__init__.py b/tensorboardX/__init__.py new file mode 100644 index 0000000..472510a --- /dev/null +++ b/tensorboardX/__init__.py @@ -0,0 +1,5 @@ +"""A module for visualization with tensorboard +""" + +from .record_writer import RecordWriter +from .writer import FileWriter, SummaryWriter diff --git a/tensorboardX/crc32c.py b/tensorboardX/crc32c.py new file mode 100644 index 0000000..0870b1c --- /dev/null +++ b/tensorboardX/crc32c.py @@ -0,0 +1,314 @@ +import array + +CRC_TABLE = ( + 0x00000000, + 0xF26B8303, + 0xE13B70F7, + 0x1350F3F4, + 0xC79A971F, + 0x35F1141C, + 0x26A1E7E8, + 0xD4CA64EB, + 0x8AD958CF, + 0x78B2DBCC, + 0x6BE22838, + 0x9989AB3B, + 0x4D43CFD0, + 0xBF284CD3, + 0xAC78BF27, + 0x5E133C24, + 0x105EC76F, + 0xE235446C, + 0xF165B798, + 0x030E349B, + 0xD7C45070, + 0x25AFD373, + 0x36FF2087, + 0xC494A384, + 0x9A879FA0, + 0x68EC1CA3, + 0x7BBCEF57, + 0x89D76C54, + 0x5D1D08BF, + 0xAF768BBC, + 0xBC267848, + 0x4E4DFB4B, + 0x20BD8EDE, + 0xD2D60DDD, + 0xC186FE29, + 0x33ED7D2A, + 0xE72719C1, + 0x154C9AC2, + 0x061C6936, + 0xF477EA35, + 0xAA64D611, + 0x580F5512, + 0x4B5FA6E6, + 0xB93425E5, + 0x6DFE410E, + 0x9F95C20D, + 0x8CC531F9, + 0x7EAEB2FA, + 0x30E349B1, + 0xC288CAB2, + 0xD1D83946, + 0x23B3BA45, + 0xF779DEAE, + 0x05125DAD, + 0x1642AE59, + 0xE4292D5A, + 0xBA3A117E, + 0x4851927D, + 0x5B016189, + 0xA96AE28A, + 0x7DA08661, + 0x8FCB0562, + 0x9C9BF696, + 0x6EF07595, + 0x417B1DBC, + 0xB3109EBF, + 0xA0406D4B, + 0x522BEE48, + 0x86E18AA3, + 0x748A09A0, + 0x67DAFA54, + 0x95B17957, + 0xCBA24573, + 0x39C9C670, + 0x2A993584, + 0xD8F2B687, + 0x0C38D26C, + 0xFE53516F, + 0xED03A29B, + 0x1F682198, + 0x5125DAD3, + 0xA34E59D0, + 0xB01EAA24, + 0x42752927, + 0x96BF4DCC, + 0x64D4CECF, + 0x77843D3B, + 0x85EFBE38, + 0xDBFC821C, + 0x2997011F, + 0x3AC7F2EB, + 0xC8AC71E8, + 0x1C661503, + 0xEE0D9600, + 0xFD5D65F4, + 0x0F36E6F7, + 0x61C69362, + 0x93AD1061, + 0x80FDE395, + 0x72966096, + 0xA65C047D, + 0x5437877E, + 0x4767748A, + 0xB50CF789, + 0xEB1FCBAD, + 0x197448AE, + 0x0A24BB5A, + 0xF84F3859, + 0x2C855CB2, + 0xDEEEDFB1, + 0xCDBE2C45, + 0x3FD5AF46, + 0x7198540D, + 0x83F3D70E, + 0x90A324FA, + 0x62C8A7F9, + 0xB602C312, + 0x44694011, + 0x5739B3E5, + 0xA55230E6, + 0xFB410CC2, + 0x092A8FC1, + 0x1A7A7C35, + 0xE811FF36, + 0x3CDB9BDD, + 0xCEB018DE, + 0xDDE0EB2A, + 0x2F8B6829, + 0x82F63B78, + 0x709DB87B, + 0x63CD4B8F, + 0x91A6C88C, + 0x456CAC67, + 0xB7072F64, + 0xA457DC90, + 0x563C5F93, + 0x082F63B7, + 0xFA44E0B4, + 0xE9141340, + 0x1B7F9043, + 0xCFB5F4A8, + 0x3DDE77AB, + 0x2E8E845F, + 0xDCE5075C, + 0x92A8FC17, + 0x60C37F14, + 0x73938CE0, + 0x81F80FE3, + 0x55326B08, + 0xA759E80B, + 0xB4091BFF, + 0x466298FC, + 0x1871A4D8, + 0xEA1A27DB, + 0xF94AD42F, + 0x0B21572C, + 0xDFEB33C7, + 0x2D80B0C4, + 0x3ED04330, + 0xCCBBC033, + 0xA24BB5A6, + 0x502036A5, + 0x4370C551, + 0xB11B4652, + 0x65D122B9, + 0x97BAA1BA, + 0x84EA524E, + 0x7681D14D, + 0x2892ED69, + 0xDAF96E6A, + 0xC9A99D9E, + 0x3BC21E9D, + 0xEF087A76, + 0x1D63F975, + 0x0E330A81, + 0xFC588982, + 0xB21572C9, + 0x407EF1CA, + 0x532E023E, + 0xA145813D, + 0x758FE5D6, + 0x87E466D5, + 0x94B49521, + 0x66DF1622, + 0x38CC2A06, + 0xCAA7A905, + 0xD9F75AF1, + 0x2B9CD9F2, + 0xFF56BD19, + 0x0D3D3E1A, + 0x1E6DCDEE, + 0xEC064EED, + 0xC38D26C4, + 0x31E6A5C7, + 0x22B65633, + 0xD0DDD530, + 0x0417B1DB, + 0xF67C32D8, + 0xE52CC12C, + 0x1747422F, + 0x49547E0B, + 0xBB3FFD08, + 0xA86F0EFC, + 0x5A048DFF, + 0x8ECEE914, + 0x7CA56A17, + 0x6FF599E3, + 0x9D9E1AE0, + 0xD3D3E1AB, + 0x21B862A8, + 0x32E8915C, + 0xC083125F, + 0x144976B4, + 0xE622F5B7, + 0xF5720643, + 0x07198540, + 0x590AB964, + 0xAB613A67, + 0xB831C993, + 0x4A5A4A90, + 0x9E902E7B, + 0x6CFBAD78, + 0x7FAB5E8C, + 0x8DC0DD8F, + 0xE330A81A, + 0x115B2B19, + 0x020BD8ED, + 0xF0605BEE, + 0x24AA3F05, + 0xD6C1BC06, + 0xC5914FF2, + 0x37FACCF1, + 0x69E9F0D5, + 0x9B8273D6, + 0x88D28022, + 0x7AB90321, + 0xAE7367CA, + 0x5C18E4C9, + 0x4F48173D, + 0xBD23943E, + 0xF36E6F75, + 0x0105EC76, + 0x12551F82, + 0xE03E9C81, + 0x34F4F86A, + 0xC69F7B69, + 0xD5CF889D, + 0x27A40B9E, + 0x79B737BA, + 0x8BDCB4B9, + 0x988C474D, + 0x6AE7C44E, + 0xBE2DA0A5, + 0x4C4623A6, + 0x5F16D052, + 0xAD7D5351, +) + + +CRC_INIT = 0 + +_MASK = 0xFFFFFFFF + + +def crc_update(crc, data): + """Update CRC-32C checksum with data. + + Args: + crc: 32-bit checksum to update as long. + data: byte array, string or iterable over bytes. + + Returns: + 32-bit updated CRC-32C as long. + """ + + if type(data) != array.array or data.itemsize != 1: + buf = array.array("B", data) + else: + buf = data + + crc ^= _MASK + for b in buf: + table_index = (crc ^ b) & 0xFF + crc = (CRC_TABLE[table_index] ^ (crc >> 8)) & _MASK + return crc ^ _MASK + + +def crc_finalize(crc): + """Finalize CRC-32C checksum. + + This function should be called as last step of crc calculation. + + Args: + crc: 32-bit checksum as long. + + Returns: + finalized 32-bit checksum as long + """ + return crc & _MASK + + +def crc32c(data): + """Compute CRC-32C checksum of the data. + + Args: + data: byte array, string or iterable over bytes. + + Returns: + 32-bit CRC-32C checksum of data as long. + """ + return crc_finalize(crc_update(CRC_INIT, data)) diff --git a/tensorboardX/embedding.py b/tensorboardX/embedding.py new file mode 100644 index 0000000..a5c8259 --- /dev/null +++ b/tensorboardX/embedding.py @@ -0,0 +1,52 @@ +import os + + +def make_tsv(metadata, save_path): + metadata = [str(x) for x in metadata] + with open(os.path.join(save_path, "metadata.tsv"), "w") as f: + for x in metadata: + f.write(x + "\n") + + +# https://github.com/tensorflow/tensorboard/issues/44 image label will be squared +def make_sprite(label_img, save_path): + import math + + import torch + import torchvision + + from .x2num import makenp + + # this ensures the sprite image has correct dimension as described in + # https://www.tensorflow.org/get_started/embedding_viz + nrow = int(math.ceil((label_img.size(0)) ** 0.5)) + + label_img = torch.from_numpy(makenp(label_img)) # for other framework + # augment images so that #images equals nrow*nrow + label_img = torch.cat((label_img, torch.randn(nrow ** 2 - label_img.size(0), *label_img.size()[1:]) * 255), 0) + + torchvision.utils.save_image(label_img, os.path.join(save_path, "sprite.png"), nrow=nrow, padding=0) + + +def append_pbtxt(metadata, label_img, save_path, global_step, tag): + with open(os.path.join(save_path, "projector_config.pbtxt"), "a") as f: + # step = os.path.split(save_path)[-1] + f.write("embeddings {\n") + f.write('tensor_name: "{}:{}"\n'.format(tag, global_step)) + f.write('tensor_path: "{}"\n'.format(os.path.join(global_step, "tensors.tsv"))) + if metadata is not None: + f.write('metadata_path: "{}"\n'.format(os.path.join(global_step, "metadata.tsv"))) + if label_img is not None: + f.write("sprite {\n") + f.write('image_path: "{}"\n'.format(os.path.join(global_step, "sprite.png"))) + f.write("single_image_dim: {}\n".format(label_img.size(3))) + f.write("single_image_dim: {}\n".format(label_img.size(2))) + f.write("}\n") + f.write("}\n") + + +def make_mat(matlist, save_path): + with open(os.path.join(save_path, "tensors.tsv"), "w") as f: + for x in matlist: + x = [str(i) for i in x] + f.write("\t".join(x) + "\n") diff --git a/tensorboardX/event_file_writer.py b/tensorboardX/event_file_writer.py new file mode 100644 index 0000000..b9d826b --- /dev/null +++ b/tensorboardX/event_file_writer.py @@ -0,0 +1,192 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Writes events to disk in a logdir.""" + +from __future__ import absolute_import, division, print_function + +import logging +import os.path +import socket +import threading +import time + +import six + +from .record_writer import RecordWriter +from .src import event_pb2 + + +def directory_check(path): + """Initialize the directory for log files.""" + # If the direcotry does not exist, create it! + if not os.path.exists(path): + os.makedirs(path) + + +class EventsWriter(object): + """Writes `Event` protocol buffers to an event file.""" + + def __init__(self, file_prefix): + """ + Events files have a name of the form + '/some/file/path/events.out.tfevents.[timestamp].[hostname]' + """ + self._file_prefix = file_prefix + ".out.tfevents." + str(time.time())[:10] + "." + socket.gethostname() + + # Open(Create) the log file with the particular form of name. + logging.basicConfig(filename=self._file_prefix) + + self._num_outstanding_events = 0 + + self._py_recordio_writer = RecordWriter(self._file_prefix) + + # Initialize an event instance. + self._event = event_pb2.Event() + + self._event.wall_time = time.time() + + self.write_event(self._event) + + def write_event(self, event): + """Append "event" to the file.""" + + # Check if event is of type event_pb2.Event proto. + if not isinstance(event, event_pb2.Event): + raise TypeError("Expected an event_pb2.Event proto, " " but got %s" % type(event)) + return self._write_serialized_event(event.SerializeToString()) + + def _write_serialized_event(self, event_str): + self._num_outstanding_events += 1 + self._py_recordio_writer.write(event_str) + + def flush(self): + """Flushes the event file to disk.""" + self._num_outstanding_events = 0 + return True + + def close(self): + """Call self.flush().""" + return_value = self.flush() + return return_value + + +class EventFileWriter(object): + """Writes `Event` protocol buffers to an event file. + The `EventFileWriter` class creates an event file in the specified directory, + and asynchronously writes Event protocol buffers to the file. The Event file + is encoded using the tfrecord format, which is similar to RecordIO. + @@__init__ + @@add_event + @@flush + @@close + """ + + def __init__(self, logdir, max_queue=10, flush_secs=120): + """Creates a `EventFileWriter` and an event file to write to. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers, which are written to + disk via the add_event method. + The other arguments to the constructor control the asynchronous writes to + the event file: + * `flush_secs`: How often, in seconds, to flush the added summaries + and events to disk. + * `max_queue`: Maximum number of summaries or events pending to be + written to disk before one of the 'add' calls block. + Args: + logdir: A string. Directory where event file will be written. + max_queue: Integer. Size of the queue for pending events and summaries. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. + """ + self._logdir = logdir + directory_check(self._logdir) + self._event_queue = six.moves.queue.Queue(max_queue) + self._ev_writer = EventsWriter(os.path.join(self._logdir, "events")) + self._closed = False + self._worker = _EventLoggerThread(self._event_queue, self._ev_writer, flush_secs) + + self._worker.start() + + def get_logdir(self): + """Returns the directory where event file will be written.""" + return self._logdir + + def reopen(self): + """Reopens the EventFileWriter. + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + Does nothing if the EventFileWriter was not closed. + """ + if self._closed: + self._closed = False + + def add_event(self, event): + """Adds an event to the event file. + Args: + event: An `Event` protocol buffer. + """ + if not self._closed: + self._event_queue.put(event) + + def flush(self): + """Flushes the event file to disk. + Call this method to make sure that all pending events have been written to + disk. + """ + self._event_queue.join() + self._ev_writer.flush() + + def close(self): + """Flushes the event file to disk and close the file. + Call this method when you do not need the summary writer anymore. + """ + self.flush() + self._ev_writer.close() + self._closed = True + + +class _EventLoggerThread(threading.Thread): + """Thread that logs events.""" + + def __init__(self, queue, ev_writer, flush_secs): + """Creates an _EventLoggerThread. + Args: + queue: A Queue from which to dequeue events. + ev_writer: An event writer. Used to log brain events for + the visualizer. + flush_secs: How often, in seconds, to flush the + pending file to disk. + """ + threading.Thread.__init__(self) + self.daemon = True + self._queue = queue + self._ev_writer = ev_writer + self._flush_secs = flush_secs + # The first event will be flushed immediately. + self._next_event_flush_time = 0 + + def run(self): + while True: + event = self._queue.get() + try: + self._ev_writer.write_event(event) + # Flush the event writer every so often. + now = time.time() + if now > self._next_event_flush_time: + self._ev_writer.flush() + # Do it again in two minutes. + self._next_event_flush_time = now + self._flush_secs + finally: + self._queue.task_done() diff --git a/tensorboardX/graph.py b/tensorboardX/graph.py new file mode 100644 index 0000000..afb5329 --- /dev/null +++ b/tensorboardX/graph.py @@ -0,0 +1,73 @@ +from distutils.version import LooseVersion + +from .src.attr_value_pb2 import AttrValue +from .src.graph_pb2 import GraphDef +from .src.node_def_pb2 import NodeDef +from .src.tensor_shape_pb2 import TensorShapeProto +from .src.versions_pb2 import VersionDef + + +def replace(name, scope): + return "/".join([scope[name], name]) + + +def parse(graph): + import torch + + scope = {} + for n in graph.nodes(): + inputs = [i.uniqueName() for i in n.inputs()] + for i in range(1, len(inputs)): + scope[inputs[i]] = n.scopeName() + + uname = next(iter(n.outputs())).uniqueName() + assert n.scopeName() != "", "{} has empty scope name".format(n) + scope[uname] = n.scopeName() + if LooseVersion(torch.__version__) >= LooseVersion("0.4"): + scope["0"] = "input" + else: + scope["1"] = "input" + + nodes = [] + for n in graph.nodes(): + attrs = {k: n[k] for k in n.attributeNames()} + attrs = str(attrs).replace("'", " ") # singlequote will be escaped by tensorboard + if any(i.uniqueName() not in scope.keys() for i in n.inputs()): # 0.3.1 workaround + continue + inputs = [replace(i.uniqueName(), scope) for i in n.inputs()] + uname = next(iter(n.outputs())).uniqueName() # FIXME: only first output is considered + nodes.append({"name": replace(uname, scope), "op": n.kind(), "inputs": inputs, "attr": attrs}) + + for n in graph.inputs(): + uname = n.uniqueName() + if uname not in scope.keys(): + scope[uname] = "unused" + nodes.append({"name": replace(uname, scope), "op": "Parameter", "inputs": [], "attr": str(n.type())}) + + return nodes + + +def graph(model, args, verbose=False): + import torch + + with torch.onnx.set_training(model, False): + trace, _ = torch.jit.trace(model, args) + if LooseVersion(torch.__version__) >= LooseVersion("0.4"): + torch.onnx._optimize_trace(trace, False) + else: + torch.onnx._optimize_trace(trace) + graph = trace.graph() + if verbose: + print(graph) + list_of_nodes = parse(graph) + nodes = [] + for node in list_of_nodes: + nodes.append( + NodeDef( + name=node["name"], + op=node["op"], + input=node["inputs"], + attr={"lanpa": AttrValue(s=node["attr"].encode(encoding="utf_8"))}, + ) + ) + return GraphDef(node=nodes, versions=VersionDef(producer=22)) diff --git a/tensorboardX/graph_onnx.py b/tensorboardX/graph_onnx.py new file mode 100644 index 0000000..ac12b16 --- /dev/null +++ b/tensorboardX/graph_onnx.py @@ -0,0 +1,159 @@ +from .src.attr_value_pb2 import AttrValue +from .src.graph_pb2 import GraphDef +from .src.node_def_pb2 import NodeDef +from .src.tensor_shape_pb2 import TensorShapeProto +from .src.versions_pb2 import VersionDef + +# from .src.onnx_pb2 import ModelProto + + +def gg(fname): + import onnx # 0.2.1 + + m = onnx.load(fname) + nodes_proto = [] + nodes = [] + g = m.graph + import itertools + + for node in itertools.chain(g.input, g.output): + nodes_proto.append(node) + + for node in nodes_proto: + shapeproto = TensorShapeProto( + dim=[TensorShapeProto.Dim(size=d.dim_value) for d in node.type.tensor_type.shape.dim] + ) + nodes.append( + NodeDef( + name=node.name, + op="Variable", + input=[], + attr={"dtype": AttrValue(type=node.type.tensor_type.elem_type), "shape": AttrValue(shape=shapeproto),}, + ) + ) + + for node in g.node: + attr = [] + for s in node.attribute: + attr.append(" = ".join([str(f[1]) for f in s.ListFields()])) + attr = ", ".join(attr).encode(encoding="utf_8") + + nodes.append( + NodeDef(name=node.output[0], op=node.op_type, input=node.input, attr={"parameters": AttrValue(s=attr)},) + ) + # two pass token replacement, appends opname to object id + mapping = {} + for node in nodes: + mapping[node.name] = node.op + "_" + node.name + + nodes, mapping = updatenodes(nodes, mapping) + mapping = smartGrouping(nodes, mapping) + nodes, mapping = updatenodes(nodes, mapping) + + return GraphDef(node=nodes, versions=VersionDef(producer=22)) + + +def updatenodes(nodes, mapping): + for node in nodes: + newname = mapping[node.name] + node.name = newname + newinput = [] + for inputnode in list(node.input): + newinput.append(mapping[inputnode]) + node.input.remove(inputnode) + node.input.extend(newinput) + newmap = {} + for k, v in mapping.items(): + newmap[v] = v + return nodes, newmap + + +def findnode(nodes, name): + """ input: node name + returns: node object + """ + for n in nodes: + if n.name == name: + return n + + +def parser(s, nodes, node): + print(s) + if len(s) == 0: + return + if len(s) > 0: + if s[0] == node.op: + print(s[0], node.name, s[1], node.input) + for n in node.input: + print(n, s[1]) + parser(s[1], nodes, findnode(nodes, n)) + else: + return False + + +# TODO: use recursive parse + + +def smartGrouping(nodes, mapping): + # a Fully Conv is: (TODO: check var1.size(0)==var2.size(0)) + # GEMM <-- Variable (c1) + # ^-- Transpose (c2) <-- Variable (c3) + + # a Conv with bias is: (TODO: check var1.size(0)==var2.size(0)) + # Add <-- Conv (c2) <-- Variable (c3) + # ^-- Variable (c1) + # + # gemm = ('Gemm', ('Variable', ('Transpose', ('Variable')))) + + FCcounter = 1 + Convcounter = 1 + for node in nodes: + if node.op == "Gemm": + c1 = c2 = c3 = False + for name_in in node.input: + n = findnode(nodes, name_in) + if n.op == "Variable": + c1 = True + c1name = n.name + if n.op == "Transpose": + c2 = True + c2name = n.name + if len(n.input) == 1: + nn = findnode(nodes, n.input[0]) + if nn.op == "Variable": + c3 = True + c3name = nn.name + # print(n.op, n.name, c1, c2, c3) + if c1 and c2 and c3: + # print(c1name, c2name, c3name) + mapping[c1name] = "FC{}/{}".format(FCcounter, c1name) + mapping[c2name] = "FC{}/{}".format(FCcounter, c2name) + mapping[c3name] = "FC{}/{}".format(FCcounter, c3name) + mapping[node.name] = "FC{}/{}".format(FCcounter, node.name) + FCcounter += 1 + continue + if node.op == "Add": + c1 = c2 = c3 = False + for name_in in node.input: + n = findnode(nodes, name_in) + if n.op == "Variable": + c1 = True + c1name = n.name + if n.op == "Conv": + c2 = True + c2name = n.name + if len(n.input) >= 1: + for nn_name in n.input: + nn = findnode(nodes, nn_name) + if nn.op == "Variable": + c3 = True + c3name = nn.name + + if c1 and c2 and c3: + # print(c1name, c2name, c3name) + mapping[c1name] = "Conv{}/{}".format(Convcounter, c1name) + mapping[c2name] = "Conv{}/{}".format(Convcounter, c2name) + mapping[c3name] = "Conv{}/{}".format(Convcounter, c3name) + mapping[node.name] = "Conv{}/{}".format(Convcounter, node.name) + Convcounter += 1 + return mapping diff --git a/tensorboardX/record_writer.py b/tensorboardX/record_writer.py new file mode 100644 index 0000000..85da287 --- /dev/null +++ b/tensorboardX/record_writer.py @@ -0,0 +1,47 @@ +""" +To write tf_record into file. Here we use it for tensorboard's event writting. +The code was borrow from https://github.com/TeamHG-Memex/tensorboard_logger +""" + +import re +import struct + +from .crc32c import crc32c + +_VALID_OP_NAME_START = re.compile("^[A-Za-z0-9.]") +_VALID_OP_NAME_PART = re.compile("[A-Za-z0-9_.\\-/]+") + + +class RecordWriter(object): + def __init__(self, path, flush_secs=2): + self._name_to_tf_name = {} + self._tf_names = set() + self.path = path + self.flush_secs = flush_secs # TODO. flush every flush_secs, not every time. + self._writer = None + self._writer = open(path, "wb") + + def write(self, event_str): + w = self._writer.write + header = struct.pack("Q", len(event_str)) + w(header) + w(struct.pack("I", masked_crc32c(header))) + w(event_str) + w(struct.pack("I", masked_crc32c(event_str))) + self._writer.flush() + + +def masked_crc32c(data): + x = u32(crc32c(data)) + return u32(((x >> 15) | u32(x << 17)) + 0xA282EAD8) + + +def u32(x): + return x & 0xFFFFFFFF + + +def make_valid_tf_name(name): + if not _VALID_OP_NAME_START.match(name): + # Must make it valid somehow, but don't want to remove stuff + name = "." + name + return "_".join(_VALID_OP_NAME_PART.findall(name)) diff --git a/tensorboardX/src/__init__.py b/tensorboardX/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tensorboardX/src/attr_value_pb2.py b/tensorboardX/src/attr_value_pb2.py new file mode 100644 index 0000000..a9900c0 --- /dev/null +++ b/tensorboardX/src/attr_value_pb2.py @@ -0,0 +1,600 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/attr_value.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2 +from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2 +from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/attr_value.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + "\n!tensorboardX/src/attr_value.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto\"\xaf\x04\n\tAttrValue\x12\x0b\n\x01s\x18\x02 \x01(\x0cH\x00\x12\x0b\n\x01i\x18\x03 \x01(\x03H\x00\x12\x0b\n\x01\x66\x18\x04 \x01(\x02H\x00\x12\x0b\n\x01\x62\x18\x05 \x01(\x08H\x00\x12%\n\x04type\x18\x06 \x01(\x0e\x32\x15.tensorboard.DataTypeH\x00\x12.\n\x05shape\x18\x07 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProtoH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x12\x30\n\x04list\x18\x01 \x01(\x0b\x32 .tensorboard.AttrValue.ListValueH\x00\x12)\n\x04\x66unc\x18\n \x01(\x0b\x32\x19.tensorboard.NameAttrListH\x00\x12\x15\n\x0bplaceholder\x18\t \x01(\tH\x00\x1a\xed\x01\n\tListValue\x12\t\n\x01s\x18\x02 \x03(\x0c\x12\r\n\x01i\x18\x03 \x03(\x03\x42\x02\x10\x01\x12\r\n\x01\x66\x18\x04 \x03(\x02\x42\x02\x10\x01\x12\r\n\x01\x62\x18\x05 \x03(\x08\x42\x02\x10\x01\x12'\n\x04type\x18\x06 \x03(\x0e\x32\x15.tensorboard.DataTypeB\x02\x10\x01\x12,\n\x05shape\x18\x07 \x03(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12(\n\x06tensor\x18\x08 \x03(\x0b\x32\x18.tensorboard.TensorProto\x12'\n\x04\x66unc\x18\t \x03(\x0b\x32\x19.tensorboard.NameAttrListB\x07\n\x05value\"\x94\x01\n\x0cNameAttrList\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04\x61ttr\x18\x02 \x03(\x0b\x32#.tensorboard.NameAttrList.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42\x30\n\x18org.tensorflow.frameworkB\x0f\x41ttrValueProtosP\x01\xf8\x01\x01\x62\x06proto3" + ), + dependencies=[ + tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR, + tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR, + tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR, + ], +) + + +_ATTRVALUE_LISTVALUE = _descriptor.Descriptor( + name="ListValue", + full_name="tensorboard.AttrValue.ListValue", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="s", + full_name="tensorboard.AttrValue.ListValue.s", + index=0, + number=2, + type=12, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="i", + full_name="tensorboard.AttrValue.ListValue.i", + index=1, + number=3, + type=3, + cpp_type=2, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="f", + full_name="tensorboard.AttrValue.ListValue.f", + index=2, + number=4, + type=2, + cpp_type=6, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="b", + full_name="tensorboard.AttrValue.ListValue.b", + index=3, + number=5, + type=8, + cpp_type=7, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="type", + full_name="tensorboard.AttrValue.ListValue.type", + index=4, + number=6, + type=14, + cpp_type=8, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="shape", + full_name="tensorboard.AttrValue.ListValue.shape", + index=5, + number=7, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tensor", + full_name="tensorboard.AttrValue.ListValue.tensor", + index=6, + number=8, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="func", + full_name="tensorboard.AttrValue.ListValue.func", + index=7, + number=9, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=462, + serialized_end=699, +) + +_ATTRVALUE = _descriptor.Descriptor( + name="AttrValue", + full_name="tensorboard.AttrValue", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="s", + full_name="tensorboard.AttrValue.s", + index=0, + number=2, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="i", + full_name="tensorboard.AttrValue.i", + index=1, + number=3, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="f", + full_name="tensorboard.AttrValue.f", + index=2, + number=4, + type=2, + cpp_type=6, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="b", + full_name="tensorboard.AttrValue.b", + index=3, + number=5, + type=8, + cpp_type=7, + label=1, + has_default_value=False, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="type", + full_name="tensorboard.AttrValue.type", + index=4, + number=6, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="shape", + full_name="tensorboard.AttrValue.shape", + index=5, + number=7, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tensor", + full_name="tensorboard.AttrValue.tensor", + index=6, + number=8, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="list", + full_name="tensorboard.AttrValue.list", + index=7, + number=1, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="func", + full_name="tensorboard.AttrValue.func", + index=8, + number=10, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="placeholder", + full_name="tensorboard.AttrValue.placeholder", + index=9, + number=9, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_ATTRVALUE_LISTVALUE,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name="value", full_name="tensorboard.AttrValue.value", index=0, containing_type=None, fields=[] + ), + ], + serialized_start=149, + serialized_end=708, +) + + +_NAMEATTRLIST_ATTRENTRY = _descriptor.Descriptor( + name="AttrEntry", + full_name="tensorboard.NameAttrList.AttrEntry", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="key", + full_name="tensorboard.NameAttrList.AttrEntry.key", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="value", + full_name="tensorboard.NameAttrList.AttrEntry.value", + index=1, + number=2, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b("8\001")), + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=792, + serialized_end=859, +) + +_NAMEATTRLIST = _descriptor.Descriptor( + name="NameAttrList", + full_name="tensorboard.NameAttrList", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="name", + full_name="tensorboard.NameAttrList.name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="attr", + full_name="tensorboard.NameAttrList.attr", + index=1, + number=2, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_NAMEATTRLIST_ATTRENTRY,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=711, + serialized_end=859, +) + +_ATTRVALUE_LISTVALUE.fields_by_name["type"].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_ATTRVALUE_LISTVALUE.fields_by_name[ + "shape" +].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_ATTRVALUE_LISTVALUE.fields_by_name["tensor"].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_ATTRVALUE_LISTVALUE.fields_by_name["func"].message_type = _NAMEATTRLIST +_ATTRVALUE_LISTVALUE.containing_type = _ATTRVALUE +_ATTRVALUE.fields_by_name["type"].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_ATTRVALUE.fields_by_name["shape"].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_ATTRVALUE.fields_by_name["tensor"].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_ATTRVALUE.fields_by_name["list"].message_type = _ATTRVALUE_LISTVALUE +_ATTRVALUE.fields_by_name["func"].message_type = _NAMEATTRLIST +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["s"]) +_ATTRVALUE.fields_by_name["s"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["i"]) +_ATTRVALUE.fields_by_name["i"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["f"]) +_ATTRVALUE.fields_by_name["f"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["b"]) +_ATTRVALUE.fields_by_name["b"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["type"]) +_ATTRVALUE.fields_by_name["type"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["shape"]) +_ATTRVALUE.fields_by_name["shape"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["tensor"]) +_ATTRVALUE.fields_by_name["tensor"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["list"]) +_ATTRVALUE.fields_by_name["list"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["func"]) +_ATTRVALUE.fields_by_name["func"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_ATTRVALUE.oneofs_by_name["value"].fields.append(_ATTRVALUE.fields_by_name["placeholder"]) +_ATTRVALUE.fields_by_name["placeholder"].containing_oneof = _ATTRVALUE.oneofs_by_name["value"] +_NAMEATTRLIST_ATTRENTRY.fields_by_name["value"].message_type = _ATTRVALUE +_NAMEATTRLIST_ATTRENTRY.containing_type = _NAMEATTRLIST +_NAMEATTRLIST.fields_by_name["attr"].message_type = _NAMEATTRLIST_ATTRENTRY +DESCRIPTOR.message_types_by_name["AttrValue"] = _ATTRVALUE +DESCRIPTOR.message_types_by_name["NameAttrList"] = _NAMEATTRLIST +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +AttrValue = _reflection.GeneratedProtocolMessageType( + "AttrValue", + (_message.Message,), + dict( + ListValue=_reflection.GeneratedProtocolMessageType( + "ListValue", + (_message.Message,), + dict( + DESCRIPTOR=_ATTRVALUE_LISTVALUE, + __module__="tensorboardX.src.attr_value_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.AttrValue.ListValue) + ), + ), + DESCRIPTOR=_ATTRVALUE, + __module__="tensorboardX.src.attr_value_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.AttrValue) + ), +) +_sym_db.RegisterMessage(AttrValue) +_sym_db.RegisterMessage(AttrValue.ListValue) + +NameAttrList = _reflection.GeneratedProtocolMessageType( + "NameAttrList", + (_message.Message,), + dict( + AttrEntry=_reflection.GeneratedProtocolMessageType( + "AttrEntry", + (_message.Message,), + dict( + DESCRIPTOR=_NAMEATTRLIST_ATTRENTRY, + __module__="tensorboardX.src.attr_value_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList.AttrEntry) + ), + ), + DESCRIPTOR=_NAMEATTRLIST, + __module__="tensorboardX.src.attr_value_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList) + ), +) +_sym_db.RegisterMessage(NameAttrList) +_sym_db.RegisterMessage(NameAttrList.AttrEntry) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\017AttrValueProtosP\001\370\001\001") +) +_ATTRVALUE_LISTVALUE.fields_by_name["i"].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name["i"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_ATTRVALUE_LISTVALUE.fields_by_name["f"].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name["f"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_ATTRVALUE_LISTVALUE.fields_by_name["b"].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name["b"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_ATTRVALUE_LISTVALUE.fields_by_name["type"].has_options = True +_ATTRVALUE_LISTVALUE.fields_by_name["type"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_NAMEATTRLIST_ATTRENTRY.has_options = True +_NAMEATTRLIST_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b("8\001")) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/event_pb2.py b/tensorboardX/src/event_pb2.py new file mode 100644 index 0000000..24d2c99 --- /dev/null +++ b/tensorboardX/src/event_pb2.py @@ -0,0 +1,506 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/event.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import summary_pb2 as tensorboardX_dot_src_dot_summary__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/event.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1ctensorboardX/src/event.proto\x12\x0btensorboard\x1a\x1etensorboardX/src/summary.proto"\xbf\x02\n\x05\x45vent\x12\x11\n\twall_time\x18\x01 \x01(\x01\x12\x0c\n\x04step\x18\x02 \x01(\x03\x12\x16\n\x0c\x66ile_version\x18\x03 \x01(\tH\x00\x12\x13\n\tgraph_def\x18\x04 \x01(\x0cH\x00\x12\'\n\x07summary\x18\x05 \x01(\x0b\x32\x14.tensorboard.SummaryH\x00\x12.\n\x0blog_message\x18\x06 \x01(\x0b\x32\x17.tensorboard.LogMessageH\x00\x12.\n\x0bsession_log\x18\x07 \x01(\x0b\x32\x17.tensorboard.SessionLogH\x00\x12=\n\x13tagged_run_metadata\x18\x08 \x01(\x0b\x32\x1e.tensorboard.TaggedRunMetadataH\x00\x12\x18\n\x0emeta_graph_def\x18\t \x01(\x0cH\x00\x42\x06\n\x04what"\x96\x01\n\nLogMessage\x12,\n\x05level\x18\x01 \x01(\x0e\x32\x1d.tensorboard.LogMessage.Level\x12\x0f\n\x07message\x18\x02 \x01(\t"I\n\x05Level\x12\x0b\n\x07UNKNOWN\x10\x00\x12\t\n\x05\x44\x45\x42UG\x10\n\x12\x08\n\x04INFO\x10\x14\x12\x08\n\x04WARN\x10\x1e\x12\t\n\x05\x45RROR\x10(\x12\t\n\x05\x46\x41TAL\x10\x32"\xb7\x01\n\nSessionLog\x12\x35\n\x06status\x18\x01 \x01(\x0e\x32%.tensorboard.SessionLog.SessionStatus\x12\x17\n\x0f\x63heckpoint_path\x18\x02 \x01(\t\x12\x0b\n\x03msg\x18\x03 \x01(\t"L\n\rSessionStatus\x12\x16\n\x12STATUS_UNSPECIFIED\x10\x00\x12\t\n\x05START\x10\x01\x12\x08\n\x04STOP\x10\x02\x12\x0e\n\nCHECKPOINT\x10\x03"6\n\x11TaggedRunMetadata\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x14\n\x0crun_metadata\x18\x02 \x01(\x0c\x42\'\n\x13org.tensorflow.utilB\x0b\x45ventProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), + dependencies=[tensorboardX_dot_src_dot_summary__pb2.DESCRIPTOR,], +) + + +_LOGMESSAGE_LEVEL = _descriptor.EnumDescriptor( + name="Level", + full_name="tensorboard.LogMessage.Level", + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor(name="UNKNOWN", index=0, number=0, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DEBUG", index=1, number=10, options=None, type=None), + _descriptor.EnumValueDescriptor(name="INFO", index=2, number=20, options=None, type=None), + _descriptor.EnumValueDescriptor(name="WARN", index=3, number=30, options=None, type=None), + _descriptor.EnumValueDescriptor(name="ERROR", index=4, number=40, options=None, type=None), + _descriptor.EnumValueDescriptor(name="FATAL", index=5, number=50, options=None, type=None), + ], + containing_type=None, + options=None, + serialized_start=477, + serialized_end=550, +) +_sym_db.RegisterEnumDescriptor(_LOGMESSAGE_LEVEL) + +_SESSIONLOG_SESSIONSTATUS = _descriptor.EnumDescriptor( + name="SessionStatus", + full_name="tensorboard.SessionLog.SessionStatus", + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor(name="STATUS_UNSPECIFIED", index=0, number=0, options=None, type=None), + _descriptor.EnumValueDescriptor(name="START", index=1, number=1, options=None, type=None), + _descriptor.EnumValueDescriptor(name="STOP", index=2, number=2, options=None, type=None), + _descriptor.EnumValueDescriptor(name="CHECKPOINT", index=3, number=3, options=None, type=None), + ], + containing_type=None, + options=None, + serialized_start=660, + serialized_end=736, +) +_sym_db.RegisterEnumDescriptor(_SESSIONLOG_SESSIONSTATUS) + + +_EVENT = _descriptor.Descriptor( + name="Event", + full_name="tensorboard.Event", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="wall_time", + full_name="tensorboard.Event.wall_time", + index=0, + number=1, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="step", + full_name="tensorboard.Event.step", + index=1, + number=2, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="file_version", + full_name="tensorboard.Event.file_version", + index=2, + number=3, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="graph_def", + full_name="tensorboard.Event.graph_def", + index=3, + number=4, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="summary", + full_name="tensorboard.Event.summary", + index=4, + number=5, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="log_message", + full_name="tensorboard.Event.log_message", + index=5, + number=6, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="session_log", + full_name="tensorboard.Event.session_log", + index=6, + number=7, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tagged_run_metadata", + full_name="tensorboard.Event.tagged_run_metadata", + index=7, + number=8, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="meta_graph_def", + full_name="tensorboard.Event.meta_graph_def", + index=8, + number=9, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name="what", full_name="tensorboard.Event.what", index=0, containing_type=None, fields=[] + ), + ], + serialized_start=78, + serialized_end=397, +) + + +_LOGMESSAGE = _descriptor.Descriptor( + name="LogMessage", + full_name="tensorboard.LogMessage", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="level", + full_name="tensorboard.LogMessage.level", + index=0, + number=1, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="message", + full_name="tensorboard.LogMessage.message", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[_LOGMESSAGE_LEVEL,], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=400, + serialized_end=550, +) + + +_SESSIONLOG = _descriptor.Descriptor( + name="SessionLog", + full_name="tensorboard.SessionLog", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="status", + full_name="tensorboard.SessionLog.status", + index=0, + number=1, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="checkpoint_path", + full_name="tensorboard.SessionLog.checkpoint_path", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="msg", + full_name="tensorboard.SessionLog.msg", + index=2, + number=3, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[_SESSIONLOG_SESSIONSTATUS,], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=553, + serialized_end=736, +) + + +_TAGGEDRUNMETADATA = _descriptor.Descriptor( + name="TaggedRunMetadata", + full_name="tensorboard.TaggedRunMetadata", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="tag", + full_name="tensorboard.TaggedRunMetadata.tag", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="run_metadata", + full_name="tensorboard.TaggedRunMetadata.run_metadata", + index=1, + number=2, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=738, + serialized_end=792, +) + +_EVENT.fields_by_name["summary"].message_type = tensorboardX_dot_src_dot_summary__pb2._SUMMARY +_EVENT.fields_by_name["log_message"].message_type = _LOGMESSAGE +_EVENT.fields_by_name["session_log"].message_type = _SESSIONLOG +_EVENT.fields_by_name["tagged_run_metadata"].message_type = _TAGGEDRUNMETADATA +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["file_version"]) +_EVENT.fields_by_name["file_version"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["graph_def"]) +_EVENT.fields_by_name["graph_def"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["summary"]) +_EVENT.fields_by_name["summary"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["log_message"]) +_EVENT.fields_by_name["log_message"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["session_log"]) +_EVENT.fields_by_name["session_log"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["tagged_run_metadata"]) +_EVENT.fields_by_name["tagged_run_metadata"].containing_oneof = _EVENT.oneofs_by_name["what"] +_EVENT.oneofs_by_name["what"].fields.append(_EVENT.fields_by_name["meta_graph_def"]) +_EVENT.fields_by_name["meta_graph_def"].containing_oneof = _EVENT.oneofs_by_name["what"] +_LOGMESSAGE.fields_by_name["level"].enum_type = _LOGMESSAGE_LEVEL +_LOGMESSAGE_LEVEL.containing_type = _LOGMESSAGE +_SESSIONLOG.fields_by_name["status"].enum_type = _SESSIONLOG_SESSIONSTATUS +_SESSIONLOG_SESSIONSTATUS.containing_type = _SESSIONLOG +DESCRIPTOR.message_types_by_name["Event"] = _EVENT +DESCRIPTOR.message_types_by_name["LogMessage"] = _LOGMESSAGE +DESCRIPTOR.message_types_by_name["SessionLog"] = _SESSIONLOG +DESCRIPTOR.message_types_by_name["TaggedRunMetadata"] = _TAGGEDRUNMETADATA +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +Event = _reflection.GeneratedProtocolMessageType( + "Event", + (_message.Message,), + dict( + DESCRIPTOR=_EVENT, + __module__="tensorboardX.src.event_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.Event) + ), +) +_sym_db.RegisterMessage(Event) + +LogMessage = _reflection.GeneratedProtocolMessageType( + "LogMessage", + (_message.Message,), + dict( + DESCRIPTOR=_LOGMESSAGE, + __module__="tensorboardX.src.event_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.LogMessage) + ), +) +_sym_db.RegisterMessage(LogMessage) + +SessionLog = _reflection.GeneratedProtocolMessageType( + "SessionLog", + (_message.Message,), + dict( + DESCRIPTOR=_SESSIONLOG, + __module__="tensorboardX.src.event_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.SessionLog) + ), +) +_sym_db.RegisterMessage(SessionLog) + +TaggedRunMetadata = _reflection.GeneratedProtocolMessageType( + "TaggedRunMetadata", + (_message.Message,), + dict( + DESCRIPTOR=_TAGGEDRUNMETADATA, + __module__="tensorboardX.src.event_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.TaggedRunMetadata) + ), +) +_sym_db.RegisterMessage(TaggedRunMetadata) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\023org.tensorflow.utilB\013EventProtosP\001\370\001\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/graph_pb2.py b/tensorboardX/src/graph_pb2.py new file mode 100644 index 0000000..ea98390 --- /dev/null +++ b/tensorboardX/src/graph_pb2.py @@ -0,0 +1,129 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/graph.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import node_def_pb2 as tensorboardX_dot_src_dot_node__def__pb2 +from tensorboardX.src import versions_pb2 as tensorboardX_dot_src_dot_versions__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/graph.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1ctensorboardX/src/graph.proto\x12\x0btensorboard\x1a\x1ftensorboardX/src/node_def.proto\x1a\x1ftensorboardX/src/versions.proto"n\n\x08GraphDef\x12"\n\x04node\x18\x01 \x03(\x0b\x32\x14.tensorboard.NodeDef\x12)\n\x08versions\x18\x04 \x01(\x0b\x32\x17.tensorboard.VersionDef\x12\x13\n\x07version\x18\x03 \x01(\x05\x42\x02\x18\x01\x42,\n\x18org.tensorflow.frameworkB\x0bGraphProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), + dependencies=[ + tensorboardX_dot_src_dot_node__def__pb2.DESCRIPTOR, + tensorboardX_dot_src_dot_versions__pb2.DESCRIPTOR, + ], +) + + +_GRAPHDEF = _descriptor.Descriptor( + name="GraphDef", + full_name="tensorboard.GraphDef", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="node", + full_name="tensorboard.GraphDef.node", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="versions", + full_name="tensorboard.GraphDef.versions", + index=1, + number=4, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="version", + full_name="tensorboard.GraphDef.version", + index=2, + number=3, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\030\001")), + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=111, + serialized_end=221, +) + +_GRAPHDEF.fields_by_name["node"].message_type = tensorboardX_dot_src_dot_node__def__pb2._NODEDEF +_GRAPHDEF.fields_by_name["versions"].message_type = tensorboardX_dot_src_dot_versions__pb2._VERSIONDEF +DESCRIPTOR.message_types_by_name["GraphDef"] = _GRAPHDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +GraphDef = _reflection.GeneratedProtocolMessageType( + "GraphDef", + (_message.Message,), + dict( + DESCRIPTOR=_GRAPHDEF, + __module__="tensorboardX.src.graph_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.GraphDef) + ), +) +_sym_db.RegisterMessage(GraphDef) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\013GraphProtosP\001\370\001\001") +) +_GRAPHDEF.fields_by_name["version"].has_options = True +_GRAPHDEF.fields_by_name["version"]._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\030\001")) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/node_def_pb2.py b/tensorboardX/src/node_def_pb2.py new file mode 100644 index 0000000..e7d3a14 --- /dev/null +++ b/tensorboardX/src/node_def_pb2.py @@ -0,0 +1,224 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/node_def.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import attr_value_pb2 as tensorboardX_dot_src_dot_attr__value__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/node_def.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1ftensorboardX/src/node_def.proto\x12\x0btensorboard\x1a!tensorboardX/src/attr_value.proto"\xb5\x01\n\x07NodeDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\n\n\x02op\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x03(\t\x12\x0e\n\x06\x64\x65vice\x18\x04 \x01(\t\x12,\n\x04\x61ttr\x18\x05 \x03(\x0b\x32\x1e.tensorboard.NodeDef.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42*\n\x18org.tensorflow.frameworkB\tNodeProtoP\x01\xf8\x01\x01\x62\x06proto3' + ), + dependencies=[tensorboardX_dot_src_dot_attr__value__pb2.DESCRIPTOR,], +) + + +_NODEDEF_ATTRENTRY = _descriptor.Descriptor( + name="AttrEntry", + full_name="tensorboard.NodeDef.AttrEntry", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="key", + full_name="tensorboard.NodeDef.AttrEntry.key", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="value", + full_name="tensorboard.NodeDef.AttrEntry.value", + index=1, + number=2, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b("8\001")), + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=198, + serialized_end=265, +) + +_NODEDEF = _descriptor.Descriptor( + name="NodeDef", + full_name="tensorboard.NodeDef", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="name", + full_name="tensorboard.NodeDef.name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="op", + full_name="tensorboard.NodeDef.op", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="input", + full_name="tensorboard.NodeDef.input", + index=2, + number=3, + type=9, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="device", + full_name="tensorboard.NodeDef.device", + index=3, + number=4, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="attr", + full_name="tensorboard.NodeDef.attr", + index=4, + number=5, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_NODEDEF_ATTRENTRY,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=84, + serialized_end=265, +) + +_NODEDEF_ATTRENTRY.fields_by_name["value"].message_type = tensorboardX_dot_src_dot_attr__value__pb2._ATTRVALUE +_NODEDEF_ATTRENTRY.containing_type = _NODEDEF +_NODEDEF.fields_by_name["attr"].message_type = _NODEDEF_ATTRENTRY +DESCRIPTOR.message_types_by_name["NodeDef"] = _NODEDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +NodeDef = _reflection.GeneratedProtocolMessageType( + "NodeDef", + (_message.Message,), + dict( + AttrEntry=_reflection.GeneratedProtocolMessageType( + "AttrEntry", + (_message.Message,), + dict( + DESCRIPTOR=_NODEDEF_ATTRENTRY, + __module__="tensorboardX.src.node_def_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.NodeDef.AttrEntry) + ), + ), + DESCRIPTOR=_NODEDEF, + __module__="tensorboardX.src.node_def_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.NodeDef) + ), +) +_sym_db.RegisterMessage(NodeDef) +_sym_db.RegisterMessage(NodeDef.AttrEntry) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\tNodeProtoP\001\370\001\001") +) +_NODEDEF_ATTRENTRY.has_options = True +_NODEDEF_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b("8\001")) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/plugin_pr_curve_pb2.py b/tensorboardX/src/plugin_pr_curve_pb2.py new file mode 100644 index 0000000..d0e0d9b --- /dev/null +++ b/tensorboardX/src/plugin_pr_curve_pb2.py @@ -0,0 +1,97 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/plugin_pr_curve.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/plugin_pr_curve.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n&tensorboardX/src/plugin_pr_curve.proto\x12\x0btensorboard"<\n\x11PrCurvePluginData\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0enum_thresholds\x18\x02 \x01(\rb\x06proto3' + ), +) + + +_PRCURVEPLUGINDATA = _descriptor.Descriptor( + name="PrCurvePluginData", + full_name="tensorboard.PrCurvePluginData", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="version", + full_name="tensorboard.PrCurvePluginData.version", + index=0, + number=1, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="num_thresholds", + full_name="tensorboard.PrCurvePluginData.num_thresholds", + index=1, + number=2, + type=13, + cpp_type=3, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=55, + serialized_end=115, +) + +DESCRIPTOR.message_types_by_name["PrCurvePluginData"] = _PRCURVEPLUGINDATA +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +PrCurvePluginData = _reflection.GeneratedProtocolMessageType( + "PrCurvePluginData", + (_message.Message,), + dict( + DESCRIPTOR=_PRCURVEPLUGINDATA, + __module__="tensorboardX.src.plugin_pr_curve_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.PrCurvePluginData) + ), +) +_sym_db.RegisterMessage(PrCurvePluginData) + + +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/resource_handle_pb2.py b/tensorboardX/src/resource_handle_pb2.py new file mode 100644 index 0000000..54bae9a --- /dev/null +++ b/tensorboardX/src/resource_handle_pb2.py @@ -0,0 +1,152 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/resource_handle.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/resource_handle.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n&tensorboardX/src/resource_handle.proto\x12\x0btensorboard"r\n\x13ResourceHandleProto\x12\x0e\n\x06\x64\x65vice\x18\x01 \x01(\t\x12\x11\n\tcontainer\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x11\n\thash_code\x18\x04 \x01(\x04\x12\x17\n\x0fmaybe_type_name\x18\x05 \x01(\tB/\n\x18org.tensorflow.frameworkB\x0eResourceHandleP\x01\xf8\x01\x01\x62\x06proto3' + ), +) + + +_RESOURCEHANDLEPROTO = _descriptor.Descriptor( + name="ResourceHandleProto", + full_name="tensorboard.ResourceHandleProto", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="device", + full_name="tensorboard.ResourceHandleProto.device", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="container", + full_name="tensorboard.ResourceHandleProto.container", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="name", + full_name="tensorboard.ResourceHandleProto.name", + index=2, + number=3, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="hash_code", + full_name="tensorboard.ResourceHandleProto.hash_code", + index=3, + number=4, + type=4, + cpp_type=4, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="maybe_type_name", + full_name="tensorboard.ResourceHandleProto.maybe_type_name", + index=4, + number=5, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=55, + serialized_end=169, +) + +DESCRIPTOR.message_types_by_name["ResourceHandleProto"] = _RESOURCEHANDLEPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +ResourceHandleProto = _reflection.GeneratedProtocolMessageType( + "ResourceHandleProto", + (_message.Message,), + dict( + DESCRIPTOR=_RESOURCEHANDLEPROTO, + __module__="tensorboardX.src.resource_handle_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.ResourceHandleProto) + ), +) +_sym_db.RegisterMessage(ResourceHandleProto) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\016ResourceHandleP\001\370\001\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/summary_pb2.py b/tensorboardX/src/summary_pb2.py new file mode 100644 index 0000000..fe92ec1 --- /dev/null +++ b/tensorboardX/src/summary_pb2.py @@ -0,0 +1,834 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/summary.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/summary.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1etensorboardX/src/summary.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto"\'\n\x12SummaryDescription\x12\x11\n\ttype_hint\x18\x01 \x01(\t"\x87\x01\n\x0eHistogramProto\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03num\x18\x03 \x01(\x01\x12\x0b\n\x03sum\x18\x04 \x01(\x01\x12\x13\n\x0bsum_squares\x18\x05 \x01(\x01\x12\x18\n\x0c\x62ucket_limit\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x12\n\x06\x62ucket\x18\x07 \x03(\x01\x42\x02\x10\x01"\x83\x01\n\x0fSummaryMetadata\x12<\n\x0bplugin_data\x18\x01 \x03(\x0b\x32\'.tensorboard.SummaryMetadata.PluginData\x1a\x32\n\nPluginData\x12\x13\n\x0bplugin_name\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t"\xe4\x04\n\x07Summary\x12)\n\x05value\x18\x01 \x03(\x0b\x32\x1a.tensorboard.Summary.Value\x1aX\n\x05Image\x12\x0e\n\x06height\x18\x01 \x01(\x05\x12\r\n\x05width\x18\x02 \x01(\x05\x12\x12\n\ncolorspace\x18\x03 \x01(\x05\x12\x1c\n\x14\x65ncoded_image_string\x18\x04 \x01(\x0c\x1a}\n\x05\x41udio\x12\x13\n\x0bsample_rate\x18\x01 \x01(\x02\x12\x14\n\x0cnum_channels\x18\x02 \x01(\x03\x12\x15\n\rlength_frames\x18\x03 \x01(\x03\x12\x1c\n\x14\x65ncoded_audio_string\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\x1a\xd4\x02\n\x05Value\x12\x11\n\tnode_name\x18\x07 \x01(\t\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12.\n\x08metadata\x18\t \x01(\x0b\x32\x1c.tensorboard.SummaryMetadata\x12\x16\n\x0csimple_value\x18\x02 \x01(\x02H\x00\x12&\n\x1cobsolete_old_style_histogram\x18\x03 \x01(\x0cH\x00\x12+\n\x05image\x18\x04 \x01(\x0b\x32\x1a.tensorboard.Summary.ImageH\x00\x12,\n\x05histo\x18\x05 \x01(\x0b\x32\x1b.tensorboard.HistogramProtoH\x00\x12+\n\x05\x61udio\x18\x06 \x01(\x0b\x32\x1a.tensorboard.Summary.AudioH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x42\x07\n\x05valueB.\n\x18org.tensorflow.frameworkB\rSummaryProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), + dependencies=[tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR,], +) + + +_SUMMARYDESCRIPTION = _descriptor.Descriptor( + name="SummaryDescription", + full_name="tensorboard.SummaryDescription", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="type_hint", + full_name="tensorboard.SummaryDescription.type_hint", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=78, + serialized_end=117, +) + + +_HISTOGRAMPROTO = _descriptor.Descriptor( + name="HistogramProto", + full_name="tensorboard.HistogramProto", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="min", + full_name="tensorboard.HistogramProto.min", + index=0, + number=1, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="max", + full_name="tensorboard.HistogramProto.max", + index=1, + number=2, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="num", + full_name="tensorboard.HistogramProto.num", + index=2, + number=3, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="sum", + full_name="tensorboard.HistogramProto.sum", + index=3, + number=4, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="sum_squares", + full_name="tensorboard.HistogramProto.sum_squares", + index=4, + number=5, + type=1, + cpp_type=5, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="bucket_limit", + full_name="tensorboard.HistogramProto.bucket_limit", + index=5, + number=6, + type=1, + cpp_type=5, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="bucket", + full_name="tensorboard.HistogramProto.bucket", + index=6, + number=7, + type=1, + cpp_type=5, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=120, + serialized_end=255, +) + + +_SUMMARYMETADATA_PLUGINDATA = _descriptor.Descriptor( + name="PluginData", + full_name="tensorboard.SummaryMetadata.PluginData", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="plugin_name", + full_name="tensorboard.SummaryMetadata.PluginData.plugin_name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="content", + full_name="tensorboard.SummaryMetadata.PluginData.content", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=339, + serialized_end=389, +) + +_SUMMARYMETADATA = _descriptor.Descriptor( + name="SummaryMetadata", + full_name="tensorboard.SummaryMetadata", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="plugin_data", + full_name="tensorboard.SummaryMetadata.plugin_data", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_SUMMARYMETADATA_PLUGINDATA,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=258, + serialized_end=389, +) + + +_SUMMARY_IMAGE = _descriptor.Descriptor( + name="Image", + full_name="tensorboard.Summary.Image", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="height", + full_name="tensorboard.Summary.Image.height", + index=0, + number=1, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="width", + full_name="tensorboard.Summary.Image.width", + index=1, + number=2, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="colorspace", + full_name="tensorboard.Summary.Image.colorspace", + index=2, + number=3, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="encoded_image_string", + full_name="tensorboard.Summary.Image.encoded_image_string", + index=3, + number=4, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=446, + serialized_end=534, +) + +_SUMMARY_AUDIO = _descriptor.Descriptor( + name="Audio", + full_name="tensorboard.Summary.Audio", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="sample_rate", + full_name="tensorboard.Summary.Audio.sample_rate", + index=0, + number=1, + type=2, + cpp_type=6, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="num_channels", + full_name="tensorboard.Summary.Audio.num_channels", + index=1, + number=2, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="length_frames", + full_name="tensorboard.Summary.Audio.length_frames", + index=2, + number=3, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="encoded_audio_string", + full_name="tensorboard.Summary.Audio.encoded_audio_string", + index=3, + number=4, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="content_type", + full_name="tensorboard.Summary.Audio.content_type", + index=4, + number=5, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=536, + serialized_end=661, +) + +_SUMMARY_VALUE = _descriptor.Descriptor( + name="Value", + full_name="tensorboard.Summary.Value", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="node_name", + full_name="tensorboard.Summary.Value.node_name", + index=0, + number=7, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tag", + full_name="tensorboard.Summary.Value.tag", + index=1, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="metadata", + full_name="tensorboard.Summary.Value.metadata", + index=2, + number=9, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="simple_value", + full_name="tensorboard.Summary.Value.simple_value", + index=3, + number=2, + type=2, + cpp_type=6, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="obsolete_old_style_histogram", + full_name="tensorboard.Summary.Value.obsolete_old_style_histogram", + index=4, + number=3, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="image", + full_name="tensorboard.Summary.Value.image", + index=5, + number=4, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="histo", + full_name="tensorboard.Summary.Value.histo", + index=6, + number=5, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="audio", + full_name="tensorboard.Summary.Value.audio", + index=7, + number=6, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tensor", + full_name="tensorboard.Summary.Value.tensor", + index=8, + number=8, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[ + _descriptor.OneofDescriptor( + name="value", full_name="tensorboard.Summary.Value.value", index=0, containing_type=None, fields=[] + ), + ], + serialized_start=664, + serialized_end=1004, +) + +_SUMMARY = _descriptor.Descriptor( + name="Summary", + full_name="tensorboard.Summary", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="value", + full_name="tensorboard.Summary.value", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_SUMMARY_IMAGE, _SUMMARY_AUDIO, _SUMMARY_VALUE,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=392, + serialized_end=1004, +) + +_SUMMARYMETADATA_PLUGINDATA.containing_type = _SUMMARYMETADATA +_SUMMARYMETADATA.fields_by_name["plugin_data"].message_type = _SUMMARYMETADATA_PLUGINDATA +_SUMMARY_IMAGE.containing_type = _SUMMARY +_SUMMARY_AUDIO.containing_type = _SUMMARY +_SUMMARY_VALUE.fields_by_name["metadata"].message_type = _SUMMARYMETADATA +_SUMMARY_VALUE.fields_by_name["image"].message_type = _SUMMARY_IMAGE +_SUMMARY_VALUE.fields_by_name["histo"].message_type = _HISTOGRAMPROTO +_SUMMARY_VALUE.fields_by_name["audio"].message_type = _SUMMARY_AUDIO +_SUMMARY_VALUE.fields_by_name["tensor"].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO +_SUMMARY_VALUE.containing_type = _SUMMARY +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["simple_value"]) +_SUMMARY_VALUE.fields_by_name["simple_value"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["obsolete_old_style_histogram"]) +_SUMMARY_VALUE.fields_by_name["obsolete_old_style_histogram"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["image"]) +_SUMMARY_VALUE.fields_by_name["image"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["histo"]) +_SUMMARY_VALUE.fields_by_name["histo"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["audio"]) +_SUMMARY_VALUE.fields_by_name["audio"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY_VALUE.oneofs_by_name["value"].fields.append(_SUMMARY_VALUE.fields_by_name["tensor"]) +_SUMMARY_VALUE.fields_by_name["tensor"].containing_oneof = _SUMMARY_VALUE.oneofs_by_name["value"] +_SUMMARY.fields_by_name["value"].message_type = _SUMMARY_VALUE +DESCRIPTOR.message_types_by_name["SummaryDescription"] = _SUMMARYDESCRIPTION +DESCRIPTOR.message_types_by_name["HistogramProto"] = _HISTOGRAMPROTO +DESCRIPTOR.message_types_by_name["SummaryMetadata"] = _SUMMARYMETADATA +DESCRIPTOR.message_types_by_name["Summary"] = _SUMMARY +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +SummaryDescription = _reflection.GeneratedProtocolMessageType( + "SummaryDescription", + (_message.Message,), + dict( + DESCRIPTOR=_SUMMARYDESCRIPTION, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.SummaryDescription) + ), +) +_sym_db.RegisterMessage(SummaryDescription) + +HistogramProto = _reflection.GeneratedProtocolMessageType( + "HistogramProto", + (_message.Message,), + dict( + DESCRIPTOR=_HISTOGRAMPROTO, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.HistogramProto) + ), +) +_sym_db.RegisterMessage(HistogramProto) + +SummaryMetadata = _reflection.GeneratedProtocolMessageType( + "SummaryMetadata", + (_message.Message,), + dict( + PluginData=_reflection.GeneratedProtocolMessageType( + "PluginData", + (_message.Message,), + dict( + DESCRIPTOR=_SUMMARYMETADATA_PLUGINDATA, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata.PluginData) + ), + ), + DESCRIPTOR=_SUMMARYMETADATA, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata) + ), +) +_sym_db.RegisterMessage(SummaryMetadata) +_sym_db.RegisterMessage(SummaryMetadata.PluginData) + +Summary = _reflection.GeneratedProtocolMessageType( + "Summary", + (_message.Message,), + dict( + Image=_reflection.GeneratedProtocolMessageType( + "Image", + (_message.Message,), + dict( + DESCRIPTOR=_SUMMARY_IMAGE, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Image) + ), + ), + Audio=_reflection.GeneratedProtocolMessageType( + "Audio", + (_message.Message,), + dict( + DESCRIPTOR=_SUMMARY_AUDIO, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Audio) + ), + ), + Value=_reflection.GeneratedProtocolMessageType( + "Value", + (_message.Message,), + dict( + DESCRIPTOR=_SUMMARY_VALUE, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.Summary.Value) + ), + ), + DESCRIPTOR=_SUMMARY, + __module__="tensorboardX.src.summary_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.Summary) + ), +) +_sym_db.RegisterMessage(Summary) +_sym_db.RegisterMessage(Summary.Image) +_sym_db.RegisterMessage(Summary.Audio) +_sym_db.RegisterMessage(Summary.Value) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\rSummaryProtosP\001\370\001\001") +) +_HISTOGRAMPROTO.fields_by_name["bucket_limit"].has_options = True +_HISTOGRAMPROTO.fields_by_name["bucket_limit"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_HISTOGRAMPROTO.fields_by_name["bucket"].has_options = True +_HISTOGRAMPROTO.fields_by_name["bucket"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/tensor_pb2.py b/tensorboardX/src/tensor_pb2.py new file mode 100644 index 0000000..f572208 --- /dev/null +++ b/tensorboardX/src/tensor_pb2.py @@ -0,0 +1,351 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/tensor.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from tensorboardX.src import resource_handle_pb2 as tensorboardX_dot_src_dot_resource__handle__pb2 +from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2 +from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2 + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/tensor.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1dtensorboardX/src/tensor.proto\x12\x0btensorboard\x1a&tensorboardX/src/resource_handle.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto"\xa6\x03\n\x0bTensorProto\x12$\n\x05\x64type\x18\x01 \x01(\x0e\x32\x15.tensorboard.DataType\x12\x33\n\x0ctensor_shape\x18\x02 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12\x16\n\x0eversion_number\x18\x03 \x01(\x05\x12\x16\n\x0etensor_content\x18\x04 \x01(\x0c\x12\x14\n\x08half_val\x18\r \x03(\x05\x42\x02\x10\x01\x12\x15\n\tfloat_val\x18\x05 \x03(\x02\x42\x02\x10\x01\x12\x16\n\ndouble_val\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x13\n\x07int_val\x18\x07 \x03(\x05\x42\x02\x10\x01\x12\x12\n\nstring_val\x18\x08 \x03(\x0c\x12\x18\n\x0cscomplex_val\x18\t \x03(\x02\x42\x02\x10\x01\x12\x15\n\tint64_val\x18\n \x03(\x03\x42\x02\x10\x01\x12\x14\n\x08\x62ool_val\x18\x0b \x03(\x08\x42\x02\x10\x01\x12\x18\n\x0c\x64\x63omplex_val\x18\x0c \x03(\x01\x42\x02\x10\x01\x12=\n\x13resource_handle_val\x18\x0e \x03(\x0b\x32 .tensorboard.ResourceHandleProtoB-\n\x18org.tensorflow.frameworkB\x0cTensorProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), + dependencies=[ + tensorboardX_dot_src_dot_resource__handle__pb2.DESCRIPTOR, + tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR, + tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR, + ], +) + + +_TENSORPROTO = _descriptor.Descriptor( + name="TensorProto", + full_name="tensorboard.TensorProto", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="dtype", + full_name="tensorboard.TensorProto.dtype", + index=0, + number=1, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tensor_shape", + full_name="tensorboard.TensorProto.tensor_shape", + index=1, + number=2, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="version_number", + full_name="tensorboard.TensorProto.version_number", + index=2, + number=3, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="tensor_content", + full_name="tensorboard.TensorProto.tensor_content", + index=3, + number=4, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b(""), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="half_val", + full_name="tensorboard.TensorProto.half_val", + index=4, + number=13, + type=5, + cpp_type=1, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="float_val", + full_name="tensorboard.TensorProto.float_val", + index=5, + number=5, + type=2, + cpp_type=6, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="double_val", + full_name="tensorboard.TensorProto.double_val", + index=6, + number=6, + type=1, + cpp_type=5, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="int_val", + full_name="tensorboard.TensorProto.int_val", + index=7, + number=7, + type=5, + cpp_type=1, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="string_val", + full_name="tensorboard.TensorProto.string_val", + index=8, + number=8, + type=12, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="scomplex_val", + full_name="tensorboard.TensorProto.scomplex_val", + index=9, + number=9, + type=2, + cpp_type=6, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="int64_val", + full_name="tensorboard.TensorProto.int64_val", + index=10, + number=10, + type=3, + cpp_type=2, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="bool_val", + full_name="tensorboard.TensorProto.bool_val", + index=11, + number=11, + type=8, + cpp_type=7, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="dcomplex_val", + full_name="tensorboard.TensorProto.dcomplex_val", + index=12, + number=12, + type=1, + cpp_type=5, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b("\020\001")), + ), + _descriptor.FieldDescriptor( + name="resource_handle_val", + full_name="tensorboard.TensorProto.resource_handle_val", + index=13, + number=14, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=154, + serialized_end=576, +) + +_TENSORPROTO.fields_by_name["dtype"].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE +_TENSORPROTO.fields_by_name["tensor_shape"].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_TENSORPROTO.fields_by_name[ + "resource_handle_val" +].message_type = tensorboardX_dot_src_dot_resource__handle__pb2._RESOURCEHANDLEPROTO +DESCRIPTOR.message_types_by_name["TensorProto"] = _TENSORPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +TensorProto = _reflection.GeneratedProtocolMessageType( + "TensorProto", + (_message.Message,), + dict( + DESCRIPTOR=_TENSORPROTO, + __module__="tensorboardX.src.tensor_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.TensorProto) + ), +) +_sym_db.RegisterMessage(TensorProto) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\014TensorProtosP\001\370\001\001") +) +_TENSORPROTO.fields_by_name["half_val"].has_options = True +_TENSORPROTO.fields_by_name["half_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["float_val"].has_options = True +_TENSORPROTO.fields_by_name["float_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["double_val"].has_options = True +_TENSORPROTO.fields_by_name["double_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["int_val"].has_options = True +_TENSORPROTO.fields_by_name["int_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["scomplex_val"].has_options = True +_TENSORPROTO.fields_by_name["scomplex_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["int64_val"].has_options = True +_TENSORPROTO.fields_by_name["int64_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["bool_val"].has_options = True +_TENSORPROTO.fields_by_name["bool_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +_TENSORPROTO.fields_by_name["dcomplex_val"].has_options = True +_TENSORPROTO.fields_by_name["dcomplex_val"]._options = _descriptor._ParseOptions( + descriptor_pb2.FieldOptions(), _b("\020\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/tensor_shape_pb2.py b/tensorboardX/src/tensor_shape_pb2.py new file mode 100644 index 0000000..6f0b680 --- /dev/null +++ b/tensorboardX/src/tensor_shape_pb2.py @@ -0,0 +1,167 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/tensor_shape.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/tensor_shape.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n#tensorboardX/src/tensor_shape.proto\x12\x0btensorboard"{\n\x10TensorShapeProto\x12.\n\x03\x64im\x18\x02 \x03(\x0b\x32!.tensorboard.TensorShapeProto.Dim\x12\x14\n\x0cunknown_rank\x18\x03 \x01(\x08\x1a!\n\x03\x44im\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\tB2\n\x18org.tensorflow.frameworkB\x11TensorShapeProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), +) + + +_TENSORSHAPEPROTO_DIM = _descriptor.Descriptor( + name="Dim", + full_name="tensorboard.TensorShapeProto.Dim", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="size", + full_name="tensorboard.TensorShapeProto.Dim.size", + index=0, + number=1, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="name", + full_name="tensorboard.TensorShapeProto.Dim.name", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=_b("").decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=142, + serialized_end=175, +) + +_TENSORSHAPEPROTO = _descriptor.Descriptor( + name="TensorShapeProto", + full_name="tensorboard.TensorShapeProto", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="dim", + full_name="tensorboard.TensorShapeProto.dim", + index=0, + number=2, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="unknown_rank", + full_name="tensorboard.TensorShapeProto.unknown_rank", + index=1, + number=3, + type=8, + cpp_type=7, + label=1, + has_default_value=False, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[_TENSORSHAPEPROTO_DIM,], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=52, + serialized_end=175, +) + +_TENSORSHAPEPROTO_DIM.containing_type = _TENSORSHAPEPROTO +_TENSORSHAPEPROTO.fields_by_name["dim"].message_type = _TENSORSHAPEPROTO_DIM +DESCRIPTOR.message_types_by_name["TensorShapeProto"] = _TENSORSHAPEPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +TensorShapeProto = _reflection.GeneratedProtocolMessageType( + "TensorShapeProto", + (_message.Message,), + dict( + Dim=_reflection.GeneratedProtocolMessageType( + "Dim", + (_message.Message,), + dict( + DESCRIPTOR=_TENSORSHAPEPROTO_DIM, + __module__="tensorboardX.src.tensor_shape_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto.Dim) + ), + ), + DESCRIPTOR=_TENSORSHAPEPROTO, + __module__="tensorboardX.src.tensor_shape_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto) + ), +) +_sym_db.RegisterMessage(TensorShapeProto) +_sym_db.RegisterMessage(TensorShapeProto.Dim) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\021TensorShapeProtosP\001\370\001\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/types_pb2.py b/tensorboardX/src/types_pb2.py new file mode 100644 index 0000000..4e3924e --- /dev/null +++ b/tensorboardX/src/types_pb2.py @@ -0,0 +1,135 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/types.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import enum_type_wrapper + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/types.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + "\n\x1ctensorboardX/src/types.proto\x12\x0btensorboard*\xc2\x05\n\x08\x44\x61taType\x12\x0e\n\nDT_INVALID\x10\x00\x12\x0c\n\x08\x44T_FLOAT\x10\x01\x12\r\n\tDT_DOUBLE\x10\x02\x12\x0c\n\x08\x44T_INT32\x10\x03\x12\x0c\n\x08\x44T_UINT8\x10\x04\x12\x0c\n\x08\x44T_INT16\x10\x05\x12\x0b\n\x07\x44T_INT8\x10\x06\x12\r\n\tDT_STRING\x10\x07\x12\x10\n\x0c\x44T_COMPLEX64\x10\x08\x12\x0c\n\x08\x44T_INT64\x10\t\x12\x0b\n\x07\x44T_BOOL\x10\n\x12\x0c\n\x08\x44T_QINT8\x10\x0b\x12\r\n\tDT_QUINT8\x10\x0c\x12\r\n\tDT_QINT32\x10\r\x12\x0f\n\x0b\x44T_BFLOAT16\x10\x0e\x12\r\n\tDT_QINT16\x10\x0f\x12\x0e\n\nDT_QUINT16\x10\x10\x12\r\n\tDT_UINT16\x10\x11\x12\x11\n\rDT_COMPLEX128\x10\x12\x12\x0b\n\x07\x44T_HALF\x10\x13\x12\x0f\n\x0b\x44T_RESOURCE\x10\x14\x12\x10\n\x0c\x44T_FLOAT_REF\x10\x65\x12\x11\n\rDT_DOUBLE_REF\x10\x66\x12\x10\n\x0c\x44T_INT32_REF\x10g\x12\x10\n\x0c\x44T_UINT8_REF\x10h\x12\x10\n\x0c\x44T_INT16_REF\x10i\x12\x0f\n\x0b\x44T_INT8_REF\x10j\x12\x11\n\rDT_STRING_REF\x10k\x12\x14\n\x10\x44T_COMPLEX64_REF\x10l\x12\x10\n\x0c\x44T_INT64_REF\x10m\x12\x0f\n\x0b\x44T_BOOL_REF\x10n\x12\x10\n\x0c\x44T_QINT8_REF\x10o\x12\x11\n\rDT_QUINT8_REF\x10p\x12\x11\n\rDT_QINT32_REF\x10q\x12\x13\n\x0f\x44T_BFLOAT16_REF\x10r\x12\x11\n\rDT_QINT16_REF\x10s\x12\x12\n\x0e\x44T_QUINT16_REF\x10t\x12\x11\n\rDT_UINT16_REF\x10u\x12\x15\n\x11\x44T_COMPLEX128_REF\x10v\x12\x0f\n\x0b\x44T_HALF_REF\x10w\x12\x13\n\x0f\x44T_RESOURCE_REF\x10xB,\n\x18org.tensorflow.frameworkB\x0bTypesProtosP\x01\xf8\x01\x01\x62\x06proto3" + ), +) + +_DATATYPE = _descriptor.EnumDescriptor( + name="DataType", + full_name="tensorboard.DataType", + filename=None, + file=DESCRIPTOR, + values=[ + _descriptor.EnumValueDescriptor(name="DT_INVALID", index=0, number=0, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_FLOAT", index=1, number=1, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_DOUBLE", index=2, number=2, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT32", index=3, number=3, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_UINT8", index=4, number=4, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT16", index=5, number=5, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT8", index=6, number=6, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_STRING", index=7, number=7, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_COMPLEX64", index=8, number=8, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT64", index=9, number=9, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_BOOL", index=10, number=10, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT8", index=11, number=11, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QUINT8", index=12, number=12, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT32", index=13, number=13, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_BFLOAT16", index=14, number=14, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT16", index=15, number=15, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QUINT16", index=16, number=16, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_UINT16", index=17, number=17, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_COMPLEX128", index=18, number=18, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_HALF", index=19, number=19, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_RESOURCE", index=20, number=20, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_FLOAT_REF", index=21, number=101, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_DOUBLE_REF", index=22, number=102, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT32_REF", index=23, number=103, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_UINT8_REF", index=24, number=104, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT16_REF", index=25, number=105, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT8_REF", index=26, number=106, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_STRING_REF", index=27, number=107, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_COMPLEX64_REF", index=28, number=108, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_INT64_REF", index=29, number=109, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_BOOL_REF", index=30, number=110, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT8_REF", index=31, number=111, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QUINT8_REF", index=32, number=112, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT32_REF", index=33, number=113, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_BFLOAT16_REF", index=34, number=114, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QINT16_REF", index=35, number=115, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_QUINT16_REF", index=36, number=116, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_UINT16_REF", index=37, number=117, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_COMPLEX128_REF", index=38, number=118, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_HALF_REF", index=39, number=119, options=None, type=None), + _descriptor.EnumValueDescriptor(name="DT_RESOURCE_REF", index=40, number=120, options=None, type=None), + ], + containing_type=None, + options=None, + serialized_start=46, + serialized_end=752, +) +_sym_db.RegisterEnumDescriptor(_DATATYPE) + +DataType = enum_type_wrapper.EnumTypeWrapper(_DATATYPE) +DT_INVALID = 0 +DT_FLOAT = 1 +DT_DOUBLE = 2 +DT_INT32 = 3 +DT_UINT8 = 4 +DT_INT16 = 5 +DT_INT8 = 6 +DT_STRING = 7 +DT_COMPLEX64 = 8 +DT_INT64 = 9 +DT_BOOL = 10 +DT_QINT8 = 11 +DT_QUINT8 = 12 +DT_QINT32 = 13 +DT_BFLOAT16 = 14 +DT_QINT16 = 15 +DT_QUINT16 = 16 +DT_UINT16 = 17 +DT_COMPLEX128 = 18 +DT_HALF = 19 +DT_RESOURCE = 20 +DT_FLOAT_REF = 101 +DT_DOUBLE_REF = 102 +DT_INT32_REF = 103 +DT_UINT8_REF = 104 +DT_INT16_REF = 105 +DT_INT8_REF = 106 +DT_STRING_REF = 107 +DT_COMPLEX64_REF = 108 +DT_INT64_REF = 109 +DT_BOOL_REF = 110 +DT_QINT8_REF = 111 +DT_QUINT8_REF = 112 +DT_QINT32_REF = 113 +DT_BFLOAT16_REF = 114 +DT_QINT16_REF = 115 +DT_QUINT16_REF = 116 +DT_UINT16_REF = 117 +DT_COMPLEX128_REF = 118 +DT_HALF_REF = 119 +DT_RESOURCE_REF = 120 + + +DESCRIPTOR.enum_types_by_name["DataType"] = _DATATYPE +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\013TypesProtosP\001\370\001\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/src/versions_pb2.py b/tensorboardX/src/versions_pb2.py new file mode 100644 index 0000000..f26386c --- /dev/null +++ b/tensorboardX/src/versions_pb2.py @@ -0,0 +1,118 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: tensorboardX/src/versions.proto + +import sys + +_b = sys.version_info[0] < 3 and (lambda x: x) or (lambda x: x.encode("latin1")) +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pb2 +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="tensorboardX/src/versions.proto", + package="tensorboard", + syntax="proto3", + serialized_pb=_b( + '\n\x1ftensorboardX/src/versions.proto\x12\x0btensorboard"K\n\nVersionDef\x12\x10\n\x08producer\x18\x01 \x01(\x05\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\x05\x12\x15\n\rbad_consumers\x18\x03 \x03(\x05\x42/\n\x18org.tensorflow.frameworkB\x0eVersionsProtosP\x01\xf8\x01\x01\x62\x06proto3' + ), +) + + +_VERSIONDEF = _descriptor.Descriptor( + name="VersionDef", + full_name="tensorboard.VersionDef", + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name="producer", + full_name="tensorboard.VersionDef.producer", + index=0, + number=1, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="min_consumer", + full_name="tensorboard.VersionDef.min_consumer", + index=1, + number=2, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + _descriptor.FieldDescriptor( + name="bad_consumers", + full_name="tensorboard.VersionDef.bad_consumers", + index=2, + number=3, + type=5, + cpp_type=1, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + options=None, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=48, + serialized_end=123, +) + +DESCRIPTOR.message_types_by_name["VersionDef"] = _VERSIONDEF +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +VersionDef = _reflection.GeneratedProtocolMessageType( + "VersionDef", + (_message.Message,), + dict( + DESCRIPTOR=_VERSIONDEF, + __module__="tensorboardX.src.versions_pb2" + # @@protoc_insertion_point(class_scope:tensorboard.VersionDef) + ), +) +_sym_db.RegisterMessage(VersionDef) + + +DESCRIPTOR.has_options = True +DESCRIPTOR._options = _descriptor._ParseOptions( + descriptor_pb2.FileOptions(), _b("\n\030org.tensorflow.frameworkB\016VersionsProtosP\001\370\001\001") +) +# @@protoc_insertion_point(module_scope) diff --git a/tensorboardX/summary.py b/tensorboardX/summary.py new file mode 100644 index 0000000..15eb5b2 --- /dev/null +++ b/tensorboardX/summary.py @@ -0,0 +1,269 @@ +# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""## Generation of summaries. +### Class for writing Summaries +@@FileWriter +@@FileWriterCache +### Summary Ops +@@tensor_summary +@@scalar +@@histogram +@@audio +@@image +@@merge +@@merge_all +## Utilities +@@get_summary_description +""" + +from __future__ import absolute_import, division, print_function + +import bisect +import logging +import re as _re + +import numpy as np +from six import StringIO +from six.moves import range + +from .src.plugin_pr_curve_pb2 import PrCurvePluginData + +# pylint: disable=unused-import +from .src.summary_pb2 import HistogramProto, Summary, SummaryMetadata +from .src.tensor_pb2 import TensorProto +from .src.tensor_shape_pb2 import TensorShapeProto +from .x2num import makenp + +_INVALID_TAG_CHARACTERS = _re.compile(r"[^-/\w\.]") + + +def _clean_tag(name): + # In the past, the first argument to summary ops was a tag, which allowed + # arbitrary characters. Now we are changing the first argument to be the node + # name. This has a number of advantages (users of summary ops now can + # take advantage of the tf name scope system) but risks breaking existing + # usage, because a much smaller set of characters are allowed in node names. + # This function replaces all illegal characters with _s, and logs a warning. + # It also strips leading slashes from the name. + if name is not None: + new_name = _INVALID_TAG_CHARACTERS.sub("_", name) + new_name = new_name.lstrip("/") # Remove leading slashes + if new_name != name: + logging.info("Summary name %s is illegal; using %s instead." % (name, new_name)) + name = new_name + return name + + +def scalar(name, scalar, collections=None): + """Outputs a `Summary` protocol buffer containing a single scalar value. + The generated Summary has a Tensor.proto containing the input Tensor. + Args: + name: A name for the generated node. Will also serve as the series name in + TensorBoard. + tensor: A real numeric Tensor containing a single value. + collections: Optional list of graph collections keys. The new summary op is + added to these collections. Defaults to `[GraphKeys.SUMMARIES]`. + Returns: + A scalar `Tensor` of type `string`. Which contains a `Summary` protobuf. + Raises: + ValueError: If tensor has the wrong shape or type. + """ + name = _clean_tag(name) + scalar = makenp(scalar) + assert scalar.squeeze().ndim == 0, "scalar should be 0D" + scalar = float(scalar) + return Summary(value=[Summary.Value(tag=name, simple_value=scalar)]) + + +def histogram(name, values, bins, collections=None): + # pylint: disable=line-too-long + """Outputs a `Summary` protocol buffer with a histogram. + The generated + [`Summary`](https://www.tensorflow.org/code/tensorflow/core/framework/summary.proto) + has one summary value containing a histogram for `values`. + This op reports an `InvalidArgument` error if any value is not finite. + Args: + name: A name for the generated node. Will also serve as a series name in + TensorBoard. + values: A real numeric `Tensor`. Any shape. Values to use to + build the histogram. + collections: Optional list of graph collections keys. The new summary op is + added to these collections. Defaults to `[GraphKeys.SUMMARIES]`. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + name = _clean_tag(name) + values = makenp(values) + hist = make_histogram(values.astype(float), bins) + return Summary(value=[Summary.Value(tag=name, histo=hist)]) + + +def make_histogram(values, bins): + """Convert values into a histogram proto using logic from histogram.cc.""" + values = values.reshape(-1) + counts, limits = np.histogram(values, bins=bins) + limits = limits[1:] + + sum_sq = values.dot(values) + return HistogramProto( + min=values.min(), + max=values.max(), + num=len(values), + sum=values.sum(), + sum_squares=sum_sq, + bucket_limit=limits, + bucket=counts, + ) + + +def image(tag, tensor): + """Outputs a `Summary` protocol buffer with images. + The summary has up to `max_images` summary values containing images. The + images are built from `tensor` which must be 3-D with shape `[height, width, + channels]` and where `channels` can be: + * 1: `tensor` is interpreted as Grayscale. + * 3: `tensor` is interpreted as RGB. + * 4: `tensor` is interpreted as RGBA. + The `name` in the outputted Summary.Value protobufs is generated based on the + name, with a suffix depending on the max_outputs setting: + * If `max_outputs` is 1, the summary value tag is '*name*/image'. + * If `max_outputs` is greater than 1, the summary value tags are + generated sequentially as '*name*/image/0', '*name*/image/1', etc. + Args: + tag: A name for the generated node. Will also serve as a series name in + TensorBoard. + tensor: A 3-D `uint8` or `float32` `Tensor` of shape `[height, width, + channels]` where `channels` is 1, 3, or 4. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + tag = _clean_tag(tag) + tensor = makenp(tensor, "IMG") + tensor = tensor.astype(np.float32) + tensor = (tensor * 255).astype(np.uint8) + image = make_image(tensor) + return Summary(value=[Summary.Value(tag=tag, image=image)]) + + +def make_image(tensor): + """Convert an numpy representation image to Image protobuf""" + from PIL import Image + + height, width, channel = tensor.shape + image = Image.fromarray(tensor) + import io + + output = io.BytesIO() + image.save(output, format="PNG") + image_string = output.getvalue() + output.close() + return Summary.Image(height=height, width=width, colorspace=channel, encoded_image_string=image_string) + + +def audio(tag, tensor, sample_rate=44100): + tensor = makenp(tensor) + tensor = tensor.squeeze() + if abs(tensor).max() > 1: + print("warning: audio amplitude out of range, auto clipped.") + tensor = tensor.clip(-1, 1) + assert tensor.ndim == 1, "input tensor should be 1 dimensional." + + tensor_list = [int(32767.0 * x) for x in tensor] + import io + import struct + import wave + + fio = io.BytesIO() + Wave_write = wave.open(fio, "wb") + Wave_write.setnchannels(1) + Wave_write.setsampwidth(2) + Wave_write.setframerate(sample_rate) + tensor_enc = b"" + for v in tensor_list: + tensor_enc += struct.pack(" 127: # wierd, value > 127 breaks protobuf + num_thresholds = 127 + data = compute_curve(labels, predictions, num_thresholds=num_thresholds, weights=weights) + pr_curve_plugin_data = PrCurvePluginData(version=0, num_thresholds=num_thresholds).SerializeToString() + PluginData = [SummaryMetadata.PluginData(plugin_name="pr_curves", content=pr_curve_plugin_data)] + smd = SummaryMetadata(plugin_data=PluginData) + tensor = TensorProto( + dtype="DT_FLOAT", + float_val=data.reshape(-1).tolist(), + tensor_shape=TensorShapeProto( + dim=[TensorShapeProto.Dim(size=data.shape[0]), TensorShapeProto.Dim(size=data.shape[1])] + ), + ) + return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)]) + + +# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/summary.py +def compute_curve(labels, predictions, num_thresholds=None, weights=None): + _MINIMUM_COUNT = 1e-7 + + if weights is None: + weights = 1.0 + + # Compute bins of true positives and false positives. + bucket_indices = np.int32(np.floor(predictions * (num_thresholds - 1))) + float_labels = labels.astype(np.float) + histogram_range = (0, num_thresholds - 1) + tp_buckets, _ = np.histogram( + bucket_indices, bins=num_thresholds, range=histogram_range, weights=float_labels * weights + ) + fp_buckets, _ = np.histogram( + bucket_indices, bins=num_thresholds, range=histogram_range, weights=(1.0 - float_labels) * weights + ) + + # Obtain the reverse cumulative sum. + tp = np.cumsum(tp_buckets[::-1])[::-1] + fp = np.cumsum(fp_buckets[::-1])[::-1] + tn = fp[0] - fp + fn = tp[0] - tp + precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp) + recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn) + return np.stack((tp, fp, tn, fn, precision, recall)) diff --git a/tensorboardX/writer.py b/tensorboardX/writer.py new file mode 100644 index 0000000..97ca7c9 --- /dev/null +++ b/tensorboardX/writer.py @@ -0,0 +1,482 @@ +# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Provides an API for generating Event protocol buffers.""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import time + +from .embedding import append_pbtxt, make_mat, make_sprite, make_tsv +from .event_file_writer import EventFileWriter +from .graph import graph +from .graph_onnx import gg +from .src import event_pb2, graph_pb2, summary_pb2 +from .summary import audio, histogram, image, pr_curve, scalar, text + + +class SummaryToEventTransformer(object): + """Abstractly implements the SummaryWriter API. + This API basically implements a number of endpoints (add_summary, + add_session_log, etc). The endpoints all generate an event protobuf, which is + passed to the contained event_writer. + @@__init__ + @@add_summary + @@add_session_log + @@add_graph + @@add_meta_graph + @@add_run_metadata + """ + + def __init__(self, event_writer, graph=None, graph_def=None): + """Creates a `SummaryWriter` and an event file. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers constructed when you + call one of the following functions: `add_summary()`, `add_session_log()`, + `add_event()`, or `add_graph()`. + If you pass a `Graph` to the constructor it is added to + the event file. (This is equivalent to calling `add_graph()` later). + TensorBoard will pick the graph from the file and display it graphically so + you can interactively explore the graph you built. You will usually pass + the graph from the session in which you launched it: + ```python + ...create a graph... + # Launch the graph in a session. + sess = tf.Session() + # Create a summary writer, add the 'graph' to the event file. + writer = tf.summary.FileWriter(, sess.graph) + ``` + Args: + event_writer: An EventWriter. Implements add_event method. + graph: A `Graph` object, such as `sess.graph`. + graph_def: DEPRECATED: Use the `graph` argument instead. + """ + self.event_writer = event_writer + # For storing used tags for session.run() outputs. + self._session_run_tags = {} + # TODO(zihaolucky). pass this an empty graph to check whether it's necessary. + # currently we don't support graph in MXNet using tensorboard. + + def add_summary(self, summary, global_step=None): + """Adds a `Summary` protocol buffer to the event file. + This method wraps the provided summary in an `Event` protocol buffer + and adds it to the event file. + You can pass the result of evaluating any summary op, using + [`Session.run()`](client.md#Session.run) or + [`Tensor.eval()`](framework.md#Tensor.eval), to this + function. Alternatively, you can pass a `tf.Summary` protocol + buffer that you populate with your own data. The latter is + commonly done to report evaluation results in event files. + Args: + summary: A `Summary` protocol buffer, optionally serialized as a string. + global_step: Number. Optional global step value to record with the + summary. + """ + if isinstance(summary, bytes): + summ = summary_pb2.Summary() + summ.ParseFromString(summary) + summary = summ + event = event_pb2.Event(summary=summary) + self._add_event(event, global_step) + + def add_graph_onnx(self, graph): + """Adds a `Graph` protocol buffer to the event file. + """ + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self._add_event(event, None) + + def add_graph(self, graph): + """Adds a `Graph` protocol buffer to the event file. + """ + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self._add_event(event, None) + + def add_session_log(self, session_log, global_step=None): + """Adds a `SessionLog` protocol buffer to the event file. + This method wraps the provided session in an `Event` protocol buffer + and adds it to the event file. + Args: + session_log: A `SessionLog` protocol buffer. + global_step: Number. Optional global step value to record with the + summary. + """ + event = event_pb2.Event(session_log=session_log) + self._add_event(event, global_step) + + def _add_event(self, event, step): + event.wall_time = time.time() + if step is not None: + event.step = int(step) + self.event_writer.add_event(event) + + +class FileWriter(SummaryToEventTransformer): + """Writes `Summary` protocol buffers to event files. + The `FileWriter` class provides a mechanism to create an event file in a + given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + @@__init__ + @@add_summary + @@add_session_log + @@add_event + @@add_graph + @@add_run_metadata + @@get_logdir + @@flush + @@close + """ + + def __init__(self, logdir, graph=None, max_queue=10, flush_secs=120, graph_def=None): + """Creates a `FileWriter` and an event file. + On construction the summary writer creates a new event file in `logdir`. + This event file will contain `Event` protocol buffers constructed when you + call one of the following functions: `add_summary()`, `add_session_log()`, + `add_event()`, or `add_graph()`. + If you pass a `Graph` to the constructor it is added to + the event file. (This is equivalent to calling `add_graph()` later). + TensorBoard will pick the graph from the file and display it graphically so + you can interactively explore the graph you built. You will usually pass + the graph from the session in which you launched it: + ```python + ...create a graph... + # Launch the graph in a session. + sess = tf.Session() + # Create a summary writer, add the 'graph' to the event file. + writer = tf.summary.FileWriter(, sess.graph) + ``` + The other arguments to the constructor control the asynchronous writes to + the event file: + * `flush_secs`: How often, in seconds, to flush the added summaries + and events to disk. + * `max_queue`: Maximum number of summaries or events pending to be + written to disk before one of the 'add' calls block. + Args: + logdir: A string. Directory where event file will be written. + graph: A `Graph` object, such as `sess.graph`. + max_queue: Integer. Size of the queue for pending events and summaries. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. + graph_def: DEPRECATED: Use the `graph` argument instead. + """ + event_writer = EventFileWriter(logdir, max_queue, flush_secs) + super(FileWriter, self).__init__(event_writer, graph, graph_def) + + def get_logdir(self): + """Returns the directory where event file will be written.""" + return self.event_writer.get_logdir() + + def add_event(self, event): + """Adds an event to the event file. + Args: + event: An `Event` protocol buffer. + """ + self.event_writer.add_event(event) + + def flush(self): + """Flushes the event file to disk. + Call this method to make sure that all pending events have been written to + disk. + """ + self.event_writer.flush() + + def close(self): + """Flushes the event file to disk and close the file. + Call this method when you do not need the summary writer anymore. + """ + self.event_writer.close() + + def reopen(self): + """Reopens the EventFileWriter. + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + Does nothing if the EventFileWriter was not closed. + """ + self.event_writer.reopen() + + +class SummaryWriter(object): + """Writes `Summary` directly to event files. + The `SummaryWriter` class provides a high-level api to create an event file in a + given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + """ + + def __init__(self, log_dir=None, comment=""): + """ + Args: + log_dir (string): save location, default is: runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each + run. Use hierarchical folder structure to compare between runs easily. e.g. 'runs/exp1', 'runs/exp2' + comment (string): comment that appends to the default log_dir + """ + if not log_dir: + import socket + from datetime import datetime + + current_time = datetime.now().strftime("%b%d_%H-%M-%S") + log_dir = os.path.join("runs", current_time + "_" + socket.gethostname() + comment) + self.file_writer = FileWriter(logdir=log_dir) + v = 1e-12 + buckets = [] + neg_buckets = [] + while v < 1e20: + buckets.append(v) + neg_buckets.append(-v) + v *= 1.1 + self.default_bins = neg_buckets[::-1] + [0] + buckets + self.text_tags = [] + # + self.all_writers = {self.file_writer.get_logdir(): self.file_writer} + self.scalar_dict = {} # {writer_id : [[timestamp, step, value],...],...} + + def __append_to_scalar_dict(self, tag, scalar_value, global_step, timestamp): + """This adds an entry to the self.scalar_dict datastructure with format + {writer_id : [[timestamp, step, value], ...], ...}. + """ + from .x2num import makenp + + if tag not in self.scalar_dict.keys(): + self.scalar_dict[tag] = [] + self.scalar_dict[tag].append([timestamp, global_step, float(makenp(scalar_value))]) + + def add_scalar(self, tag, scalar_value, global_step=None): + """Add scalar data to summary. + + Args: + tag (string): Data identifier + scalar_value (float): Value to save + global_step (int): Global step value to record + """ + self.file_writer.add_summary(scalar(tag, scalar_value), global_step) + self.__append_to_scalar_dict(tag, scalar_value, global_step, time.time()) + + def add_scalars(self, main_tag, tag_scalar_dict, global_step=None): + """Adds many scalar data to summary. + + Args: + tag (string): Data identifier + main_tag (string): The parent name for the tags + tag_scalar_dict (dict): Key-value pair storing the tag and corresponding values + global_step (int): Global step value to record + + Examples:: + + writer.add_scalars('run_14h',{'xsinx':i*np.sin(i/r), + 'xcosx':i*np.cos(i/r), + 'arctanx': numsteps*np.arctan(i/r)}, i) + # This function adds three values to the same scalar plot with the tag + # 'run_14h' in TensorBoard's scalar section. + """ + timestamp = time.time() + fw_logdir = self.file_writer.get_logdir() + for tag, scalar_value in tag_scalar_dict.items(): + fw_tag = fw_logdir + "/" + main_tag + "/" + tag + if fw_tag in self.all_writers.keys(): + fw = self.all_writers[fw_tag] + else: + fw = FileWriter(logdir=fw_tag) + self.all_writers[fw_tag] = fw + fw.add_summary(scalar(main_tag, scalar_value), global_step) + self.__append_to_scalar_dict(fw_tag, scalar_value, global_step, timestamp) + + def export_scalars_to_json(self, path): + """Exports to the given path an ASCII file containing all the scalars written + so far by this instance, with the following format: + {writer_id : [[timestamp, step, value], ...], ...} + """ + with open(path, "w") as f: + json.dump(self.scalar_dict, f) + + def add_histogram(self, tag, values, global_step=None, bins="tensorflow"): + """Add histogram to summary. + + Args: + tag (string): Data identifier + values (numpy.array): Values to build histogram + global_step (int): Global step value to record + bins (string): one of {'tensorflow','auto', 'fd', ...}, this determines how the bins are made. You can find + other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html + """ + if bins == "tensorflow": + bins = self.default_bins + self.file_writer.add_summary(histogram(tag, values, bins), global_step) + + def add_image(self, tag, img_tensor, global_step=None): + """Add image data to summary. + + Note that this requires the ``pillow`` package. + + Args: + tag (string): Data identifier + img_tensor (torch.Tensor): Image data + global_step (int): Global step value to record + Shape: + img_tensor: :math:`(3, H, W)`. Use ``torchvision.utils.make_grid()`` to prepare it is a good idea. + """ + self.file_writer.add_summary(image(tag, img_tensor), global_step) + + def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100): + """Add audio data to summary. + + Args: + tag (string): Data identifier + snd_tensor (torch.Tensor): Sound data + global_step (int): Global step value to record + sample_rate (int): sample rate in Hz + + Shape: + snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1]. + """ + self.file_writer.add_summary(audio(tag, snd_tensor, sample_rate=sample_rate), global_step) + + def add_text(self, tag, text_string, global_step=None): + """Add text data to summary. + + Args: + tag (string): Data identifier + text_string (string): String to save + global_step (int): Global step value to record + + Examples:: + + writer.add_text('lstm', 'This is an lstm', 0) + writer.add_text('rnn', 'This is an rnn', 10) + """ + self.file_writer.add_summary(text(tag, text_string), global_step) + if tag not in self.text_tags: + self.text_tags.append(tag) + extension_dir = self.file_writer.get_logdir() + "/plugins/tensorboard_text/" + if not os.path.exists(extension_dir): + os.makedirs(extension_dir) + with open(extension_dir + "tensors.json", "w") as fp: + json.dump(self.text_tags, fp) + + def add_graph_onnx(self, prototxt): + self.file_writer.add_graph_onnx(gg(prototxt)) + + def add_graph(self, model, input_to_model, verbose=False): + # prohibit second call? + # no, let tensorboard handles it and show its warning message. + """Add graph data to summary. + + Args: + model (torch.nn.Module): model to draw. + input_to_model (torch.autograd.Variable): a variable or a tuple of variables to be fed. + + """ + from distutils.version import LooseVersion + + import torch + + if LooseVersion(torch.__version__) >= LooseVersion("0.3.1"): + pass + else: + if LooseVersion(torch.__version__) >= LooseVersion("0.3.0"): + print("You are using PyTorch==0.3.0, use add_graph_onnx()") + return + if not hasattr(torch.autograd.Variable, "grad_fn"): + print("add_graph() only supports PyTorch v0.2.") + return + self.file_writer.add_graph(graph(model, input_to_model, verbose)) + + def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, tag="default"): + """Add embedding projector data to summary. + + Args: + mat (torch.Tensor): A matrix which each row is the feature vector of the data point + metadata (list): A list of labels, each element will be convert to string + label_img (torch.Tensor): Images correspond to each data point + global_step (int): Global step value to record + tag (string): Name for the embedding + Shape: + mat: :math:`(N, D)`, where N is number of data and D is feature dimension + + label_img: :math:`(N, C, H, W)` + + Examples:: + + import keyword + import torch + meta = [] + while len(meta)<100: + meta = meta+keyword.kwlist # get some strings + meta = meta[:100] + + for i, v in enumerate(meta): + meta[i] = v+str(i) + + label_img = torch.rand(100, 3, 10, 32) + for i in range(100): + label_img[i]*=i/100.0 + + writer.add_embedding(torch.randn(100, 5), metadata=meta, label_img=label_img) + writer.add_embedding(torch.randn(100, 5), label_img=label_img) + writer.add_embedding(torch.randn(100, 5), metadata=meta) + """ + if global_step is None: + global_step = 0 + # clear pbtxt? + save_path = os.path.join(self.file_writer.get_logdir(), str(global_step).zfill(5)) + try: + os.makedirs(save_path) + except OSError: + print("warning: Embedding dir exists, did you set global_step for add_embedding()?") + if metadata is not None: + assert mat.size(0) == len(metadata), "#labels should equal with #data points" + make_tsv(metadata, save_path) + if label_img is not None: + assert mat.size(0) == label_img.size(0), "#images should equal with #data points" + make_sprite(label_img, save_path) + assert mat.dim() == 2, "mat should be 2D, where mat.size(0) is the number of data points" + make_mat(mat.tolist(), save_path) + # new funcion to append to the config file a new embedding + append_pbtxt(metadata, label_img, self.file_writer.get_logdir(), str(global_step).zfill(5), tag) + + def add_pr_curve(self, tag, labels, predictions, global_step=None, num_thresholds=127, weights=None): + """Adds precision recall curve. + + Args: + tag (string): Data identifier + labels (torch.Tensor): Ground truth data. Binary label for each element. + predictions (torch.Tensor): The probability that an element be classified as true. Value should in [0, 1] + global_step (int): Global step value to record + num_thresholds (int): Number of thresholds used to draw the curve. + + """ + from .x2num import makenp + + labels = makenp(labels) + predictions = makenp(predictions) + self.file_writer.add_summary(pr_curve(tag, labels, predictions, num_thresholds, weights), global_step) + + def close(self): + if self.file_writer is None: + return # ignore double close + self.file_writer.flush() + self.file_writer.close() + for path, writer in self.all_writers.items(): + writer.flush() + writer.close() + self.file_writer = self.all_writers = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() diff --git a/tensorboardX/x2num.py b/tensorboardX/x2num.py new file mode 100644 index 0000000..d81c5c7 --- /dev/null +++ b/tensorboardX/x2num.py @@ -0,0 +1,93 @@ +# DO NOT alter/distruct/free input object ! + +import numpy as np + + +def makenp(x, modality=None): + # if already numpy, return + if isinstance(x, np.ndarray): + if modality == "IMG" and x.dtype == np.uint8: + return x.astype(np.float32) / 255.0 + return x + if np.isscalar(x): + return np.array([x]) + if "torch" in str(type(x)): + return pytorch_np(x, modality) + if "chainer" in str(type(x)): + return chainer_np(x, modality) + if "mxnet" in str(type(x)): + return mxnet_np(x, modality) + + +def pytorch_np(x, modality): + import torch + + if isinstance(x, torch.autograd.Variable): + x = x.data + x = x.cpu().numpy() + if modality == "IMG": + x = _prepare_image(x) + return x + + +def theano_np(x): + import theano + + pass + + +def caffe2_np(x): + pass + + +def mxnet_np(x, modality): + x = x.asnumpy() + if modality == "IMG": + x = _prepare_image(x) + return x + + +def chainer_np(x, modality): + import chainer + + x = chainer.cuda.to_cpu(x.data) + if modality == "IMG": + x = _prepare_image(x) + return x + + +def make_grid(I, ncols=8): + assert isinstance(I, np.ndarray), "plugin error, should pass numpy array here" + assert I.ndim == 4 and I.shape[1] == 3 + nimg = I.shape[0] + H = I.shape[2] + W = I.shape[3] + ncols = min(nimg, ncols) + nrows = int(np.ceil(float(nimg) / ncols)) + canvas = np.zeros((3, H * nrows, W * ncols)) + i = 0 + for y in range(nrows): + for x in range(ncols): + if i >= nimg: + break + canvas[:, y * H : (y + 1) * H, x * W : (x + 1) * W] = I[i] + i = i + 1 + return canvas + + +def _prepare_image(I): + assert isinstance(I, np.ndarray), "plugin error, should pass numpy array here" + assert I.ndim == 2 or I.ndim == 3 or I.ndim == 4 + if I.ndim == 4: # NCHW + if I.shape[1] == 1: # N1HW + I = np.concatenate((I, I, I), 1) # N3HW + assert I.shape[1] == 3 + I = make_grid(I) # 3xHxW + if I.ndim == 3 and I.shape[0] == 1: # 1xHxW + I = np.concatenate((I, I, I), 0) # 3xHxW + if I.ndim == 2: # HxW + I = np.expand_dims(I, 0) # 1xHxW + I = np.concatenate((I, I, I), 0) # 3xHxW + I = I.transpose(1, 2, 0) + + return I diff --git a/test.py b/test.py new file mode 100644 index 0000000..4fcf8bf --- /dev/null +++ b/test.py @@ -0,0 +1,202 @@ +from __future__ import print_function + +import argparse +import glob +import os +import time + +import cv2 +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torchvision.transforms as transform_lib +from PIL import Image + +import lib.TestTransforms as transforms +from models.ColorVidNet import ColorVidNet +from models.FrameColor import frame_colorization +from models.NonlocalNet import VGG19_feature_color_torchversion, WarpNet +from utils.tensor_lab2rgb import tensor_lab2rgb +from utils.util import batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not, save_frames, uncenter_l +from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor + +os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +torch.cuda.set_device(0) + + +def colorize_video(opt, input_path, reference_path, output_path, nonlocal_net, colornet, vggnet): + # parameters for wls filter + wls_filter_on = True + lambda_value = 500 + sigma_color = 4 + + # processing folders + mkdir_if_not(output_path) + files = glob.glob(output_path + "*") + print("processing the folder:", input_path) + path, dirs, filenames = os.walk(input_path).__next__() + file_count = len(filenames) + filenames.sort() + + # NOTE: resize frames to 216*384 + transform = transforms.Compose( + [CenterPad(opt.image_size), transform_lib.CenterCrop(opt.image_size), RGB2Lab(), ToTensor(), Normalize()] + ) + + # if frame propagation: use the first frame as reference + # otherwise, use the specified reference image + ref_name = input_path + filenames[0] if opt.frame_propagate else reference_path + print("reference name:", ref_name) + frame_ref = Image.open(ref_name) + + iter_num = 0 + total_time = 0 + I_last_lab_predict = None + + for index in range(file_count): + frame_name = filenames[index] + print("input =", frame_name) + print("reference =", ref_name) + frame1 = Image.open(os.path.join(input_path, frame_name)) + + IA_lab_large = transform(frame1).unsqueeze(0).cuda() + IB_lab_large = transform(frame_ref).unsqueeze(0).cuda() + IA_lab = torch.nn.functional.interpolate(IA_lab_large, scale_factor=0.5, mode="bilinear") + IB_lab = torch.nn.functional.interpolate(IB_lab_large, scale_factor=0.5, mode="bilinear") + + IA_l = IA_lab[:, 0:1, :, :] + IA_ab = IA_lab[:, 1:3, :, :] + IB_l = IB_lab[:, 0:1, :, :] + IB_ab = IB_lab[:, 1:3, :, :] + if I_last_lab_predict is None: + if opt.frame_propagate: + I_last_lab_predict = IB_lab + else: + I_last_lab_predict = torch.zeros_like(IA_lab).cuda() + + # start the frame colorization + with torch.no_grad(): + I_current_lab = IA_lab + I_reference_lab = IB_lab + I_reference_l = I_reference_lab[:, 0:1, :, :] + I_reference_ab = I_reference_lab[:, 1:3, :, :] + I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) + + t_start = time.clock() + torch.cuda.synchronize() + + features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) + I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( + I_current_lab, + I_reference_lab, + I_last_lab_predict, + features_B, + vggnet, + nonlocal_net, + colornet, + feature_noise=0, + temperature=1e-10, + ) + I_last_lab_predict = torch.cat((IA_l, I_current_ab_predict), dim=1) + + # update timing + torch.cuda.synchronize() + delta_t = time.clock() - t_start + print("runtime:", delta_t) + if iter_num > 0: + total_time += delta_t + print("%.2g second" % (total_time / iter_num)) + + # upsampling + curr_bs_l = IA_lab_large[:, 0:1, :, :] + curr_predict = ( + torch.nn.functional.interpolate(I_current_ab_predict.data.cpu(), scale_factor=2, mode="bilinear") * 1.25 + ) + + # filtering + if wls_filter_on: + guide_image = uncenter_l(curr_bs_l) * 255 / 100 + wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter( + guide_image[0, 0, :, :].cpu().numpy().astype(np.uint8), lambda_value, sigma_color + ) + curr_predict_a = wls_filter.filter(curr_predict[0, 0, :, :].cpu().numpy()) + curr_predict_b = wls_filter.filter(curr_predict[0, 1, :, :].cpu().numpy()) + curr_predict_a = torch.from_numpy(curr_predict_a).unsqueeze(0).unsqueeze(0) + curr_predict_b = torch.from_numpy(curr_predict_b).unsqueeze(0).unsqueeze(0) + curr_predict_filter = torch.cat((curr_predict_a, curr_predict_b), dim=1) + IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict_filter[:32, ...]) + else: + IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict[:32, ...]) + + # save the frames + save_frames(IA_predict_rgb, output_path, index) + iter_num = iter_num + 1 + + # output video + video_name = "video.avi" + folder2vid(image_folder=output_path, output_dir=".", filename=video_name) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--frame_propagate", default=False, type=bool) + parser.add_argument("--image_size", type=int, default=[216 * 2, 384 * 2]) + parser.add_argument("--cuda", action="store_false") + parser.add_argument("--gpu_ids", type=str, default="0", help="separate by comma") + parser.add_argument("--clips_path", type=str, default="legacy_videos/legacy_clips", help="path of input clips") + parser.add_argument( + "--ref_path", type=str, default="legacy_videos/legacy_clips_ref", help="path of refernce images" + ) + parser.add_argument( + "--output_path", type=str, default="legacy_videos/legacy_clips_output", help="path of output clips" + ) + opt = parser.parse_args(args=[]) + opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")] + cudnn.benchmark = True + print("running on GPU", opt.gpu_ids) + + clips = os.listdir(opt.clips_path) + clips.sort() + clips.reverse() + + refs = os.listdir(opt.ref_path) + refs.sort() + + nonlocal_net = WarpNet(1) + colornet = ColorVidNet(7) + vggnet = VGG19_feature_color_torchversion() + vggnet.load_state_dict(torch.load("data/vgg19_conv.pth")) + for param in vggnet.parameters(): + param.requires_grad = False + + nonlocal_test_path = os.path.join("checkpoints/", "video_moredata_l1/nonlocal_net_iter_76000.pth") + color_test_path = os.path.join("checkpoints/", "video_moredata_l1/colornet_iter_76000.pth") + print("succesfully load nonlocal model: ", nonlocal_test_path) + print("succesfully load color model: ", color_test_path) + nonlocal_net.load_state_dict(torch.load(nonlocal_test_path)) + colornet.load_state_dict(torch.load(color_test_path)) + + nonlocal_net.eval() + colornet.eval() + vggnet.eval() + nonlocal_net.cuda() + colornet.cuda() + vggnet.cuda() + + for clip in clips: + clip_name = clip.split("_")[0] + for ref_name in refs: + try: + colorize_video( + opt, + os.path.join(opt.clips_path, clip), + os.path.join(opt.ref_path, ref_name), + os.path.join(opt.output_path, clip_name + "_" + ref_name.split(".")[0]), + nonlocal_net, + colornet, + vggnet, + ) + except Exception as error: + print("error when colorizing the video " + ref_name) + print(error) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/utils/data_augmentation.py b/utils/data_augmentation.py new file mode 100644 index 0000000..1362ed4 --- /dev/null +++ b/utils/data_augmentation.py @@ -0,0 +1,154 @@ +import random + +import cv2 +import lib.functional as F +import numpy as np +from PIL import Image +from scipy.misc import imread, imresize, imsave +from scipy.ndimage.filters import gaussian_filter +from skimage import color + +cv2.setNumThreads(0) + + +# %% Dataloader +class RGB2Lab(object): + def __init__(self): + pass + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs) + return image_lab + + +class Guassian_noise(object): + """Elastic distortion + """ + + def __init__(self, noise_sigma=0.1): + self.noise_sigma = noise_sigma + + def __call__(self, inputs): + h = inputs.shape[0] + w = inputs.shape[1] + noisy_image = inputs + noise = np.random.randn(h, w) * self.noise_sigma + noisy_image[:, :, 0] = inputs[:, :, 0] + noise + + return noisy_image + + +class Distortion(object): + """Elastic distortion + """ + + def __init__(self, distortion_level=3, flip_probability=0.1): + self.alpha_max = distortion_level + self.flip_probability = flip_probability + + def __call__(self, inputs): + if np.random.rand() < self.flip_probability: + inputs = inputs.transpose(Image.FLIP_LEFT_RIGHT) + + inputs = np.array(inputs) + alpha = np.random.rand() * self.alpha_max + sigma = 50 + random_state = np.random.RandomState(None) + shape = inputs.shape[0], inputs.shape[1] + + dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000 + dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000 + + x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij") + remap_image = cv2.remap( + inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR + ) + + return Image.fromarray(remap_image) + + +class Normalize(object): + def __init__(self): + pass + + def __call__(self, inputs): + inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1) + inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1)) + return inputs + + +class ToTensor(object): + def __init__(self): + pass + + def __call__(self, inputs): + outputs = F.to_mytensor(inputs) + return outputs + + +class RandomErasing(object): + """ + Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al. + ------------------------------------------------------------------------------------- + probability: The probability that the operation will be performed. + sl: min erasing area + sh: max erasing area + r1: min aspect ratio + mean: erasing value + ------------------------------------------------------------------------------------- + """ + + def __init__(self, probability=0.6, sl=0.05, sh=0.6): + self.probability = probability + self.sl = sl + self.sh = sh + + def __call__(self, img): + img = np.array(img) + if random.uniform(0, 1) > self.probability: + return Image.fromarray(img) + + area = img.shape[0] * img.shape[1] + h0 = img.shape[0] + w0 = img.shape[1] + channel = img.shape[2] + + h = int(round(random.uniform(self.sl, self.sh) * h0)) + w = int(round(random.uniform(self.sl, self.sh) * w0)) + + if w < img.shape[1] and h < img.shape[0]: + x1 = random.randint(0, img.shape[0] - h) + y1 = random.randint(0, img.shape[1] - w) + img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255 + return Image.fromarray(img) + + return Image.fromarray(img) + + +class CenteredPad(object): + """ + pad the frame with black border, + make square image for processing + """ + + def __init__(self, image_size): + self.image_size = image_size + + def __call__(self, img): + img = np.array(img) + width = np.size(img, 1) + height = np.size(img, 0) + old_size = [height, width] + + ratio = float(self.image_size) / max(height, width) + new_size = [int(x * ratio) for x in old_size] + I_resize = imresize(img, new_size) + width_new = np.size(I_resize, 1) + height_new = np.size(I_resize, 0) + + I_pad = np.zeros((self.image_size, self.image_size, 3)) + start_height = (self.image_size - new_size[0]) // 2 + start_width = (self.image_size - new_size[1]) // 2 + I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize + + return Image.fromarray(I_pad.astype(np.uint8)) diff --git a/utils/flowlib.py b/utils/flowlib.py new file mode 100644 index 0000000..f982829 --- /dev/null +++ b/utils/flowlib.py @@ -0,0 +1,522 @@ +#!/usr/bin/python +""" +# ============================== +# flowlib.py +# library for optical flow processing +# Author: Ruoteng Li +# Date: 6th Aug 2016 +# ============================== +""" +import matplotlib.colors as cl +import matplotlib.pyplot as plt +import numpy as np +import png +from PIL import Image + +UNKNOWN_FLOW_THRESH = 1e7 +SMALLFLOW = 0.0 +LARGEFLOW = 1e8 + +""" +============= +Flow Section +============= +""" + + +def show_flow(filename): + """ + visualize optical flow map using matplotlib + :param filename: optical flow file + :return: None + """ + flow = read_flow(filename) + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + + +def visualize_flow(flow, mode="Y"): + """ + this function visualize the input flow + :param flow: input flow in array + :param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color) + :return: None + """ + if mode == "Y": + # Ccbcr color wheel + img = flow_to_image(flow) + plt.imshow(img) + plt.show() + elif mode == "RGB": + (h, w) = flow.shape[0:2] + du = flow[:, :, 0] + dv = flow[:, :, 1] + valid = flow[:, :, 2] + max_flow = max(np.max(du), np.max(dv)) + img = np.zeros((h, w, 3), dtype=np.float64) + # angle layer + img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi) + # magnitude layer, normalized to 1 + img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow + # phase layer + img[:, :, 2] = 8 - img[:, :, 1] + # clip to [0,1] + small_idx = img[:, :, 0:3] < 0 + large_idx = img[:, :, 0:3] > 1 + img[small_idx] = 0 + img[large_idx] = 1 + # convert to rgb + img = cl.hsv_to_rgb(img) + # remove invalid point + img[:, :, 0] = img[:, :, 0] * valid + img[:, :, 1] = img[:, :, 1] * valid + img[:, :, 2] = img[:, :, 2] * valid + # show + plt.imshow(img) + plt.show() + + return None + + +def read_flow(filename): + """ + read optical flow from Middlebury .flo file + :param filename: name of the flow file + :return: optical flow data in matrix + """ + f = open(filename, "rb") + try: + magic = np.fromfile(f, np.float32, count=1)[0] # For Python3.x + except: + magic = np.fromfile(f, np.float32, count=1) # For Python2.x + data2d = None + + if 202021.25 != magic: + print("Magic number incorrect. Invalid .flo file") + else: + w = np.fromfile(f, np.int32, count=1)[0] + h = np.fromfile(f, np.int32, count=1)[0] + # print("Reading %d x %d flo file" % (h, w)) + data2d = np.fromfile(f, np.float32, count=2 * w * h) + # reshape data into 3D array (columns, rows, channels) + data2d = np.resize(data2d, (h, w, 2)) + f.close() + return data2d + + +def read_flow_png(flow_file): + """ + Read optical flow from KITTI .png file + :param flow_file: name of the flow file + :return: optical flow data in matrix + """ + flow_object = png.Reader(filename=flow_file) + flow_direct = flow_object.asDirect() + flow_data = list(flow_direct[2]) + (w, h) = flow_direct[3]["size"] + flow = np.zeros((h, w, 3), dtype=np.float64) + for i in range(len(flow_data)): + flow[i, :, 0] = flow_data[i][0::3] + flow[i, :, 1] = flow_data[i][1::3] + flow[i, :, 2] = flow_data[i][2::3] + + invalid_idx = flow[:, :, 2] == 0 + flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0 + flow[invalid_idx, 0] = 0 + flow[invalid_idx, 1] = 0 + return flow + + +def write_flow(flow, filename): + """ + write optical flow in Middlebury .flo format + :param flow: optical flow map + :param filename: optical flow file path to be saved + :return: None + """ + f = open(filename, "wb") + magic = np.array([202021.25], dtype=np.float32) + (height, width) = flow.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + flow.tofile(f) + f.close() + + +def segment_flow(flow): + h = flow.shape[0] + w = flow.shape[1] + u = flow[:, :, 0] + v = flow[:, :, 1] + + idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW) + idx2 = abs(u) == SMALLFLOW + class0 = (v == 0) & (u == 0) + u[idx2] = 0.00001 + tan_value = v / u + + class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0) + class2 = (tan_value >= 1) & (u >= 0) & (v >= 0) + class3 = (tan_value < -1) & (u <= 0) & (v >= 0) + class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0) + class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0) + class7 = (tan_value < -1) & (u >= 0) & (v <= 0) + class6 = (tan_value >= 1) & (u <= 0) & (v <= 0) + class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0) + + seg = np.zeros((h, w)) + + seg[class1] = 1 + seg[class2] = 2 + seg[class3] = 3 + seg[class4] = 4 + seg[class5] = 5 + seg[class6] = 6 + seg[class7] = 7 + seg[class8] = 8 + seg[class0] = 0 + seg[idx] = 0 + + return seg + + +def flow_error(tu, tv, u, v): + """ + Calculate average end point error + :param tu: ground-truth horizontal flow map + :param tv: ground-truth vertical flow map + :param u: estimated horizontal flow map + :param v: estimated vertical flow map + :return: End point error of the estimated flow + """ + smallflow = 0.0 + """ + stu = tu[bord+1:end-bord,bord+1:end-bord] + stv = tv[bord+1:end-bord,bord+1:end-bord] + su = u[bord+1:end-bord,bord+1:end-bord] + sv = v[bord+1:end-bord,bord+1:end-bord] + """ + stu = tu[:] + stv = tv[:] + su = u[:] + sv = v[:] + + idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH) + stu[idxUnknow] = 0 + stv[idxUnknow] = 0 + su[idxUnknow] = 0 + sv[idxUnknow] = 0 + + ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)] + index_su = su[ind2] + index_sv = sv[ind2] + an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1) + un = index_su * an + vn = index_sv * an + + index_stu = stu[ind2] + index_stv = stv[ind2] + tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1) + tun = index_stu * tn + tvn = index_stv * tn + + """ + angle = un * tun + vn * tvn + (an * tn) + index = [angle == 1.0] + angle[index] = 0.999 + ang = np.arccos(angle) + mang = np.mean(ang) + mang = mang * 180 / np.pi + """ + + epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2) + epe = epe[ind2] + mepe = np.mean(epe) + return mepe + + +def flow_to_image(flow, display=False): + """ + Convert flow into middlebury color code image + :param flow: optical flow map + :return: optical flow image in middlebury color + """ + u = flow[:, :, 0] + v = flow[:, :, 1] + + maxu = -999.0 + maxv = -999.0 + minu = 999.0 + minv = 999.0 + + idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH) + u[idxUnknow] = 0 + v[idxUnknow] = 0 + + maxu = max(maxu, np.max(u)) + minu = min(minu, np.min(u)) + + maxv = max(maxv, np.max(v)) + minv = min(minv, np.min(v)) + + rad = np.sqrt(u ** 2 + v ** 2) + maxrad = max(-1, np.max(rad)) + + if display: + print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv)) + + u = u / (maxrad + np.finfo(float).eps) + v = v / (maxrad + np.finfo(float).eps) + + img = compute_color(u, v) + + idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2) + img[idx] = 0 + + return np.uint8(img) + + +def evaluate_flow_file(gt, pred): + """ + evaluate the estimated optical flow end point error according to ground truth provided + :param gt: ground truth file path + :param pred: estimated optical flow file path + :return: end point error, float32 + """ + # Read flow files and calculate the errors + gt_flow = read_flow(gt) # ground truth flow + eva_flow = read_flow(pred) # predicted flow + # Calculate errors + average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1]) + return average_pe + + +def evaluate_flow(gt_flow, pred_flow): + """ + gt: ground-truth flow + pred: estimated flow + """ + average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], pred_flow[:, :, 0], pred_flow[:, :, 1]) + return average_pe + + +""" +============== +Disparity Section +============== +""" + + +def read_disp_png(file_name): + """ + Read optical flow from KITTI .png file + :param file_name: name of the flow file + :return: optical flow data in matrix + """ + image_object = png.Reader(filename=file_name) + image_direct = image_object.asDirect() + image_data = list(image_direct[2]) + (w, h) = image_direct[3]["size"] + channel = len(image_data[0]) / w + flow = np.zeros((h, w, channel), dtype=np.uint16) + for i in range(len(image_data)): + for j in range(channel): + flow[i, :, j] = image_data[i][j::channel] + return flow[:, :, 0] / 256 + + +def disp_to_flowfile(disp, filename): + """ + Read KITTI disparity file in png format + :param disp: disparity matrix + :param filename: the flow file name to save + :return: None + """ + f = open(filename, "wb") + magic = np.array([202021.25], dtype=np.float32) + (height, width) = disp.shape[0:2] + w = np.array([width], dtype=np.int32) + h = np.array([height], dtype=np.int32) + empty_map = np.zeros((height, width), dtype=np.float32) + data = np.dstack((disp, empty_map)) + magic.tofile(f) + w.tofile(f) + h.tofile(f) + data.tofile(f) + f.close() + + +""" +============== +Image Section +============== +""" + + +def read_image(filename): + """ + Read normal image of any format + :param filename: name of the image file + :return: image data in matrix uint8 type + """ + img = Image.open(filename) + im = np.array(img) + return im + + +def warp_image(im, flow): + """ + Use optical flow to warp image to the next + :param im: image to warp + :param flow: optical flow + :return: warped image + """ + from scipy import interpolate + + image_height = im.shape[0] + image_width = im.shape[1] + flow_height = flow.shape[0] + flow_width = flow.shape[1] + n = image_height * image_width + (iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width]) + (fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width]) + fx += flow[:, :, 0].astype(np.float) + fy += flow[:, :, 1] + mask = np.logical_or(fx < 0, fx > flow_width) + mask = np.logical_or(mask, fy < 0) + mask = np.logical_or(mask, fy > flow_height) + fx = np.minimum(np.maximum(fx, 0), flow_width) + fy = np.minimum(np.maximum(fy, 0), flow_height) + points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1) + xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1) + warp = np.zeros((image_height, image_width, im.shape[2])) + for i in range(im.shape[2]): + channel = im[:, :, i] + # plt.imshow(channel, cmap='gray') + values = channel.reshape(n, 1) + new_channel = interpolate.griddata(points, values, xi, method="cubic") + new_channel = np.reshape(new_channel, [flow_height, flow_width]) + new_channel[mask] = 1 + warp[:, :, i] = new_channel.astype(np.uint8) + + return warp.astype(np.uint8) + + +""" +============== +Others +============== +""" + + +def scale_image(image, new_range): + """ + Linearly scale the image into desired range + :param image: input image + :param new_range: the new range to be aligned + :return: image normalized in new range + """ + min_val = np.min(image).astype(np.float32) + max_val = np.max(image).astype(np.float32) + min_val_new = np.array(min(new_range), dtype=np.float32) + max_val_new = np.array(max(new_range), dtype=np.float32) + scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new + return scaled_image.astype(np.uint8) + + +def compute_color(u, v): + """ + compute optical flow color map + :param u: optical flow horizontal map + :param v: optical flow vertical map + :return: optical flow in color code + """ + [h, w] = u.shape + img = np.zeros([h, w, 3]) + nanIdx = np.isnan(u) | np.isnan(v) + u[nanIdx] = 0 + v[nanIdx] = 0 + + colorwheel = make_color_wheel() + ncols = np.size(colorwheel, 0) + + rad = np.sqrt(u ** 2 + v ** 2) + + a = np.arctan2(-v, -u) / np.pi + + fk = (a + 1) / 2 * (ncols - 1) + 1 + + k0 = np.floor(fk).astype(int) + + k1 = k0 + 1 + k1[k1 == ncols + 1] = 1 + f = fk - k0 + + for i in range(0, np.size(colorwheel, 1)): + tmp = colorwheel[:, i] + col0 = tmp[k0 - 1] / 255 + col1 = tmp[k1 - 1] / 255 + col = (1 - f) * col0 + f * col1 + + idx = rad <= 1 + col[idx] = 1 - rad[idx] * (1 - col[idx]) + notidx = np.logical_not(idx) + + col[notidx] *= 0.75 + img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx))) + + return img + + +def make_color_wheel(): + """ + Generate color wheel according Middlebury color code + :return: Color wheel + """ + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + + colorwheel = np.zeros([ncols, 3]) + + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY)) + col += RY + + # YG + colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG)) + colorwheel[col : col + YG, 1] = 255 + col += YG + + # GC + colorwheel[col : col + GC, 1] = 255 + colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC)) + col += GC + + # CB + colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB)) + colorwheel[col : col + CB, 2] = 255 + col += CB + + # BM + colorwheel[col : col + BM, 2] = 255 + colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM)) + col += +BM + + # MR + colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) + colorwheel[col : col + MR, 0] = 255 + + return colorwheel diff --git a/utils/frame2video.py b/utils/frame2video.py new file mode 100644 index 0000000..8c5bb25 --- /dev/null +++ b/utils/frame2video.py @@ -0,0 +1,42 @@ +import argparse +import gc +import glob +import math +import os +import platform +import random +import shutil +import sys +import time +from datetime import datetime +from pdb import set_trace + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +from skimage import color, io + +video_name = "video.mp4" +frame_folder = "../../results/seq1_new/" + +images = [ + img for img in os.listdir(frame_folder) if img.endswith(".jpg") or img.endswith(".bmp") or img.endswith(".png") +] +images.sort() +frame = cv2.imread(os.path.join(frame_folder, images[0])) +height, width, layers = frame.shape +video = cv2.VideoWriter("video.avi", cv2.VideoWriter_fourcc("X", "V", "I", "D"), 15, (width, height)) +# video = cv2.VideoWriter('video.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (width, height)) +for image in images: + video.write(cv2.imread(os.path.join(frame_folder, image))) + +cv2.destroyAllWindows() +video.release() + +import imageio + +frames = [] +for image in images: + frames.append(imageio.imread(os.path.join(frame_folder, image))) +imageio.mimsave("movie.gif", frames) + diff --git a/utils/hist_match.py b/utils/hist_match.py new file mode 100644 index 0000000..c9271ed --- /dev/null +++ b/utils/hist_match.py @@ -0,0 +1,101 @@ +import numpy as np +from matplotlib import pyplot as plt + + +def hist_match_module(source, template): + """ + Adjust the pixel values of a grayscale image such that its histogram + matches that of a target image + + Arguments: + ----------- + source: np.ndarray + Image to transform; the histogram is computed over the flattened + array + template: np.ndarray + Template image; can have different dimensions to source + Returns: + ----------- + matched: np.ndarray + The transformed output image + """ + + oldshape = source.shape + source = source.ravel() + template = template.ravel() + + # get the set of unique pixel values and their corresponding indices and + # counts + s_values, bin_idx, s_counts = np.unique(source, return_inverse=True, return_counts=True) + t_values, t_counts = np.unique(template, return_counts=True) + + # take the cumsum of the counts and normalize by the number of pixels to + # get the empirical cumulative distribution functions for the source and + # template images (maps pixel value --> quantile) + s_quantiles = np.cumsum(s_counts).astype(np.float64) + s_quantiles /= s_quantiles[-1] + t_quantiles = np.cumsum(t_counts).astype(np.float64) + t_quantiles /= t_quantiles[-1] + + # interpolate linearly to find the pixel values in the template image + # that correspond most closely to the quantiles in the source image + interp_t_values = np.interp(s_quantiles, t_quantiles, t_values) + + return interp_t_values[bin_idx].reshape(oldshape) + + +def hist_match(image_name1, image_name2, output_name): + """ + Adjust the pixel values of a grayscale image such that its histogram + matches that of a target image + + Arguments: + ----------- + image name1, image name2 + Returns: + ----------- + output dir + imagename + """ + image1 = plt.imread(image_name1) + image2 = plt.imread(image_name2) + # source = face()[:,:,0] + # template = ascent() + # matched = hist_match(source, template) + + # def ecdf(x): + # """convenience function for computing the empirical CDF""" + # vals, counts = np.unique(x, return_counts=True) + # ecdf = np.cumsum(counts).astype(np.float64) + # ecdf /= ecdf[-1] + # return vals, ecdf + + # x1, y1 = ecdf(source.ravel()) + # x2, y2 = ecdf(template.ravel()) + # x3, y3 = ecdf(matched.ravel()) + + # fig = plt.figure() + # gs = plt.GridSpec(2, 3) + # ax1 = fig.add_subplot(gs[0, 0]) + # ax2 = fig.add_subplot(gs[0, 1], sharex=ax1, sharey=ax1) + # ax3 = fig.add_subplot(gs[0, 2], sharex=ax1, sharey=ax1) + # ax4 = fig.add_subplot(gs[1, :]) + # for aa in (ax1, ax2, ax3): + # aa.set_axis_off() + + # ax1.imshow(source, cmap=plt.cm.gray) + # ax1.set_title('Source') + # ax2.imshow(template, cmap=plt.cm.gray) + # ax2.set_title('template') + # ax3.imshow(matched, cmap=plt.cm.gray) + # ax3.set_title('Matched') + + # ax4.plot(x1, y1 * 100, '-r', lw=3, label='Source') + # ax4.plot(x2, y2 * 100, '-k', lw=3, label='Template') + # ax4.plot(x3, y3 * 100, '--r', lw=3, label='Matched') + # ax4.set_xlim(x1[0], x1[-1]) + # ax4.set_xlabel('Pixel value') + # ax4.set_ylabel('Cumulative %') + # ax4.legend(loc=5) + + # plt.show() + # print('finished testing') diff --git a/utils/model_util.py b/utils/model_util.py new file mode 100755 index 0000000..d00925a --- /dev/null +++ b/utils/model_util.py @@ -0,0 +1,62 @@ +import torch.nn as nn +import torchvision.models as models + +vgg_cfg = { + "A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], + "B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], + "D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"], + "E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"], +} + + +def _vgg_make_conv_layer_names(cfg, batch_norm=False): + layers = [] + block_idx = 1 + layer_idx = 1 + for v in cfg: + if v == "M": + layers += ["pool%d" % block_idx] + block_idx += 1 + layer_idx = 1 + else: + if batch_norm: + layers += [ + "conv%d_%d" % (block_idx, layer_idx), + "conv%d_%dnorm" % (block_idx, layer_idx), + "relu%d_%d" % (block_idx, layer_idx), + ] + else: + layers += ["conv%d_%d" % (block_idx, layer_idx), "relu%d_%d" % (block_idx, layer_idx)] + layer_idx += 1 + return layers + + +def _vgg_make_fc_layer_names(): + layers = ["fc6", "relu6", "dropout6", "fc7", "relu7", "dropout7", "fc8"] + return layers + + +class VGG19Before(nn.Module): + def __init__(self, layer_name, part="conv"): + super(VGG19Before, self).__init__() + assert part == "conv" or part == "fc", "invalid part name" + vgg_19 = models.vgg19(pretrained=True) + + if part == "conv": + vgg19_conv_layers = _vgg_make_conv_layer_names(vgg_cfg["E"], batch_norm=False) + idx = vgg19_conv_layers.index(layer_name) + assert 0 <= idx < len(vgg19_conv_layers), "cannot find %s in part %s" % (layer_name, part) + self.features = nn.Sequential(*list(vgg_19.features.children())[: (idx + 1)]) + else: + vgg19_fc_layers = _vgg_make_fc_layer_names() + idx = vgg19_fc_layers.index(layer_name) + assert 0 <= idx < len(vgg19_fc_layers), "cannot find %s in part %s" % (layer_name, part) + self.features = vgg_19.features + self.classifier = nn.Sequential(*list(vgg_19.classifier.children())[: (idx + 1)]) + + def forward(self, x): + x = self.features(x) + if hasattr(self, "classifier"): + x = self.classifier(x) + return x + diff --git a/utils/parallel.py b/utils/parallel.py new file mode 100644 index 0000000..af9cf5a --- /dev/null +++ b/utils/parallel.py @@ -0,0 +1,286 @@ +##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ +## Created by: Hang Zhang, Rutgers University, Email: zhang.hang@rutgers.edu +## Modified by Thomas Wolf, HuggingFace Inc., Email: thomas@huggingface.co +## Copyright (c) 2017-2018 +## +## This source code is licensed under the MIT-style license found in the +## LICENSE file in the root directory of this source tree +##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ +"""Encoding Data Parallel""" +import functools +import threading + +import torch +import torch.cuda.comm as comm +from torch.autograd import Function, Variable +from torch.nn.parallel import DistributedDataParallel +from torch.nn.parallel._functions import Broadcast, ReduceAddCoalesced +from torch.nn.parallel.data_parallel import DataParallel +from torch.nn.parallel.parallel_apply import get_a_var +from torch.nn.parallel.scatter_gather import gather + +torch_ver = torch.__version__[:3] + +__all__ = ["allreduce", "DataParallelModel", "DataParallelCriterion", "patch_replication_callback"] + + +def allreduce(*inputs): + """Cross GPU all reduce autograd operation for calculate mean and + variance in SyncBN. + """ + return AllReduce.apply(*inputs) + + +class AllReduce(Function): + @staticmethod + def forward(ctx, num_inputs, *inputs): + ctx.num_inputs = num_inputs + ctx.target_gpus = [inputs[i].get_device() for i in range(0, len(inputs), num_inputs)] + inputs = [inputs[i : i + num_inputs] for i in range(0, len(inputs), num_inputs)] + # sort before reduce sum + inputs = sorted(inputs, key=lambda i: i[0].get_device()) + results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0]) + outputs = comm.broadcast_coalesced(results, ctx.target_gpus) + return tuple([t for tensors in outputs for t in tensors]) + + @staticmethod + def backward(ctx, *inputs): + inputs = [i.data for i in inputs] + inputs = [inputs[i : i + ctx.num_inputs] for i in range(0, len(inputs), ctx.num_inputs)] + results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0]) + outputs = comm.broadcast_coalesced(results, ctx.target_gpus) + return (None,) + tuple([Variable(t) for tensors in outputs for t in tensors]) + + +class Reduce(Function): + @staticmethod + def forward(ctx, *inputs): + ctx.target_gpus = [inputs[i].get_device() for i in range(len(inputs))] + inputs = sorted(inputs, key=lambda i: i.get_device()) + return comm.reduce_add(inputs) + + @staticmethod + def backward(ctx, gradOutput): + return Broadcast.apply(ctx.target_gpus, gradOutput) + + +class DistributedDataParallelModel(DistributedDataParallel): + """Implements data parallelism at the module level for the DistributedDataParallel module. + This container parallelizes the application of the given module by + splitting the input across the specified devices by chunking in the + batch dimension. + In the forward pass, the module is replicated on each device, + and each replica handles a portion of the input. During the backwards pass, + gradients from each replica are summed into the original module. + Note that the outputs are not gathered, please use compatible + :class:`encoding.parallel.DataParallelCriterion`. + The batch size should be larger than the number of GPUs used. It should + also be an integer multiple of the number of GPUs so that each chunk is + the same size (so that each GPU processes the same number of samples). + Args: + module: module to be parallelized + device_ids: CUDA devices (default: all devices) + Reference: + Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi, + Amit Agrawal. “Context Encoding for Semantic Segmentation. + *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018* + Example:: + >>> net = encoding.nn.DistributedDataParallelModel(model, device_ids=[0, 1, 2]) + >>> y = net(x) + """ + + def gather(self, outputs, output_device): + return outputs + + +class DataParallelModel(DataParallel): + """Implements data parallelism at the module level. + + This container parallelizes the application of the given module by + splitting the input across the specified devices by chunking in the + batch dimension. + In the forward pass, the module is replicated on each device, + and each replica handles a portion of the input. During the backwards pass, + gradients from each replica are summed into the original module. + Note that the outputs are not gathered, please use compatible + :class:`encoding.parallel.DataParallelCriterion`. + + The batch size should be larger than the number of GPUs used. It should + also be an integer multiple of the number of GPUs so that each chunk is + the same size (so that each GPU processes the same number of samples). + + Args: + module: module to be parallelized + device_ids: CUDA devices (default: all devices) + + Reference: + Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi, + Amit Agrawal. “Context Encoding for Semantic Segmentation. + *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018* + + Example:: + + >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2]) + >>> y = net(x) + """ + + def gather(self, outputs, output_device): + return outputs + + def replicate(self, module, device_ids): + modules = super(DataParallelModel, self).replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + +class DataParallelCriterion(DataParallel): + """ + Calculate loss in multiple-GPUs, which balance the memory usage. + The targets are splitted across the specified devices by chunking in + the batch dimension. Please use together with :class:`encoding.parallel.DataParallelModel`. + + Reference: + Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi, + Amit Agrawal. “Context Encoding for Semantic Segmentation. + *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018* + + Example:: + + >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2]) + >>> criterion = encoding.nn.DataParallelCriterion(criterion, device_ids=[0, 1, 2]) + >>> y = net(x) + >>> loss = criterion(y, target) + """ + + def forward(self, inputs, *targets, **kwargs): + # input should be already scatterd + # scattering the targets instead + if not self.device_ids: + return self.module(inputs, *targets, **kwargs) + targets, kwargs = self.scatter(targets, kwargs, self.device_ids) + if len(self.device_ids) == 1: + return self.module(inputs, *targets[0], **kwargs[0]) + replicas = self.replicate(self.module, self.device_ids[: len(inputs)]) + outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs) + # return Reduce.apply(*outputs) / len(outputs) + # return self.gather(outputs, self.output_device).mean() + return self.gather(outputs, self.output_device) + + +def _criterion_parallel_apply(modules, inputs, targets, kwargs_tup=None, devices=None): + assert len(modules) == len(inputs) + assert len(targets) == len(inputs) + if kwargs_tup: + assert len(modules) == len(kwargs_tup) + else: + kwargs_tup = ({},) * len(modules) + if devices is not None: + assert len(modules) == len(devices) + else: + devices = [None] * len(modules) + + lock = threading.Lock() + results = {} + if torch_ver != "0.3": + grad_enabled = torch.is_grad_enabled() + + def _worker(i, module, input, target, kwargs, device=None): + if torch_ver != "0.3": + torch.set_grad_enabled(grad_enabled) + if device is None: + device = get_a_var(input).get_device() + try: + with torch.cuda.device(device): + # this also avoids accidental slicing of `input` if it is a Tensor + if not isinstance(input, (list, tuple)): + input = (input,) + if not isinstance(target, (list, tuple)): + target = (target,) + output = module(*(input + target), **kwargs) + with lock: + results[i] = output + except Exception as e: + with lock: + results[i] = e + + if len(modules) > 1: + threads = [ + threading.Thread(target=_worker, args=(i, module, input, target, kwargs, device),) + for i, (module, input, target, kwargs, device) in enumerate( + zip(modules, inputs, targets, kwargs_tup, devices) + ) + ] + + for thread in threads: + thread.start() + for thread in threads: + thread.join() + else: + _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0]) + + outputs = [] + for i in range(len(inputs)): + output = results[i] + if isinstance(output, Exception): + raise output + outputs.append(output) + return outputs + + +########################################################################### +# Adapted from Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# +class CallbackContext(object): + pass + + +def execute_replication_callbacks(modules): + """ + Execute an replication callback `__data_parallel_replicate__` on each module created + by original replication. + + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Note that, as all modules are isomorphism, we assign each sub-module with a context + (shared among multiple copies of this module on different devices). + Through this context, different copies can share some information. + + We guarantee that the callback on the master copy (the first copy) will be called ahead + of calling the callback of any slave copies. + """ + master_copy = modules[0] + nr_modules = len(list(master_copy.modules())) + ctxs = [CallbackContext() for _ in range(nr_modules)] + + for i, module in enumerate(modules): + for j, m in enumerate(module.modules()): + if hasattr(m, "__data_parallel_replicate__"): + m.__data_parallel_replicate__(ctxs[j], i) + + +def patch_replication_callback(data_parallel): + """ + Monkey-patch an existing `DataParallel` object. Add the replication callback. + Useful when you have customized `DataParallel` implementation. + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) + > patch_replication_callback(sync_bn) + # this is equivalent to + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + """ + + assert isinstance(data_parallel, DataParallel) + + old_replicate = data_parallel.replicate + + @functools.wraps(old_replicate) + def new_replicate(module, device_ids): + modules = old_replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + data_parallel.replicate = new_replicate diff --git a/utils/tb_image_recorder.py b/utils/tb_image_recorder.py new file mode 100644 index 0000000..26d3ee1 --- /dev/null +++ b/utils/tb_image_recorder.py @@ -0,0 +1,29 @@ +import threading + +import numpy as np + + +class TBImageRecorder(threading.Thread): + """ + TBImageRecorder + """ + + def __init__(self, tb_writer, func, queue): + super(TBImageRecorder, self).__init__() + self._tb_writer = tb_writer + # self._tb_writer = SummaryWriter() + self._func = func + self._queue = queue + + def run(self): + while True: + msgs, iter_index = self._queue.get() + if msgs: + img_info = self._func(*msgs) + + print("logging the images") + for tag, images in img_info.items(): + if images is not None: + self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index) + else: + break diff --git a/utils/tensor_lab2rgb.py b/utils/tensor_lab2rgb.py new file mode 100755 index 0000000..f162543 --- /dev/null +++ b/utils/tensor_lab2rgb.py @@ -0,0 +1,67 @@ +""""xyz2rgb(lab2xyz(lab, illuminant, observer))""" +import numpy as np +import torch +from scipy import linalg +from torch.autograd import Variable + +xyz_from_rgb = np.array( + [[0.412453, 0.357580, 0.180423], [0.212671, 0.715160, 0.072169], [0.019334, 0.119193, 0.950227]] +) + +# rgb_from_xyz = linalg.inv(xyz_from_rgb).T +rgb_from_xyz = np.array( + [[3.24048134, -0.96925495, 0.05564664], [-1.53715152, 1.87599, -0.20404134], [-0.49853633, 0.04155593, 1.05731107]] +) + + +def tensor_lab2rgb(input): + """ + n * 3* h *w + """ + input_trans = input.transpose(1, 2).transpose(2, 3) # n * h * w * 3 + L, a, b = input_trans[:, :, :, 0:1], input_trans[:, :, :, 1:2], input_trans[:, :, :, 2:] + y = (L + 16.0) / 116.0 + x = (a / 500.0) + y + z = y - (b / 200.0) + + neg_mask = z.data < 0 + z[neg_mask] = 0 + # if neg_mask.any(): + # z[neg_mask] = 0 + xyz = torch.cat((x, y, z), dim=3) + + mask = xyz.data > 0.2068966 + mask_xyz = xyz.clone() + mask_xyz[mask] = torch.pow(xyz[mask], 3.0) + mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.0) / 7.787 + # if mask.any(): + # mask_xyz[mask] = torch.pow(xyz[mask], 3.) + # if not mask.all(): + # mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.) / 7.787 + mask_xyz[:, :, :, 0] = mask_xyz[:, :, :, 0] * 0.95047 + mask_xyz[:, :, :, 2] = mask_xyz[:, :, :, 2] * 1.08883 + + rgb_trans = torch.mm(mask_xyz.view(-1, 3), torch.from_numpy(rgb_from_xyz).type_as(xyz)).view( + input.size(0), input.size(2), input.size(3), 3 + ) + rgb = rgb_trans.transpose(2, 3).transpose(1, 2) + + mask = rgb > 0.0031308 + mask_rgb = rgb.clone() + mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055 + mask_rgb[~mask] = rgb[~mask] * 12.92 + # if mask.any(): + # mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055 + # if not mask.all(): + # mask_rgb[~mask] = rgb[~mask] * 12.92 + + neg_mask = mask_rgb.data < 0 + large_mask = mask_rgb.data > 1 + mask_rgb[neg_mask] = 0 + mask_rgb[large_mask] = 1 + # if neg_mask.any(): + # mask_rgb[neg_mask] = 0 + # if large_mask.any(): + # mask_rgb[large_mask] = 1 + + return mask_rgb diff --git a/utils/tensorboard_logger.py b/utils/tensorboard_logger.py new file mode 100755 index 0000000..78acc37 --- /dev/null +++ b/utils/tensorboard_logger.py @@ -0,0 +1,88 @@ +# Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514 +import numpy as np +import scipy.misc +import tensorflow as tf + +try: + from StringIO import StringIO # Python 2.7 +except ImportError: + from io import BytesIO # Python 3.x + + +class TBLogger(object): + def __init__(self, log_dir): + """Create a summary writer logging to log_dir.""" + self.writer = tf.summary.FileWriter(log_dir) + + def scalar_summary(self, tag, value, step): + """Log a scalar variable.""" + summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) + self.writer.add_summary(summary, step) + + def image_summaries(self, tag, images, step): + """Log a list of images.""" + + img_summaries = [] + for i, img in enumerate(images): + # Write the image to a string + try: + s = StringIO() + except: + s = BytesIO() + scipy.misc.toimage(img).save(s, format="png") + + # Create an Image object + img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=img.shape[0], width=img.shape[1]) + # Create a Summary value + img_summaries.append(tf.Summary.Value(tag="%s/%d" % (tag, i), image=img_sum)) + + # Create and write Summary + summary = tf.Summary(value=img_summaries) + self.writer.add_summary(summary, step) + + def image_summary(self, tag, image, step): + """Log a list of images.""" + + img_summaries = [] + + # Write the image to a string + try: + s = StringIO() + except: + s = BytesIO() + + scipy.misc.toimage(image).save(s, format="png") + + # Create an Image object + img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=image.shape[0], width=image.shape[1]) + # Create and write Summary + summary = tf.Summary(value=[tf.Summary.Value(tag=tag, image=img_sum)]) + self.writer.add_summary(summary, step) + + def histo_summary(self, tag, values, step, bins=1000): + """Log a histogram of the tensor of values.""" + + # Create a histogram using numpy + counts, bin_edges = np.histogram(values, bins=bins) + + # Fill the fields of the histogram proto + hist = tf.HistogramProto() + hist.min = float(np.min(values)) + hist.max = float(np.max(values)) + hist.num = int(np.prod(values.shape)) + hist.sum = float(np.sum(values)) + hist.sum_squares = float(np.sum(values ** 2)) + + # Drop the start of the first bin + bin_edges = bin_edges[1:] + + # Add bin edges and counts + for edge in bin_edges: + hist.bucket_limit.append(edge) + for c in counts: + hist.bucket.append(c) + + # Create and write Summary + summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) + self.writer.add_summary(summary, step) + self.writer.flush() diff --git a/utils/test_html.py b/utils/test_html.py new file mode 100644 index 0000000..1df27c5 --- /dev/null +++ b/utils/test_html.py @@ -0,0 +1,66 @@ +import os + +import dominate +from dominate.tags import * + + +class HTML: + def __init__(self, web_dir, title, reflesh=0): + self.title = title + self.web_dir = web_dir + self.img_dir = os.path.join(self.web_dir, "images") + if not os.path.exists(self.web_dir): + os.makedirs(self.web_dir) + if not os.path.exists(self.img_dir): + os.makedirs(self.img_dir) + # print(self.img_dir) + + self.doc = dominate.document(title=title) + if reflesh > 0: + with self.doc.head: + meta(http_equiv="reflesh", content=str(reflesh)) + + def get_image_dir(self): + return self.img_dir + + def add_header(self, str): + with self.doc: + h3(str) + + def add_table(self, border=1): + self.t = table(border=border, style="table-layout: fixed;") + self.doc.add(self.t) + + def add_images(self, ims, txts, links, width=400): + self.add_table() + with self.t: + with tr(): + for im, txt, link in zip(ims, txts, links): + with td(style="word-wrap: break-word;", halign="center", valign="top"): + with p(): + with a(href=os.path.join("images", link)): + img(style="width:%dpx" % width, src=os.path.join("images", im)) + br() + p(txt) + + def save(self): + html_file = "%s/index.html" % self.web_dir + f = open(html_file, "wt") + f.write(self.doc.render()) + f.close() + + +if __name__ == "__main__": + html = HTML("web/", "test_html") + html.add_header("hello world") + + ims = [] + txts = [] + links = [] + for n in range(4): + ims.append("image_%d.jpg" % n) + txts.append("text_%d" % n) + links.append("image_%d.jpg" % n) + html.add_images(ims, txts, links) + html.add_images(ims, txts, links) + html.save() diff --git a/utils/test_lms_edge.py b/utils/test_lms_edge.py new file mode 100644 index 0000000..6a04206 --- /dev/null +++ b/utils/test_lms_edge.py @@ -0,0 +1,48 @@ +import matplotlib.pyplot as plt +import numpy as np +import skimage +from skimage import color, feature, io, morphology +from skimage.filters import prewitt, roberts, scharr, sobel +from skimage.morphology import square + + +def rgb2lms(rgb_image): + lms_image = rgb_image + r = rgb_image[:, :, 0] + g = rgb_image[:, :, 1] + b = rgb_image[:, :, 2] + + l = 17.8824 * r + 43.5161 * g + 4.1194 * b + m = 3.4557 * r + 27.1554 * g + 3.8671 * b + s = 0.0300 * r + 0.1843 * g + 1.4671 * b + lms_image = np.stack((l, m, s), axis=2) + return l, m, s + + +def image_normalize(rg_image): + rg_image = rg_image - rg_image.min() + rg_image = rg_image / (rg_image.max() - rg_image.min()) * 255 + return rg_image + + +image_rgb = io.imread("ref5.jpg") +plt.imshow(image_rgb, cmap="gray") +L, M, S = rgb2lms(image_rgb) +image_lab = color.rgb2lab(image_rgb) +a_edge = sobel(image_lab[:, :, 1]) +b_edge = sobel(image_lab[:, :, 2]) +color_edge = np.sqrt(a_edge ** 2 + b_edge ** 2) +luminance_edge = sobel(image_lab[:, :, 0]) +canny_edge_a = feature.canny(image_lab[:, :, 1], sigma=3, low_threshold=3) +canny_edge_b = feature.canny(image_lab[:, :, 2], sigma=3, low_threshold=3) + +plt.figure() +# edge = morphology.opening(np.logical_and(color_edge > 8, luminance_edge > 4)) +edge = morphology.thin( + morphology.closing((morphology.opening(np.logical_and(color_edge > 10, luminance_edge > 2))), square(1)) +) +plt.imshow(edge, cmap="gray") + +plt.figure() +plt.imshow(np.logical_and(canny_edge_a, canny_edge_b), cmap="gray") +plt.show() diff --git a/utils/util.py b/utils/util.py new file mode 100755 index 0000000..0a30212 --- /dev/null +++ b/utils/util.py @@ -0,0 +1,422 @@ +import numbers +import os +import shutil +import sys +import time + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import torch +import torchvision.utils as vutils +from skimage import color, io +from torch.autograd import Variable + +cv2.setNumThreads(0) + +# l: [-50,50] +# ab: [-128, 128] +l_norm, ab_norm = 1.0, 1.0 +l_mean, ab_mean = 50.0, 0 + + +###### utility ###### +def to_np(x): + return x.data.cpu().numpy() + + +def utf8_str(in_str): + try: + in_str = in_str.decode("UTF-8") + except Exception: + in_str = in_str.encode("UTF-8").decode("UTF-8") + return in_str + + +class MovingAvg(object): + def __init__(self, pool_size=100): + from queue import Queue + + self.pool = Queue(maxsize=pool_size) + self.sum = 0 + self.curr_pool_size = 0 + + def set_curr_val(self, val): + if not self.pool.full(): + self.curr_pool_size += 1 + self.pool.put_nowait(val) + else: + last_first_val = self.pool.get_nowait() + self.pool.put_nowait(val) + self.sum -= last_first_val + + self.sum += val + return self.sum / self.curr_pool_size + + +###### image normalization ###### +# normalization for l +def center_l(l): + l_mc = (l - l_mean) / l_norm + return l_mc + + +# denormalization for l +def uncenter_l(l): + return l * l_norm + l_mean + + +# normalization for ab +def center_ab(ab): + ab_mc = (ab - ab_mean) / ab_norm + return ab_mc + + +# normalization for lab image +def center_lab_img(img_lab): + img_lab_mc = ( + img_lab / np.array((l_norm, ab_norm, ab_norm))[:, np.newaxis, np.newaxis] + - np.array((l_mean / l_norm, ab_mean / ab_norm, ab_mean / ab_norm))[:, np.newaxis, np.newaxis] + ) + return img_lab_mc + + +###### color space transformation ###### +def rgb2lab_transpose(img_rgb): + img_lab = color.rgb2lab(img_rgb).transpose((2, 0, 1)) + return img_lab + + +def lab2rgb(img_l, img_ab): + """ INPUTS + img_l XxXx1 [0,100] + img_ab XxXx2 [-100,100] + OUTPUTS + returned value is XxXx3 """ + pred_lab = np.concatenate((img_l, img_ab), axis=2).astype("float64") + pred_rgb = color.lab2rgb(pred_lab) + pred_rgb = (np.clip(pred_rgb, 0, 1) * 255).astype("uint8") + return pred_rgb + + +def gray2rgb_batch(l): + # gray image tensor to rgb image tensor + l_uncenter = uncenter_l(l) + l_uncenter = l_uncenter / (2 * l_mean) + return torch.cat((l_uncenter, l_uncenter, l_uncenter), dim=1) + + +def lab2rgb_transpose(img_l, img_ab): + """ INPUTS + img_l 1xXxX [0,100] + img_ab 2xXxX [-100,100] + OUTPUTS + returned value is XxXx3 """ + pred_lab = np.concatenate((img_l, img_ab), axis=0).transpose((1, 2, 0)) + pred_rgb = (np.clip(color.lab2rgb(pred_lab), 0, 1) * 255).astype("uint8") + return pred_rgb + + +def lab2rgb_transpose_mc(img_l_mc, img_ab_mc): + if isinstance(img_l_mc, Variable): + img_l_mc = img_l_mc.data.cpu() + if isinstance(img_ab_mc, Variable): + img_ab_mc = img_ab_mc.data.cpu() + + if img_l_mc.is_cuda: + img_l_mc = img_l_mc.cpu() + if img_ab_mc.is_cuda: + img_ab_mc = img_ab_mc.cpu() + + assert img_l_mc.dim() == 3 and img_ab_mc.dim() == 3, "only for batch input" + + img_l = img_l_mc * l_norm + l_mean + img_ab = img_ab_mc * ab_norm + ab_mean + pred_lab = torch.cat((img_l, img_ab), dim=0) + grid_lab = pred_lab.numpy().astype("float64") + grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8") + return grid_rgb + + +def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8): + if isinstance(img_l_mc, Variable): + img_l_mc = img_l_mc.data.cpu() + if isinstance(img_ab_mc, Variable): + img_ab_mc = img_ab_mc.data.cpu() + + if img_l_mc.is_cuda: + img_l_mc = img_l_mc.cpu() + if img_ab_mc.is_cuda: + img_ab_mc = img_ab_mc.cpu() + + assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input" + + img_l = img_l_mc * l_norm + l_mean + img_ab = img_ab_mc * ab_norm + ab_mean + pred_lab = torch.cat((img_l, img_ab), dim=1) + grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64") + grid_rgb = (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8") + return grid_rgb + + +###### loss functions ###### +def feature_normalize(feature_in): + feature_in_norm = torch.norm(feature_in, 2, 1, keepdim=True) + sys.float_info.epsilon + feature_in_norm = torch.div(feature_in, feature_in_norm) + return feature_in_norm + + +def statistics_matching(feature1, feature2): + N, C, H, W = feature1.shape + feature1 = feature1.view(N, C, -1) + feature2 = feature2.view(N, C, -1) + + mean1 = feature1.mean(dim=-1) + mean2 = feature2.mean(dim=-1) + std1 = feature1.var(dim=-1).sqrt() + std2 = feature2.var(dim=-1).sqrt() + + return mse_loss(mean1, mean2) + mse_loss(std1, std2) + + +def cosine_similarity(input, target): + input_norm = torch.norm(input, 2, 1, keepdim=True) + sys.float_info.epsilon + target_norm = torch.norm(target, 2, 1, keepdim=True) + sys.float_info.epsilon + normalized_input = torch.div(input, input_norm) + normalized_target = torch.div(target, target_norm) + cos_similarity = torch.mul(normalized_input, normalized_target) + return torch.sum(cos_similarity, dim=1, keepdim=True) + + +def mse_loss(input, target=0): + return torch.mean((input - target) ** 2) + + +def l1_loss(input, target=0): + return torch.mean(torch.abs(input - target)) + + +def calc_ab_gradient(input_ab): + x_grad = input_ab[:, :, :, 1:] - input_ab[:, :, :, :-1] + y_grad = input_ab[:, :, 1:, :] - input_ab[:, :, :-1, :] + return x_grad, y_grad + + +def calc_tv_loss(input): + x_grad = input[:, :, :, 1:] - input[:, :, :, :-1] + y_grad = input[:, :, 1:, :] - input[:, :, :-1, :] + loss = torch.sum(x_grad ** 2) / x_grad.nelement() + torch.sum(y_grad ** 2) / y_grad.nelement() + return loss + + +def calc_cosine_dist_loss(input, target): + input_norm = torch.norm(input, 2, 1, keepdim=True) + sys.float_info.epsilon + target_norm = torch.norm(target, 2, 1, keepdim=True) + sys.float_info.epsilon + normalized_input = torch.div(input, input_norm) + normalized_target = torch.div(target, target_norm) + cos_dist = torch.mul(normalized_input, normalized_target) + return torch.mean(1 - torch.sum(cos_dist, dim=1)) + + +def weighted_mse_loss(input, target, weights): + out = (input - target) ** 2 + out = out * weights.expand_as(out) + loss = out.mean() + return loss + + +def weighted_l1_loss(input, target, weights): + out = torch.abs(input - target) + out = out * weights.expand_as(out) + loss = out.mean() + return loss + + +def colorfulness(input_ab): + """ + according to the paper: Measuring colourfulness in natural images + input is batches of ab tensors in lab space + """ + N, C, H, W = input_ab.shape + a = input_ab[:, 0:1, :, :] + b = input_ab[:, 1:2, :, :] + + a = a.view(N, -1) + b = b.view(N, -1) + + sigma_a = torch.std(a, dim=-1) + sigma_b = torch.std(b, dim=-1) + + mean_a = torch.mean(a, dim=-1) + mean_b = torch.mean(b, dim=-1) + + colorfulness = torch.sqrt(sigma_a ** 2 + sigma_b ** 2) + 0.37 * torch.sqrt(mean_a ** 2 + mean_b ** 2) + return colorfulness + + +###### video related ####### +def save_frames(image, image_folder, index=None, image_name=None): + if image is not None: + image = np.clip(image, 0, 255).astype(np.uint8) + if image_name: + io.imsave(os.path.join(image_folder, image_name), image) + else: + io.imsave(os.path.join(image_folder, str(index).zfill(5) + ".jpg"), image) + + +def folder2vid(image_folder, output_dir, filename): + images = [img for img in os.listdir(image_folder) if img.endswith(".jpg")] + images.sort() + print(images) + frame = cv2.imread(os.path.join(image_folder, images[0])) + height, width, layers = frame.shape + # video = cv2.VideoWriter(os.path.join(output_dir, filename), 0x00000021, 24, (width, height)) + print("writing to video " + os.path.join(output_dir, filename)) + video = cv2.VideoWriter( + os.path.join(output_dir, filename), cv2.VideoWriter_fourcc("D", "I", "V", "X"), 24, (width, height) + ) + # video = cv2.VideoWriter(os.path.join(output_dir, filename), cv2.VideoWriter_fourcc('H', '2', '6', '4'), 24, (width, height)) + for image in images: + video.write(cv2.imread(os.path.join(image_folder, image))) + + # cv2.destroyAllWindows() + video.release() + + # import imageio + # frames = [] + # for image in images: + # frames.append(imageio.imread(os.path.join(image_folder, image))) + # imageio.mimsave('movie.gif', frames) + + +###### file system ###### +def get_size(start_path="."): + total_size = 0 + for dirpath, dirnames, filenames in os.walk(start_path): + for f in filenames: + fp = os.path.join(dirpath, f) + total_size += os.path.getsize(fp) + return total_size + + +def mkdir_if_not(dir_path): + if not os.path.exists(dir_path): + os.makedirs(dir_path) + + +def parse(parser, save=True): + opt = parser.parse_args(args=[]) + args = vars(opt) + + from time import gmtime, strftime + + print("------------ Options -------------") + for k, v in sorted(args.items()): + print("%s: %s" % (str(k), str(v))) + print("-------------- End ----------------") + + # save to the disk + if save: + file_name = os.path.join("opt.txt") + with open(file_name, "wt") as opt_file: + opt_file.write(os.path.basename(sys.argv[0]) + " " + strftime("%Y-%m-%d %H:%M:%S", gmtime()) + "\n") + opt_file.write("------------ Options -------------\n") + for k, v in sorted(args.items()): + opt_file.write("%s: %s\n" % (str(k), str(v))) + opt_file.write("-------------- End ----------------\n") + return opt + + +###### interactive ###### +def clean_tensorboard(directory): + folder_list = os.walk(directory).__next__()[1] + for folder in folder_list: + folder = directory + folder + if get_size(folder) < 10000000: + print("delete the folder of " + folder) + shutil.rmtree(folder) + + +def imshow(input_image, title=None, type_conversion=False): + inp = input_image + if type_conversion or type(input_image) is torch.Tensor: + inp = input_image.numpy() + else: + inp = input_image + fig = plt.figure() + if inp.ndim == 2: + fig = plt.imshow(inp, cmap="gray", clim=[0, 255]) + else: + fig = plt.imshow(np.transpose(inp, [1, 2, 0]).astype(np.uint8)) + plt.axis("off") + fig.axes.get_xaxis().set_visible(False) + fig.axes.get_yaxis().set_visible(False) + plt.title(title) + + +def imshow_lab(input_lab): + plt.imshow((batch_lab2rgb_transpose_mc(input_lab[:32, 0:1, :, :], input_lab[:32, 1:3, :, :])).astype(np.uint8)) + + +###### vgg preprocessing ###### +def vgg_preprocess(tensor): + # input is RGB tensor which ranges in [0,1] + # output is BGR tensor which ranges in [0,255] + tensor_bgr = torch.cat((tensor[:, 2:3, :, :], tensor[:, 1:2, :, :], tensor[:, 0:1, :, :]), dim=1) + # tensor_bgr = tensor[:, [2, 1, 0], ...] + tensor_bgr_ml = tensor_bgr - torch.Tensor([0.40760392, 0.45795686, 0.48501961]).type_as(tensor_bgr).view(1, 3, 1, 1) + tensor_rst = tensor_bgr_ml * 255 + return tensor_rst + + +def torch_vgg_preprocess(tensor): + # pytorch version normalization + # note that both input and output are RGB tensors; + # input and output ranges in [0,1] + # normalize the tensor with mean and variance + tensor_mc = tensor - torch.Tensor([0.485, 0.456, 0.406]).type_as(tensor).view(1, 3, 1, 1) + tensor_mc_norm = tensor_mc / torch.Tensor([0.229, 0.224, 0.225]).type_as(tensor_mc).view(1, 3, 1, 1) + return tensor_mc_norm + + +def network_gradient(net, gradient_on=True): + if gradient_on: + for param in net.parameters(): + param.requires_grad = True + else: + for param in net.parameters(): + param.requires_grad = False + return net + + +# def load_gray_image(img_path): +# img_path = utf8_str(img_path) +# img_gray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) +# return img_gray + +# def load_rgb_image(img_path): +# img_path = utf8_str(img_path) +# im = cv2.imread(img_path, cv2.IMREAD_COLOR) +# if im.shape[2] == 1: +# print('find gray images: ', img_path) +# im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR) +# img_rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) +# return img_rgb + +# def resize_img_longside(img_rgb, long_size, interpolation): +# h, w = img_rgb.shape[:2] +# scalar = long_size / max(h, w) +# resized_img_rgb = cv2.resize(img_rgb, (int(w * scalar), int(h * scalar)), interpolation=interpolation) +# return resized_img_rgb + +# def resize_img_shortside(img_rgb, short_size, interpolation): +# h, w = img_rgb.shape[:2] +# scalar = short_size / min(h, w) +# resized_img_rgb = cv2.resize(img_rgb, (int(w * scalar), int(h * scalar)), interpolation=interpolation) +# return resized_img_rgb + +# def resize_img(img_rgb, dsize, interpolation): +# if isinstance(dsize, numbers.Number): +# dsize = (int(dsize), int(dsize)) +# resized_img_rgb = cv2.resize(img_rgb, dsize, interpolation=interpolation) +# return resized_img_rgb diff --git a/utils/util_distortion.py b/utils/util_distortion.py new file mode 100644 index 0000000..7e18b96 --- /dev/null +++ b/utils/util_distortion.py @@ -0,0 +1,615 @@ +import math +import random + +import cv2 +import lib.functional as F +import matplotlib.pyplot as plt +import numpy as np +from PIL import Image +from scipy.misc import imread, imsave +from scipy.ndimage.filters import gaussian_filter +from skimage import color +from skimage import transform as skimage_transform +from skimage.draw import random_shapes +from skimage.filters import gaussian +from skimage.transform import resize, rotate + +cv2.setNumThreads(0) + +import time + +import scipy.io as scio +import skimage.io as io +from numba import jit, u1, u2 + + +class RGB2Lab(object): + def __init__(self): + pass + + def __call__(self, inputs): + image_lab = color.rgb2lab(inputs) + return image_lab + + +class Guassian_noise(object): + """Elastic distortion + """ + + def __init__(self, noise_sigma=0.1): + self.noise_sigma = noise_sigma + + def __call__(self, inputs): + h = inputs.shape[0] + w = inputs.shape[1] + noisy_image = inputs + noise = np.random.randn(h, w) * self.noise_sigma + noisy_image[:, :, 0] = inputs[:, :, 0] + noise + + return noisy_image + + +class Distortion(object): + """Elastic distortion + """ + + def __init__(self, distortion_level=3, flip_probability=0): + self.alpha_max = distortion_level + self.flip_probability = flip_probability + + def __call__(self, inputs): + if np.random.rand() < self.flip_probability: + inputs = inputs.transpose(Image.FLIP_LEFT_RIGHT) + + inputs = np.array(inputs) + alpha = np.random.rand() * self.alpha_max + sigma = 50 + random_state = np.random.RandomState(None) + shape = inputs.shape[0], inputs.shape[1] + + dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000 + dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000 + + x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij") + remap_image = cv2.remap( + inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR + ) + + return Image.fromarray(remap_image) + + +class Distortion_with_flow(object): + """Elastic distortion + """ + + def __init__(self): + return + + def __call__(self, inputs, dx, dy): + inputs = np.array(inputs) + shape = inputs.shape[0], inputs.shape[1] + inputs = np.array(inputs) + # x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij') + # remap_image = cv2.remap(inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR) + # backward mapping + + remap_image = forward_mapping(inputs, dy, dx, maxIter=3, precision=1e-3) + # remap_image = cv2.remap(inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT) + # remap_image = inputs + + return Image.fromarray(remap_image) + + +class Normalize(object): + def __init__(self): + pass + + def __call__(self, inputs): + inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1) + inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1)) + return inputs + + +class ToTensor(object): + def __init__(self): + pass + + def __call__(self, inputs): + outputs = F.to_mytensor(inputs) + return outputs + + +class RandomErasing(object): + """ + Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al. + ------------------------------------------------------------------------------------- + probability: The probability that the operation will be performed. + sl: min erasing area + sh: max erasing area + r1: min aspect ratio + mean: erasing value + ------------------------------------------------------------------------------------- + """ + + def __init__(self, probability=0.6, sl=0.05, sh=0.6): + self.probability = probability + self.sl = sl + self.sh = sh + + def __call__(self, img): + img = np.array(img) + if random.uniform(0, 1) > self.probability: + return Image.fromarray(img) + + area = img.shape[0] * img.shape[1] + h0 = img.shape[0] + w0 = img.shape[1] + channel = img.shape[2] + + h = int(round(random.uniform(self.sl, self.sh) * h0)) + w = int(round(random.uniform(self.sl, self.sh) * w0)) + + if w < img.shape[1] and h < img.shape[0]: + x1 = random.randint(0, img.shape[0] - h) + y1 = random.randint(0, img.shape[1] - w) + img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255 + return Image.fromarray(img) + + return Image.fromarray(img) + + +class CenteredPad(object): + """ + pad the frame with black border, + make square image for processing + """ + + def __init__(self, image_size): + self.image_size = image_size + + def __call__(self, img): + img = np.array(img) + width = np.size(img, 1) + height = np.size(img, 0) + old_size = [height, width] + + ratio = float(self.image_size) / max(height, width) + new_size = [int(x * ratio) for x in old_size] + I_resize = resize(img, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_new = np.size(I_resize, 1) + height_new = np.size(I_resize, 0) + + I_pad = np.zeros((self.image_size, self.image_size, 3)) + start_height = (self.image_size - new_size[0]) // 2 + start_width = (self.image_size - new_size[1]) // 2 + I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize + + return Image.fromarray(I_pad.astype(np.uint8)) + + +class centeredPad_with_height(object): + """ + pad the image according to the height + """ + + def __init__(self, image_size): + self.height = image_size[0] + self.width = image_size[1] + + def __call__(self, image): + # pad the image to 16:9 + # pad height + I = np.array(image) + + # for padded input + if I.shape[0] == I.shape[1] and I.shape[1] == self.width: + return Image.fromarray(I.astype(np.uint8)) + + width0 = np.size(I, 1) + height0 = np.size(I, 0) + old_size = [height0, width0] + height = self.height + width = self.width + + ratio = height / height0 + new_size = [int(x * ratio) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_new = np.size(I_resize, 1) + height_new = np.size(I_resize, 0) + + # if exceed the expected width + if width_new > width: + I_resize = I_resize[:, math.floor(width_new - width) // 2 : (math.floor(width_new - width) // 2 + width), :] + width_new = np.size(I_resize, 1) + height_new = np.size(I_resize, 0) + + # lines: 56~200 + I_pad = np.zeros((width, width, 3)) + start_height = (width - height_new) // 2 + start_width = (width - width_new) // 2 + I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize + + return Image.fromarray(I_pad.astype(np.uint8)) + + +class CenterPad(object): + def __init__(self, image_size): + self.height = image_size[0] + self.width = image_size[1] + + def __call__(self, image): + # pad the image to 16:9 + # pad height + I = np.array(image) + + # for padded input + height_old = np.size(I, 0) + width_old = np.size(I, 1) + old_size = [height_old, width_old] + height = self.height + width = self.width + I_pad = np.zeros((height, width, np.size(I, 2))) + + ratio = height / width + if height_old / width_old == ratio: + if height_old == height: + return Image.fromarray(I.astype(np.uint8)) + else: + new_size = [int(x * height / height_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + return Image.fromarray(I_resize.astype(np.uint8)) + + if height_old / width_old > ratio: # pad the width and crop + new_size = [int(x * width / width_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_height = (height_resize - height) // 2 + I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :] + else: # pad the height and crop + new_size = [int(x * height / height_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_width = (width_resize - width) // 2 + I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :] + + return Image.fromarray(I_pad.astype(np.uint8)) + + +class CenterPad_threshold(object): + def __init__(self, image_size, threshold=3 / 4): + self.height = image_size[0] + self.width = image_size[1] + self.threshold = threshold + + def __call__(self, image): + # pad the image to 16:9 + # pad height + I = np.array(image) + + # for padded input + height_old = np.size(I, 0) + width_old = np.size(I, 1) + old_size = [height_old, width_old] + height = self.height + width = self.width + I_pad = np.zeros((height, width, np.size(I, 2))) + + ratio = height / width + + if height_old / width_old == ratio: + if height_old == height: + return Image.fromarray(I.astype(np.uint8)) + else: + new_size = [int(x * height / height_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + return Image.fromarray(I_resize.astype(np.uint8)) + + if height_old / width_old > self.threshold: + width_new, height_new = width_old, int(width_old * self.threshold) + height_margin = height_old - height_new + height_crop_start = height_margin // 2 + I_crop = I[height_crop_start : (height_crop_start + height_new), :, :] + I_resize = resize( + I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True + ) + + # debug: + # height_new, width_new, channel = I_crop.shape + # print(width_new / height_new) + # import matplotlib.pyplot as plt + # plt.imshow(I / 255) + # plt.figure() + # plt.imshow(I_resize / 255) + return Image.fromarray(I_resize.astype(np.uint8)) + + if height_old / width_old > ratio: # pad the width and crop + new_size = [int(x * width / width_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_height = (height_resize - height) // 2 + I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :] + else: # pad the height and crop + new_size = [int(x * height / height_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_width = (width_resize - width) // 2 + I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :] + + return Image.fromarray(I_pad.astype(np.uint8)) + + +class CenterPadCrop_numpy(object): + """ + pad the image according to the height + """ + + def __init__(self, image_size): + self.height = image_size[0] + self.width = image_size[1] + + def __call__(self, image, threshold=3 / 4): + # pad the image to 16:9 + # pad height + I = np.array(image) + # for padded input + height_old = np.size(I, 0) + width_old = np.size(I, 1) + old_size = [height_old, width_old] + height = self.height + width = self.width + padding_size = width + if image.ndim == 2: + I_pad = np.zeros((width, width)) + else: + I_pad = np.zeros((width, width, I.shape[2])) + + ratio = height / width + if height_old / width_old == ratio: + return I + + if height_old / width_old > threshold: + width_new, height_new = width_old, int(width_old * threshold) + height_margin = height_old - height_new + height_crop_start = height_margin // 2 + I_crop = I[height_start : (height_start + height_new), :] + I_resize = resize( + I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True + ) + return I_resize + + if height_old / width_old > ratio: # pad the width and crop + new_size = [int(x * width / width_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_height = (height_resize - height) // 2 + start_height_block = (padding_size - height) // 2 + if image.ndim == 2: + I_pad[start_height_block : (start_height_block + height), :] = I_resize[ + start_height : (start_height + height), : + ] + else: + I_pad[start_height_block : (start_height_block + height), :, :] = I_resize[ + start_height : (start_height + height), :, : + ] + else: # pad the height and crop + new_size = [int(x * height / height_old) for x in old_size] + I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True) + width_resize = np.size(I_resize, 1) + height_resize = np.size(I_resize, 0) + start_width = (width_resize - width) // 2 + start_width_block = (padding_size - width) // 2 + if image.ndim == 2: + I_pad[:, start_width_block : (start_width_block + width)] = I_resize[ + :, start_width : (start_width + width) + ] + + else: + I_pad[:, start_width_block : (start_width_block + width), :] = I_resize[ + :, start_width : (start_width + width), : + ] + + crop_start_height = (I_pad.shape[0] - height) // 2 + crop_start_width = (I_pad.shape[1] - width) // 2 + + if image.ndim == 2: + return I_pad[ + crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width) + ] + else: + return I_pad[ + crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width), : + ] + + +@jit(nopython=True) +def iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision): + # print('processing location', (xr, yr)) + # + if abs(padu[yr, xr]) < precision and abs(padv[yr, xr]) < precision: + return xr, yr + + else: + # Our initialize method in this paper, can see the overleaf for detail + if (xr + 1) <= (W - 1): + dif = padu[yr, xr + 1] - padu[yr, xr] + u_next = padu[yr, xr] / (1 + dif) + else: + dif = padu[yr, xr] - padu[yr, xr - 1] + u_next = padu[yr, xr] / (1 + dif) + + if (yr + 1) <= (H - 1): + dif = padv[yr + 1, xr] - padv[yr, xr] + v_next = padv[yr, xr] / (1 + dif) + else: + dif = padv[yr, xr] - padv[yr - 1, xr] + v_next = padv[yr, xr] / (1 + dif) + + i = xr - u_next + j = yr - v_next + i_int = int(i) + j_int = int(j) + + # The same as traditional iterative search method + for iter in range(maxIter): + if 0 <= i <= (W - 1) and 0 <= j <= (H - 1): + u11 = padu[j_int, i_int] + v11 = padv[j_int, i_int] + + u12 = padu[j_int, i_int + 1] + v12 = padv[j_int, i_int + 1] + + int1 = padu[j_int + 1, i_int] + v21 = padv[j_int + 1, i_int] + + int2 = padu[j_int + 1, i_int + 1] + v22 = padv[j_int + 1, i_int + 1] + + u = ( + u11 * (i_int + 1 - i) * (j_int + 1 - j) + + u12 * (i - i_int) * (j_int + 1 - j) + + int1 * (i_int + 1 - i) * (j - j_int) + + int2 * (i - i_int) * (j - j_int) + ) + + v = ( + v11 * (i_int + 1 - i) * (j_int + 1 - j) + + v12 * (i - i_int) * (j_int + 1 - j) + + v21 * (i_int + 1 - i) * (j - j_int) + + v22 * (i - i_int) * (j - j_int) + ) + + i_next = xr - u + j_next = yr - v + + if abs(i - i_next) < precision and abs(j - j_next) < precision: + return i, j + + i = i_next + j = j_next + + else: + return i, j + + # if the search doesn't converge within max iter, it will return the last iter result + return i_next, j_next + + +# Bilinear interpolation +@jit(nopython=True) +def biInterpolation(distorted, i, j): + i = u2(i) + j = u2(j) + Q11 = distorted[j, i] + Q12 = distorted[j, i + 1] + Q21 = distorted[j + 1, i] + Q22 = distorted[j + 1, i + 1] + + pixel = u1( + Q11 * (i + 1 - i) * (j + 1 - j) + + Q12 * (i - i) * (j + 1 - j) + + Q21 * (i + 1 - i) * (j - j) + + Q22 * (i - i) * (j - j) + ) + return pixel + + +@jit(nopython=True) +def iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter=5, precision=1e-2): + for xr in range(W): + for yr in range(H): + # (xr, yr) is the point in result image, (i, j) is the search result in distorted image + i, j = iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision) + + # reflect the pixels outside the border + if i > W - 1: + i = 2 * W - 1 - i + if i < 0: + i = -i + if j > H - 1: + j = 2 * H - 1 - j + if j < 0: + j = -j + + # Bilinear interpolation to get the pixel at (i, j) in distorted image + resultImg[yr, xr, 0] = biInterpolation(distortImg[:, :, 0], i, j,) + resultImg[yr, xr, 1] = biInterpolation(distortImg[:, :, 1], i, j,) + resultImg[yr, xr, 2] = biInterpolation(distortImg[:, :, 2], i, j,) + return None + + +def forward_mapping(source_image, u, v, maxIter=5, precision=1e-2): + """ + warp the image according to the forward flow + u: horizontal + v: vertical + """ + H = source_image.shape[0] + W = source_image.shape[1] + + distortImg = np.array(np.zeros((H + 1, W + 1, 3)), dtype=np.uint8) + distortImg[0:H, 0:W] = source_image[0:H, 0:W] + distortImg[H, 0:W] = source_image[H - 1, 0:W] + distortImg[0:H, W] = source_image[0:H, W - 1] + distortImg[H, W] = source_image[H - 1, W - 1] + + padu = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32) + padu[0:H, 0:W] = u[0:H, 0:W] + padu[H, 0:W] = u[H - 1, 0:W] + padu[0:H, W] = u[0:H, W - 1] + padu[H, W] = u[H - 1, W - 1] + + padv = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32) + padv[0:H, 0:W] = v[0:H, 0:W] + padv[H, 0:W] = v[H - 1, 0:W] + padv[0:H, W] = v[0:H, W - 1] + padv[H, W] = v[H - 1, W - 1] + + resultImg = np.array(np.zeros((H, W, 3)), dtype=np.uint8) + iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter, precision) + return resultImg + + +def random_mask(H, W, mask_size=200): + """ + mask: ranges in [0,1] + """ + masked_image = np.zeros([H, W, 3]) + mask = random_shapes( + [H, W], + max_shapes=1, + min_shapes=1, + max_size=mask_size, + min_size=mask_size / 2, + multichannel=False, + intensity_range=[0, 0], + )[0] + mask = np.stack((mask, mask, mask), axis=-1) + random_state = np.random.RandomState(None) + distortion_range = 50 + alpha = np.random.rand() * 6 + forward_dx = ( + gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000 + ) + forward_dy = ( + gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000 + ) + mask = forward_mapping(mask, forward_dy, forward_dx, maxIter=3, precision=1e-3) / 255 + mask = 1 - gaussian(mask, sigma=0, preserve_range=True, multichannel=False, anti_aliasing=True) + mask = mask[:, :, 0] + # masked_image = mask * image_foreground + (1 - mask) * image_background + # masked_image = mask * image_foreground + + # randomly transform the mask + # rotate_degree = np.random.random() * 3.14 + # translation = [np.random.random() * H, np.random.random() * W] + # tform = skimage_transform.SimilarityTransform(scale=1, rotation=rotate_degree, translation=translation) + + # visualization + # plt.imshow(masked_image / 255) + # plt.show(block=True) + return mask diff --git a/utils/util_tensorboard.py b/utils/util_tensorboard.py new file mode 100644 index 0000000..2a30b60 --- /dev/null +++ b/utils/util_tensorboard.py @@ -0,0 +1,14 @@ +from utils.util import to_np + + +def histogram_logger(tb_writer, iter_index, net=None): + if net is not None: + for tag, value in net.named_parameters(): + tag = tag.replace(".", "/") + tb_writer.add_histogram(tag, to_np(value), iter_index) + tb_writer.add_histogram(tag + "/grad", to_np(value.grad), iter_index) + + +def value_logger(tb_writer, iter_index, loss_info): + for tag, value in loss_info.items(): + tb_writer.add_scalar(tag, value, iter_index) diff --git a/utils/vgg_util.py b/utils/vgg_util.py new file mode 100755 index 0000000..dfe5a31 --- /dev/null +++ b/utils/vgg_util.py @@ -0,0 +1,136 @@ +import collections +import os + +import torch +import torch.nn as nn +import torchvision.models as models +import torchvision.transforms as transforms +from PIL import Image + + +def preprocess(img, scale_size=None): + """PILimg: RGB: HxWxC""" + if scale_size is not None: + scale_transforms = MaxScale(scale_size) + img = scale_transforms(img) + + prep_transforms = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to BGR + transforms.Normalize(mean=[0.40760392, 0.45795686, 0.48501961], std=[1, 1, 1]), # subtract imagenet mean + transforms.Lambda(lambda x: x.mul_(255)), + ] + ) + img = prep_transforms(img) + return img + + +def deprocess(img): + post_transforms_a = transforms.Compose( + [ + transforms.Lambda(lambda x: x.mul_(1.0 / 255)), + transforms.Normalize(mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1]), # add imagenet mean + transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB + ] + ) + post_transforms_b = transforms.Compose([transforms.ToPILImage()]) + img = post_transforms_a(img) + img[img > 1] = 1 + img[img < 0] = 0 + img = post_transforms_b(img) + return img + + +def get_renamed_vgg(): + cache_file = "data/vgg19_conv.pth" + vgg = models.vgg19().features + renamed_vgg = nn.Sequential() + part_idx, layer_idx = 1, 1 + for layer in list(vgg): + if isinstance(layer, nn.Conv2d): + name = "conv{}_{}".format(part_idx, layer_idx) + renamed_vgg.add_module(name, layer) + elif isinstance(layer, nn.ReLU): + name = "relu{}_{}".format(part_idx, layer_idx) + renamed_vgg.add_module(name, layer) + layer_idx += 1 + elif isinstance(layer, nn.MaxPool2d): + name = "pool{}".format(part_idx) + renamed_vgg.add_module(name, layer) + part_idx += 1 + layer_idx = 1 + renamed_vgg.load_state_dict(torch.load(cache_file)) + + return renamed_vgg + + +def get_renamed_vgg_johnson(): + cache_file = "data/vgg19-d01eb7cb.pth" + vgg = models.vgg19() + model_dict = vgg.state_dict() + pretrained_dict = torch.load(cache_file) + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + model_dict.update(pretrained_dict) + vgg.load_state_dict(model_dict) + vgg = vgg.features + renamed_vgg = nn.Sequential() + part_idx, layer_idx = 1, 1 + for layer in list(vgg): + if isinstance(layer, nn.Conv2d): + name = "conv{}_{}".format(part_idx, layer_idx) + renamed_vgg.add_module(name, layer) + elif isinstance(layer, nn.ReLU): + name = "relu{}_{}".format(part_idx, layer_idx) + renamed_vgg.add_module(name, layer) + layer_idx += 1 + elif isinstance(layer, nn.MaxPool2d): + name = "pool{}".format(part_idx) + renamed_vgg.add_module(name, layer) + part_idx += 1 + layer_idx = 1 + + return renamed_vgg + + +class MaxScale(object): + """Rescale the input PIL.Image to the given size. + + Args: + size (sequence or int): Desired output size. If size is a sequence like + (w, h), output size will be matched to this. If size is an int, + smaller edge of the image will be matched to this number. + i.e, if height > width, then image will be rescaled to + (size * height / width, size) + interpolation (int, optional): Desired interpolation. Default is + ``PIL.Image.BILINEAR`` + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2) + self.size = size + self.interpolation = interpolation + + def __call__(self, img): + """ + Args: + img (PIL.Image): Image to be scaled. + + Returns: + PIL.Image: Rescaled image. + """ + if isinstance(self.size, int): + w, h = img.size + if (w <= h == self.size) or (h <= w == self.size): + return img + if w < h: + oh = self.size + ow = int(self.size * w / h) + return img.resize((ow, oh), self.interpolation) + else: + ow = self.size + oh = int(self.size * h / w) + return img.resize((ow, oh), self.interpolation) + else: + return img.resize(self.size, self.interpolation) + diff --git a/utils/visualize_analogy.py b/utils/visualize_analogy.py new file mode 100644 index 0000000..81d9241 --- /dev/null +++ b/utils/visualize_analogy.py @@ -0,0 +1,105 @@ +import glob +import os +import os.path as osp +import struct + +import cv2 +import numpy as np +import torch.utils.data as data +from PIL import Image + +cv2.setNumThreads(0) +import pdb + +import lib.functional as F +import matplotlib.pyplot as plt +import torch +import torchvision.transforms as transforms +from PIL import Image +from skimage import color + +from utils.util import ( + MovingAvg, + batch_lab2rgb_transpose_mc, + calc_ab_gradient, + calc_cosine_dist_loss, + calc_tv_loss, + clean_tensorboard, + feature_normalize, + imshow, + imshow_lab, + mkdir_if_not, + to_np, + uncenter_l, + vgg_preprocess, + weighted_l1_loss, +) + + +def combo5_loader_nonoptimize(path): + f = open(path, "rb") + + # width, height + d = f.read(4) + im_sz = struct.unpack("i", d) + h = im_sz[0] + + d = f.read(4) + im_sz = struct.unpack("i", d) + w = im_sz[0] + + # warp_ba_layer 4 + d = f.read(4) + im_sz = struct.unpack("i", d) + d = f.read(im_sz[0]) + file_bytes = np.asarray(bytearray(d), dtype=np.uint8) + img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) + img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB) + warp_ba = Image.fromarray(img_data_ndarray) + # warp_ba.save('/home/v-mingmh/Documents/gray/test_input/warp_ba.PNG') + + # warp_aba_layer 4 + d = f.read(4) + im_sz = struct.unpack("i", d) + d = f.read(im_sz[0]) + file_bytes = np.asarray(bytearray(d), dtype=np.uint8) + img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) + img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB) + warp_aba = Image.fromarray(img_data_ndarray) + + # 5 layers: err_aba, err_ba, err_ab + errs = [] + for l in range(5): + d = f.read(4) + im_sz = struct.unpack("i", d) + d = f.read(im_sz[0]) + file_bytes = np.asarray(bytearray(d), dtype=np.uint8) + img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_GRAYSCALE) + err_ba = Image.fromarray(img_data_ndarray) + + # err_ba.save('/home/v-mingmh/Documents/gray/test_input/input_err_ba' + str(l) + '.PNG') + + d = f.read(4) + im_sz = struct.unpack("i", d) + d = f.read(im_sz[0]) + file_bytes = np.asarray(bytearray(d), dtype=np.uint8) + img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_GRAYSCALE) + err_ab = Image.fromarray(img_data_ndarray) + + # err_ab.save('/home/v-mingmh/Documents/gray/test_input/input_err_ab' + str(l) + '.PNG') + + errs.append([err_ba, err_ab]) + + f.close() + + return errs, warp_ba, warp_aba + + +# combo_folder = '../../movie/seq1/combo_new/' +combo_folder = "../../movie/seq2/combo_new" +output_folder = "analogy_output2/" +path, dirs, filenames = os.walk(combo_folder).__next__() + +for filename in filenames: + errs, warp_ba, warp_aba = combo5_loader_nonoptimize(combo_folder + "/" + filename) + warp_ba.save(output_folder + filename.split(".")[0] + ".jpg") diff --git a/utils/warping.py b/utils/warping.py new file mode 100644 index 0000000..251586f --- /dev/null +++ b/utils/warping.py @@ -0,0 +1,34 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable + + +def get_grid(x): + torchHorizontal = ( + torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + torchVertical = ( + torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3)) + ) + grid = torch.cat([torchHorizontal, torchVertical], 1).cuda() + + return grid + + +class WarpingLayer(nn.Module): + def __init__(self, device): + super(WarpingLayer, self).__init__() + self.device = device + + def forward(self, x, flow): + # WarpingLayer uses F.grid_sample, which expects normalized grid + # we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code + # so here we need to denormalize the flow + flow_for_grip = torch.zeros_like(flow).cuda() + flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0) + flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0) + + grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1) + x_warp = F.grid_sample(x, grid) + return x_warp