Repository: ibaiGorordo/HITNET-Stereo-Depth-estimation
Branch: main
Commit: a64564954bfa
Files: 11
Total size: 13.0 KB
Directory structure:
gitextract_i5i3wbqa/
├── .gitignore
├── LICENSE
├── README.md
├── drivingStereoTest.py
├── hitnet/
│ ├── __init__.py
│ ├── hitnet.py
│ └── utils_hitnet.py
├── imageDepthEstimation.py
├── models/
│ └── .gitkeep
├── requirements.txt
└── videoDepthEstimation.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2021 Ibai Gorordo
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: README.md
================================================
# HITNET-Stereo-Depth-estimation
Python scripts for performing stereo depth estimation using the [HITNET Tensorflow model from Google Research](https://github.com/google-research/google-research/tree/master/hitnet).
*Stereo depth estimation on the cones images from the Middlebury dataset (https://vision.middlebury.edu/stereo/data/scenes2003/)*
# Requirements
* **OpenCV**, **numpy** and **tensorflo**. **pafy** (`pip install git+https://github.com/zizo-pro/pafy@b8976f22c19e4ab5515cacbfae0a3970370c102b`) and **youtube-dl** are required for youtube video inference.
* For the drivingStereo dataset, download the data from: https://drivingstereo-dataset.github.io/
# Tensorflow models
Download the tensorflow models from the [original repository](https://github.com/google-research/google-research/tree/master/hitnet) and save them into the **[models](https://github.com/ibaiGorordo/HITNET-Stereo-Depth-estimation/tree/main/models)** folder.
# Examples
* **Image inference**:
```
python imageDepthEstimation.py
```
* **Video inference**:
```
python videoDepthEstimation.py
```
* **DrivingStereo dataset inference**:
```
python drivingStereoTest.py
```
# [Inference video Example](https://youtu.be/ge2iN8Ga4Dg)
# References:
* Hitnet model: https://github.com/google-research/google-research/tree/master/hitnet
* DrivingStereo dataset: https://drivingstereo-dataset.github.io/
* Original paper: https://arxiv.org/abs/2007.12140
================================================
FILE: drivingStereoTest.py
================================================
import cv2
import pafy
import tensorflow as tf
import numpy as np
import glob
from hitnet import HitNet, ModelType, draw_disparity, draw_depth, CameraConfig
out = cv2.VideoWriter('outpy2.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30, (881*3,400))
# Get image list
left_images = glob.glob('DrivingStereo images/left/*.jpg')
left_images.sort()
right_images = glob.glob('DrivingStereo images/right/*.jpg')
right_images.sort()
depth_images = glob.glob('DrivingStereo images/depth/*.png')
depth_images.sort()
# Select model type
model_type = ModelType.middlebury
# model_type = ModelType.flyingthings
# model_type = ModelType.eth3d
if model_type == ModelType.middlebury:
model_path = "models/middlebury_d400.pb"
elif model_type == ModelType.flyingthings:
model_path = "models/flyingthings_finalpass_xl.pb"
elif model_type == ModelType.eth3d:
model_path = "models/eth3d.pb"
camera_config = CameraConfig(0.546, 1000)
max_distance = 50
# Initialize model
hitnet_depth = HitNet(model_path, model_type, camera_config)
cv2.namedWindow("Estimated depth", cv2.WINDOW_NORMAL)
for left_path, right_path, depth_path in zip(left_images[1500:1700:2], right_images[1500:1700:2], depth_images[1500:1700:2]):
# Read frame from the video
left_img = cv2.imread(left_path)
right_img = cv2.imread(right_path)
depth_img = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED).astype(np.float32)/256
# Estimate the depth
disparity_map = hitnet_depth(left_img, right_img)
depth_map = hitnet_depth.get_depth()
color_disparity = draw_disparity(disparity_map)
color_depth = draw_depth(depth_map, max_distance)
color_real_depth = draw_depth(depth_img, max_distance)
cobined_image = np.hstack((left_img,color_real_depth, color_depth))
out.write(cobined_image)
cv2.imshow("Estimated depth", cobined_image)
# Press key q to stop
if cv2.waitKey(1) == ord('q'):
break
out.release()
cv2.destroyAllWindows()
================================================
FILE: hitnet/__init__.py
================================================
from hitnet.hitnet import HitNet
from hitnet.utils_hitnet import *
================================================
FILE: hitnet/hitnet.py
================================================
import tensorflow as tf
import numpy as np
import time
import cv2
from hitnet.utils_hitnet import *
drivingStereo_config = CameraConfig(0.546, 1000)
class HitNet():
def __init__(self, model_path, model_type=ModelType.eth3d, camera_config=drivingStereo_config):
self.fps = 0
self.timeLastPrediction = time.time()
self.frameCounter = 0
self.camera_config = camera_config
# Initialize model
self.model = self.initialize_model(model_path, model_type)
def __call__(self, left_img, right_img):
return self.estimate_disparity(left_img, right_img)
def initialize_model(self, model_path, model_type):
self.model_type = model_type
with tf.io.gfile.GFile(model_path, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
loaded = graph_def.ParseFromString(f.read())
# Wrap frozen graph to ConcreteFunctions
if self.model_type == ModelType.flyingthings:
model = wrap_frozen_graph(graph_def=graph_def,
inputs="input:0",
outputs=["reference_output_disparity:0","secondary_output_disparity:0"])
else:
model = wrap_frozen_graph(graph_def=graph_def,
inputs="input:0",
outputs="reference_output_disparity:0")
return model
def estimate_disparity(self, left_img, right_img):
input_tensor = self.prepare_input(left_img, right_img)
# Perform inference on the image
if self.model_type == ModelType.flyingthings:
left_disparity, right_disparity = self.inference(input_tensor)
self.disparity_map = left_disparity
else:
self.disparity_map = self.inference(input_tensor)
return self.disparity_map
def get_depth(self):
return self.camera_config.f*self.camera_config.baseline/self.disparity_map
def prepare_input(self, left_img, right_img):
if (self.model_type == ModelType.eth3d):
# Shape (1, None, None, 2)
left_img = cv2.cvtColor(left_img, cv2.COLOR_BGR2GRAY)
right_img = cv2.cvtColor(right_img, cv2.COLOR_BGR2GRAY)
left_img = np.expand_dims(left_img,2)
right_img = np.expand_dims(right_img,2)
combined_img = np.concatenate((left_img, right_img), axis=-1) / 255.0
else:
# Shape (1, None, None, 6)
left_img = cv2.cvtColor(left_img, cv2.COLOR_BGR2RGB)
right_img = cv2.cvtColor(right_img, cv2.COLOR_BGR2RGB)
combined_img = np.concatenate((left_img, right_img), axis=-1) / 255.0
return tf.convert_to_tensor(np.expand_dims(combined_img, 0), dtype=tf.float32)
def inference(self, input_tensor):
output = self.model(input_tensor)
return np.squeeze(output)
================================================
FILE: hitnet/utils_hitnet.py
================================================
from enum import Enum
import tensorflow as tf
import numpy as np
import cv2
import urllib
from dataclasses import dataclass
class ModelType(Enum):
eth3d = 0
middlebury = 1
flyingthings = 2
@dataclass
class CameraConfig:
baseline: float
f: float
def wrap_frozen_graph(graph_def, inputs, outputs):
def _imports_graph_def():
tf.compat.v1.import_graph_def(graph_def, name="")
wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
import_graph = wrapped_import.graph
return wrapped_import.prune(
tf.nest.map_structure(import_graph.as_graph_element, inputs),
tf.nest.map_structure(import_graph.as_graph_element, outputs))
def draw_disparity(disparity_map):
disparity_map = disparity_map.astype(np.uint8)
norm_disparity_map = (255*((disparity_map-np.min(disparity_map))/(np.max(disparity_map) - np.min(disparity_map))))
return cv2.applyColorMap(cv2.convertScaleAbs(norm_disparity_map,1), cv2.COLORMAP_MAGMA)
def draw_depth(depth_map, max_dist):
norm_depth_map = 255*(1-depth_map/max_dist)
norm_depth_map[norm_depth_map < 0] =0
norm_depth_map[depth_map == 0] =0
return cv2.applyColorMap(cv2.convertScaleAbs(norm_depth_map,1), cv2.COLORMAP_MAGMA)
def load_img(url):
req = urllib.request.urlopen(url)
arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
return cv2.imdecode(arr, -1) # 'Load it as it is'
================================================
FILE: imageDepthEstimation.py
================================================
import cv2
import tensorflow as tf
import numpy as np
from hitnet import HitNet, ModelType, draw_disparity, draw_depth, CameraConfig, load_img
# Select model type
# model_type = ModelType.middlebury
# model_type = ModelType.flyingthings
model_type = ModelType.eth3d
if model_type == ModelType.middlebury:
model_path = "models/middlebury_d400.pb"
elif model_type == ModelType.flyingthings:
model_path = "models/flyingthings_finalpass_xl.pb"
elif model_type == ModelType.eth3d:
model_path = "models/eth3d.pb"
# Initialize model
hitnet_depth = HitNet(model_path, model_type)
# Load images
left_img = load_img("https://vision.middlebury.edu/stereo/data/scenes2003/newdata/cones/im2.png")
right_img = load_img("https://vision.middlebury.edu/stereo/data/scenes2003/newdata/cones/im6.png")
# Estimate the depth
disparity_map = hitnet_depth(left_img, right_img)
color_disparity = draw_disparity(disparity_map)
cobined_image = np.hstack((left_img, right_img, color_disparity))
cv2.namedWindow("Estimated disparity", cv2.WINDOW_NORMAL)
cv2.imshow("Estimated disparity", cobined_image)
cv2.waitKey(0)
cv2.imwrite("out.jpg", cobined_image)
cv2.destroyAllWindows()
================================================
FILE: models/.gitkeep
================================================
================================================
FILE: requirements.txt
================================================
tensorflow>=2.6
opencv-python
numpy
imread-from-url
================================================
FILE: videoDepthEstimation.py
================================================
import cv2
import pafy
import tensorflow as tf
import numpy as np
from hitnet import HitNet, ModelType, draw_disparity, draw_depth, CameraConfig
# Initialize video
# cap = cv2.VideoCapture("video.mp4")
videoUrl = 'https://youtu.be/Yui48w71SG0'
videoPafy = pafy.new(videoUrl)
print(videoPafy.streams)
cap = cv2.VideoCapture(videoPafy.getbestvideo().url)
# Select model type
# model_type = ModelType.middlebury
# model_type = ModelType.flyingthings
model_type = ModelType.eth3d
if model_type == ModelType.middlebury:
model_path = "models/middlebury_d400.pb"
elif model_type == ModelType.flyingthings:
model_path = "models/flyingthings_finalpass_xl.pb"
elif model_type == ModelType.eth3d:
model_path = "models/eth3d.pb"
# Store baseline (m) and focal length (pixel)
camera_config = CameraConfig(0.1, 320)
max_distance = 5
# Initialize model
hitnet_depth = HitNet(model_path, model_type, camera_config)
cv2.namedWindow("Estimated depth", cv2.WINDOW_NORMAL)
while cap.isOpened():
try:
# Read frame from the video
ret, frame = cap.read()
if not ret:
break
except:
continue
# Extract the left and right images
left_img = frame[:,:frame.shape[1]//3]
right_img = frame[:,frame.shape[1]//3:frame.shape[1]*2//3]
color_real_depth = frame[:,frame.shape[1]*2//3:]
# Estimate the depth
disparity_map = hitnet_depth(left_img, right_img)
depth_map = hitnet_depth.get_depth()
color_disparity = draw_disparity(disparity_map)
color_depth = draw_depth(depth_map, max_distance)
cobined_image = np.hstack((left_img,color_real_depth, color_depth))
cv2.imshow("Estimated depth", cobined_image)
# Press key q to stop
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
gitextract_i5i3wbqa/ ├── .gitignore ├── LICENSE ├── README.md ├── drivingStereoTest.py ├── hitnet/ │ ├── __init__.py │ ├── hitnet.py │ └── utils_hitnet.py ├── imageDepthEstimation.py ├── models/ │ └── .gitkeep ├── requirements.txt └── videoDepthEstimation.py
SYMBOL INDEX (14 symbols across 2 files)
FILE: hitnet/hitnet.py
class HitNet (line 10) | class HitNet():
method __init__ (line 12) | def __init__(self, model_path, model_type=ModelType.eth3d, camera_conf...
method __call__ (line 22) | def __call__(self, left_img, right_img):
method initialize_model (line 26) | def initialize_model(self, model_path, model_type):
method estimate_disparity (line 47) | def estimate_disparity(self, left_img, right_img):
method get_depth (line 60) | def get_depth(self):
method prepare_input (line 63) | def prepare_input(self, left_img, right_img):
method inference (line 83) | def inference(self, input_tensor):
FILE: hitnet/utils_hitnet.py
class ModelType (line 8) | class ModelType(Enum):
class CameraConfig (line 14) | class CameraConfig:
function wrap_frozen_graph (line 18) | def wrap_frozen_graph(graph_def, inputs, outputs):
function draw_disparity (line 27) | def draw_disparity(disparity_map):
function draw_depth (line 33) | def draw_depth(depth_map, max_dist):
function load_img (line 41) | def load_img(url):
Condensed preview — 11 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (15K chars).
[
{
"path": ".gitignore",
"chars": 1799,
"preview": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packagi"
},
{
"path": "LICENSE",
"chars": 1069,
"preview": "MIT License\n\nCopyright (c) 2021 Ibai Gorordo\n\nPermission is hereby granted, free of charge, to any person obtaining a co"
},
{
"path": "README.md",
"chars": 1743,
"preview": "# HITNET-Stereo-Depth-estimation\nPython scripts for performing stereo depth estimation using the [HITNET Tensorflow mode"
},
{
"path": "drivingStereoTest.py",
"chars": 1896,
"preview": "import cv2\nimport pafy\nimport tensorflow as tf\nimport numpy as np\nimport glob\nfrom hitnet import HitNet, ModelType, draw"
},
{
"path": "hitnet/__init__.py",
"chars": 66,
"preview": "from hitnet.hitnet import HitNet\nfrom hitnet.utils_hitnet import *"
},
{
"path": "hitnet/hitnet.py",
"chars": 2503,
"preview": "import tensorflow as tf\nimport numpy as np\nimport time\nimport cv2\nfrom hitnet.utils_hitnet import *\n\n\ndrivingStereo_conf"
},
{
"path": "hitnet/utils_hitnet.py",
"chars": 1354,
"preview": "from enum import Enum\nimport tensorflow as tf\nimport numpy as np\nimport cv2\nimport urllib\nfrom dataclasses import datacl"
},
{
"path": "imageDepthEstimation.py",
"chars": 1166,
"preview": "import cv2\nimport tensorflow as tf\nimport numpy as np\n\nfrom hitnet import HitNet, ModelType, draw_disparity, draw_depth,"
},
{
"path": "models/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "requirements.txt",
"chars": 51,
"preview": "tensorflow>=2.6\nopencv-python\nnumpy\nimread-from-url"
},
{
"path": "videoDepthEstimation.py",
"chars": 1715,
"preview": "import cv2\nimport pafy\nimport tensorflow as tf\nimport numpy as np\n\nfrom hitnet import HitNet, ModelType, draw_disparity,"
}
]
About this extraction
This page contains the full source code of the ibaiGorordo/HITNET-Stereo-Depth-estimation GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 11 files (13.0 KB), approximately 3.9k tokens, and a symbol index with 14 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.