-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
6 changed files
with
407 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,44 @@ | ||
import cv2 | ||
import numpy as np | ||
|
||
def preprocess_frame(self, frame): | ||
""" | ||
Extracts the red channel from the frame and magnifies its values. | ||
""" | ||
original_frame = frame | ||
inverted_image = cv2.bitwise_not(original_frame) | ||
|
||
original_array = np.array(original_frame, dtype=np.int64) | ||
inv_original_array = np.array(inverted_image, dtype=np.int64) | ||
|
||
red_channel = original_array[:, :, 2] | ||
result_image = original_array.copy() | ||
|
||
# MIN MAX NORM | ||
result_image[:, :, 0] = ( | ||
(result_image[:, :, 0] - np.min(result_image[:, :, 0])) | ||
/ (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0])) | ||
* 255 | ||
) | ||
|
||
red_channel = result_image[:, :, 2] | ||
|
||
result_image[:, :, 0] = red_channel * ( | ||
red_channel + inv_original_array[:, :, 0] | ||
) | ||
result_image[:, :, 1] = 0 | ||
result_image[:, :, 2] = 0 | ||
|
||
# MIN MAX NORM | ||
result_image[:, :, 0] = ( | ||
(result_image[:, :, 0] - np.min(result_image[:, :, 0])) | ||
/ (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0])) | ||
* 255 | ||
) | ||
|
||
result_image = np.clip(result_image, 0, 255) | ||
result_image = result_image.astype(np.uint8) | ||
|
||
self.image_pub_pre.publish(np.array(result_image)) | ||
|
||
return result_image |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,91 @@ | ||
import torch | ||
import onnx | ||
from onnx import helper | ||
import tensorflow as tf | ||
from torchvision import transforms | ||
import onnx_tf | ||
from ..yolov7.models.experimental import attempt_load | ||
from PIL import Image | ||
import os | ||
import shutil | ||
|
||
def convert_pt_to_tflite(weights_path, output_path, sample_image_path): | ||
""" | ||
Input: | ||
weights_path : string (e.g. path/to/weights.pt) | ||
output_path : string (e.g. path/to/tflite/weights.tflite) | ||
sample_image_path : string (e.g. path/to/sample/image.png) | ||
Output: | ||
A .tflite weights file that can be processed. | ||
""" | ||
# Load the PyTorch ResNet50 model | ||
pytorch_model = attempt_load(weights_path) | ||
pytorch_model.eval() | ||
|
||
# Export the PyTorch model to ONNX format | ||
image_path = sample_image_path | ||
input_data = Image.open(image_path).convert('RGB') | ||
input_data = input_data.resize((960, 608)) | ||
img_transform = transforms.Compose([ | ||
transforms.ToTensor() | ||
]) | ||
|
||
dummy_input = img_transform(input_data).unsqueeze(0) | ||
|
||
# print(dummy_input) | ||
onnx_model_path = 'temp.onnx' | ||
torch.onnx.export(pytorch_model, dummy_input, onnx_model_path, verbose=False, opset_version=12) | ||
|
||
# Load the ONNX model | ||
onnx_model = onnx.load(onnx_model_path) | ||
|
||
# Define a mapping from old names to new names | ||
name_map = {"input.1": "input_1"} | ||
|
||
# Initialize a list to hold the new inputs | ||
new_inputs = [] | ||
|
||
# Iterate over the inputs and change their names if needed | ||
for inp in onnx_model.graph.input: | ||
if inp.name in name_map: | ||
# Create a new ValueInfoProto with the new name | ||
new_inp = helper.make_tensor_value_info(name_map[inp.name], | ||
inp.type.tensor_type.elem_type, | ||
[dim.dim_value for dim in inp.type.tensor_type.shape.dim]) | ||
new_inputs.append(new_inp) | ||
else: | ||
new_inputs.append(inp) | ||
|
||
# Clear the old inputs and add the new ones | ||
onnx_model.graph.ClearField("input") | ||
onnx_model.graph.input.extend(new_inputs) | ||
|
||
# Go through all nodes in the model and replace the old input name with the new one | ||
for node in onnx_model.graph.node: | ||
for i, input_name in enumerate(node.input): | ||
if input_name in name_map: | ||
node.input[i] = name_map[input_name] | ||
|
||
# Remove ONNX model file | ||
try: | ||
os.remove(onnx_model_path) | ||
except OSError as e: | ||
print(f"Could not delete {onnx_model_path}\nERROR:{e}") | ||
|
||
# Convert the ONNX model to TensorFlow format | ||
tf_model_path = 'temp.pb' | ||
tf_rep = onnx_tf.backend.prepare(onnx_model) | ||
tf_rep.export_graph(tf_model_path) | ||
|
||
# Convert the TensorFlow model to TensorFlow Lite format | ||
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(tf_model_path) | ||
tflite_model = converter.convert() | ||
|
||
try: | ||
shutil.rmtree(tf_model_path) | ||
except OSError as e: | ||
print(f"Could not delete {tf_model_path}\nERROR:{e}") | ||
|
||
# Save the TensorFlow Lite model to a file | ||
with open(output_path, 'wb') as f: | ||
f.write(tflite_model) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
##Weights will be placed here |
Submodule yolov7
added at
a20784
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,205 @@ | ||
cmake_minimum_required(VERSION 3.0.2) | ||
project(rosviz) | ||
|
||
## Compile as C++11, supported in ROS Kinetic and newer | ||
# add_compile_options(-std=c++11) | ||
|
||
## Find catkin macros and libraries | ||
## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) | ||
## is used, also find other catkin packages | ||
find_package(catkin REQUIRED COMPONENTS | ||
roscpp | ||
rospy | ||
) | ||
|
||
## System dependencies are found with CMake's conventions | ||
# find_package(Boost REQUIRED COMPONENTS system) | ||
|
||
|
||
## Uncomment this if the package has a setup.py. This macro ensures | ||
## modules and global scripts declared therein get installed | ||
## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html | ||
# catkin_python_setup() | ||
|
||
################################################ | ||
## Declare ROS messages, services and actions ## | ||
################################################ | ||
|
||
## To declare and build messages, services or actions from within this | ||
## package, follow these steps: | ||
## * Let MSG_DEP_SET be the set of packages whose message types you use in | ||
## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...). | ||
## * In the file package.xml: | ||
## * add a build_depend tag for "message_generation" | ||
## * add a build_depend and a exec_depend tag for each package in MSG_DEP_SET | ||
## * If MSG_DEP_SET isn't empty the following dependency has been pulled in | ||
## but can be declared for certainty nonetheless: | ||
## * add a exec_depend tag for "message_runtime" | ||
## * In this file (CMakeLists.txt): | ||
## * add "message_generation" and every package in MSG_DEP_SET to | ||
## find_package(catkin REQUIRED COMPONENTS ...) | ||
## * add "message_runtime" and every package in MSG_DEP_SET to | ||
## catkin_package(CATKIN_DEPENDS ...) | ||
## * uncomment the add_*_files sections below as needed | ||
## and list every .msg/.srv/.action file to be processed | ||
## * uncomment the generate_messages entry below | ||
## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...) | ||
|
||
## Generate messages in the 'msg' folder | ||
# add_message_files( | ||
# FILES | ||
# Message1.msg | ||
# Message2.msg | ||
# ) | ||
|
||
## Generate services in the 'srv' folder | ||
# add_service_files( | ||
# FILES | ||
# Service1.srv | ||
# Service2.srv | ||
# ) | ||
|
||
## Generate actions in the 'action' folder | ||
# add_action_files( | ||
# FILES | ||
# Action1.action | ||
# Action2.action | ||
# ) | ||
|
||
## Generate added messages and services with any dependencies listed here | ||
# generate_messages( | ||
# DEPENDENCIES | ||
# std_msgs # Or other packages containing msgs | ||
# ) | ||
|
||
################################################ | ||
## Declare ROS dynamic reconfigure parameters ## | ||
################################################ | ||
|
||
## To declare and build dynamic reconfigure parameters within this | ||
## package, follow these steps: | ||
## * In the file package.xml: | ||
## * add a build_depend and a exec_depend tag for "dynamic_reconfigure" | ||
## * In this file (CMakeLists.txt): | ||
## * add "dynamic_reconfigure" to | ||
## find_package(catkin REQUIRED COMPONENTS ...) | ||
## * uncomment the "generate_dynamic_reconfigure_options" section below | ||
## and list every .cfg file to be processed | ||
|
||
## Generate dynamic reconfigure parameters in the 'cfg' folder | ||
# generate_dynamic_reconfigure_options( | ||
# cfg/DynReconf1.cfg | ||
# cfg/DynReconf2.cfg | ||
# ) | ||
|
||
################################### | ||
## catkin specific configuration ## | ||
################################### | ||
## The catkin_package macro generates cmake config files for your package | ||
## Declare things to be passed to dependent projects | ||
## INCLUDE_DIRS: uncomment this if your package contains header files | ||
## LIBRARIES: libraries you create in this project that dependent projects also need | ||
## CATKIN_DEPENDS: catkin_packages dependent projects also need | ||
## DEPENDS: system dependencies of this project that dependent projects also need | ||
catkin_package( | ||
# INCLUDE_DIRS include | ||
# LIBRARIES rosviz | ||
# CATKIN_DEPENDS roscpp rospy | ||
# DEPENDS system_lib | ||
) | ||
|
||
########### | ||
## Build ## | ||
########### | ||
|
||
## Specify additional locations of header files | ||
## Your package locations should be listed before other locations | ||
include_directories( | ||
# include | ||
${catkin_INCLUDE_DIRS} | ||
) | ||
|
||
## Declare a C++ library | ||
# add_library(${PROJECT_NAME} | ||
# src/${PROJECT_NAME}/rosviz.cpp | ||
# ) | ||
|
||
## Add cmake target dependencies of the library | ||
## as an example, code may need to be generated before libraries | ||
## either from message generation or dynamic reconfigure | ||
# add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) | ||
|
||
## Declare a C++ executable | ||
## With catkin_make all packages are built within a single CMake context | ||
## The recommended prefix ensures that target names across packages don't collide | ||
# add_executable(${PROJECT_NAME}_node src/rosviz_node.cpp) | ||
|
||
## Rename C++ executable without prefix | ||
## The above recommended prefix causes long target names, the following renames the | ||
## target back to the shorter version for ease of user use | ||
## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node" | ||
# set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "") | ||
|
||
## Add cmake target dependencies of the executable | ||
## same as for the library above | ||
# add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) | ||
|
||
## Specify libraries to link a library or executable target against | ||
# target_link_libraries(${PROJECT_NAME}_node | ||
# ${catkin_LIBRARIES} | ||
# ) | ||
|
||
############# | ||
## Install ## | ||
############# | ||
|
||
# all install targets should use catkin DESTINATION variables | ||
# See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html | ||
|
||
## Mark executable scripts (Python etc.) for installation | ||
## in contrast to setup.py, you can choose the destination | ||
# catkin_install_python(PROGRAMS | ||
# scripts/my_python_script | ||
# DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} | ||
# ) | ||
|
||
## Mark executables for installation | ||
## See http://docs.ros.org/melodic/api/catkin/html/howto/format1/building_executables.html | ||
# install(TARGETS ${PROJECT_NAME}_node | ||
# RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} | ||
# ) | ||
|
||
## Mark libraries for installation | ||
## See http://docs.ros.org/melodic/api/catkin/html/howto/format1/building_libraries.html | ||
# install(TARGETS ${PROJECT_NAME} | ||
# ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} | ||
# LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} | ||
# RUNTIME DESTINATION ${CATKIN_GLOBAL_BIN_DESTINATION} | ||
# ) | ||
|
||
## Mark cpp header files for installation | ||
# install(DIRECTORY include/${PROJECT_NAME}/ | ||
# DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} | ||
# FILES_MATCHING PATTERN "*.h" | ||
# PATTERN ".svn" EXCLUDE | ||
# ) | ||
|
||
## Mark other files for installation (e.g. launch and bag files, etc.) | ||
# install(FILES | ||
# # myfile1 | ||
# # myfile2 | ||
# DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION} | ||
# ) | ||
|
||
############# | ||
## Testing ## | ||
############# | ||
|
||
## Add gtest based cpp test target and link libraries | ||
# catkin_add_gtest(${PROJECT_NAME}-test test/test_rosviz.cpp) | ||
# if(TARGET ${PROJECT_NAME}-test) | ||
# target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME}) | ||
# endif() | ||
|
||
## Add folders to be run by python nosetests | ||
# catkin_add_nosetests(test) |
Oops, something went wrong.