Skip to content

Commit b3ef334

Browse files
DaniParrcbrxyz
authored andcommitted
first commit
0 parents  commit b3ef334

File tree

6 files changed

+407
-0
lines changed

6 files changed

+407
-0
lines changed

layers/rgbChannelExtraction.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import cv2
2+
import numpy as np
3+
4+
def preprocess_frame(self, frame):
5+
"""
6+
Extracts the red channel from the frame and magnifies its values.
7+
"""
8+
original_frame = frame
9+
inverted_image = cv2.bitwise_not(original_frame)
10+
11+
original_array = np.array(original_frame, dtype=np.int64)
12+
inv_original_array = np.array(inverted_image, dtype=np.int64)
13+
14+
red_channel = original_array[:, :, 2]
15+
result_image = original_array.copy()
16+
17+
# MIN MAX NORM
18+
result_image[:, :, 0] = (
19+
(result_image[:, :, 0] - np.min(result_image[:, :, 0]))
20+
/ (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0]))
21+
* 255
22+
)
23+
24+
red_channel = result_image[:, :, 2]
25+
26+
result_image[:, :, 0] = red_channel * (
27+
red_channel + inv_original_array[:, :, 0]
28+
)
29+
result_image[:, :, 1] = 0
30+
result_image[:, :, 2] = 0
31+
32+
# MIN MAX NORM
33+
result_image[:, :, 0] = (
34+
(result_image[:, :, 0] - np.min(result_image[:, :, 0]))
35+
/ (np.max(result_image[:, :, 0]) - np.min(result_image[:, :, 0]))
36+
* 255
37+
)
38+
39+
result_image = np.clip(result_image, 0, 255)
40+
result_image = result_image.astype(np.uint8)
41+
42+
self.image_pub_pre.publish(np.array(result_image))
43+
44+
return result_image

ml/converter/ptToTflite.py

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import torch
2+
import onnx
3+
from onnx import helper
4+
import tensorflow as tf
5+
from torchvision import transforms
6+
import onnx_tf
7+
from ..yolov7.models.experimental import attempt_load
8+
from PIL import Image
9+
import os
10+
import shutil
11+
12+
def convert_pt_to_tflite(weights_path, output_path, sample_image_path):
13+
"""
14+
Input:
15+
weights_path : string (e.g. path/to/weights.pt)
16+
output_path : string (e.g. path/to/tflite/weights.tflite)
17+
sample_image_path : string (e.g. path/to/sample/image.png)
18+
Output:
19+
A .tflite weights file that can be processed.
20+
"""
21+
# Load the PyTorch ResNet50 model
22+
pytorch_model = attempt_load(weights_path)
23+
pytorch_model.eval()
24+
25+
# Export the PyTorch model to ONNX format
26+
image_path = sample_image_path
27+
input_data = Image.open(image_path).convert('RGB')
28+
input_data = input_data.resize((960, 608))
29+
img_transform = transforms.Compose([
30+
transforms.ToTensor()
31+
])
32+
33+
dummy_input = img_transform(input_data).unsqueeze(0)
34+
35+
# print(dummy_input)
36+
onnx_model_path = 'temp.onnx'
37+
torch.onnx.export(pytorch_model, dummy_input, onnx_model_path, verbose=False, opset_version=12)
38+
39+
# Load the ONNX model
40+
onnx_model = onnx.load(onnx_model_path)
41+
42+
# Define a mapping from old names to new names
43+
name_map = {"input.1": "input_1"}
44+
45+
# Initialize a list to hold the new inputs
46+
new_inputs = []
47+
48+
# Iterate over the inputs and change their names if needed
49+
for inp in onnx_model.graph.input:
50+
if inp.name in name_map:
51+
# Create a new ValueInfoProto with the new name
52+
new_inp = helper.make_tensor_value_info(name_map[inp.name],
53+
inp.type.tensor_type.elem_type,
54+
[dim.dim_value for dim in inp.type.tensor_type.shape.dim])
55+
new_inputs.append(new_inp)
56+
else:
57+
new_inputs.append(inp)
58+
59+
# Clear the old inputs and add the new ones
60+
onnx_model.graph.ClearField("input")
61+
onnx_model.graph.input.extend(new_inputs)
62+
63+
# Go through all nodes in the model and replace the old input name with the new one
64+
for node in onnx_model.graph.node:
65+
for i, input_name in enumerate(node.input):
66+
if input_name in name_map:
67+
node.input[i] = name_map[input_name]
68+
69+
# Remove ONNX model file
70+
try:
71+
os.remove(onnx_model_path)
72+
except OSError as e:
73+
print(f"Could not delete {onnx_model_path}\nERROR:{e}")
74+
75+
# Convert the ONNX model to TensorFlow format
76+
tf_model_path = 'temp.pb'
77+
tf_rep = onnx_tf.backend.prepare(onnx_model)
78+
tf_rep.export_graph(tf_model_path)
79+
80+
# Convert the TensorFlow model to TensorFlow Lite format
81+
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(tf_model_path)
82+
tflite_model = converter.convert()
83+
84+
try:
85+
shutil.rmtree(tf_model_path)
86+
except OSError as e:
87+
print(f"Could not delete {tf_model_path}\nERROR:{e}")
88+
89+
# Save the TensorFlow Lite model to a file
90+
with open(output_path, 'wb') as f:
91+
f.write(tflite_model)

ml/weights/README.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
##Weights will be placed here

ml/yolov7

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Subproject commit a207844b1ce82d204ab36d87d496728d3d2348e7

ros/rosviz/CMakeLists.txt

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
cmake_minimum_required(VERSION 3.0.2)
2+
project(rosviz)
3+
4+
## Compile as C++11, supported in ROS Kinetic and newer
5+
# add_compile_options(-std=c++11)
6+
7+
## Find catkin macros and libraries
8+
## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz)
9+
## is used, also find other catkin packages
10+
find_package(catkin REQUIRED COMPONENTS
11+
roscpp
12+
rospy
13+
)
14+
15+
## System dependencies are found with CMake's conventions
16+
# find_package(Boost REQUIRED COMPONENTS system)
17+
18+
19+
## Uncomment this if the package has a setup.py. This macro ensures
20+
## modules and global scripts declared therein get installed
21+
## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html
22+
# catkin_python_setup()
23+
24+
################################################
25+
## Declare ROS messages, services and actions ##
26+
################################################
27+
28+
## To declare and build messages, services or actions from within this
29+
## package, follow these steps:
30+
## * Let MSG_DEP_SET be the set of packages whose message types you use in
31+
## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...).
32+
## * In the file package.xml:
33+
## * add a build_depend tag for "message_generation"
34+
## * add a build_depend and a exec_depend tag for each package in MSG_DEP_SET
35+
## * If MSG_DEP_SET isn't empty the following dependency has been pulled in
36+
## but can be declared for certainty nonetheless:
37+
## * add a exec_depend tag for "message_runtime"
38+
## * In this file (CMakeLists.txt):
39+
## * add "message_generation" and every package in MSG_DEP_SET to
40+
## find_package(catkin REQUIRED COMPONENTS ...)
41+
## * add "message_runtime" and every package in MSG_DEP_SET to
42+
## catkin_package(CATKIN_DEPENDS ...)
43+
## * uncomment the add_*_files sections below as needed
44+
## and list every .msg/.srv/.action file to be processed
45+
## * uncomment the generate_messages entry below
46+
## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...)
47+
48+
## Generate messages in the 'msg' folder
49+
# add_message_files(
50+
# FILES
51+
# Message1.msg
52+
# Message2.msg
53+
# )
54+
55+
## Generate services in the 'srv' folder
56+
# add_service_files(
57+
# FILES
58+
# Service1.srv
59+
# Service2.srv
60+
# )
61+
62+
## Generate actions in the 'action' folder
63+
# add_action_files(
64+
# FILES
65+
# Action1.action
66+
# Action2.action
67+
# )
68+
69+
## Generate added messages and services with any dependencies listed here
70+
# generate_messages(
71+
# DEPENDENCIES
72+
# std_msgs # Or other packages containing msgs
73+
# )
74+
75+
################################################
76+
## Declare ROS dynamic reconfigure parameters ##
77+
################################################
78+
79+
## To declare and build dynamic reconfigure parameters within this
80+
## package, follow these steps:
81+
## * In the file package.xml:
82+
## * add a build_depend and a exec_depend tag for "dynamic_reconfigure"
83+
## * In this file (CMakeLists.txt):
84+
## * add "dynamic_reconfigure" to
85+
## find_package(catkin REQUIRED COMPONENTS ...)
86+
## * uncomment the "generate_dynamic_reconfigure_options" section below
87+
## and list every .cfg file to be processed
88+
89+
## Generate dynamic reconfigure parameters in the 'cfg' folder
90+
# generate_dynamic_reconfigure_options(
91+
# cfg/DynReconf1.cfg
92+
# cfg/DynReconf2.cfg
93+
# )
94+
95+
###################################
96+
## catkin specific configuration ##
97+
###################################
98+
## The catkin_package macro generates cmake config files for your package
99+
## Declare things to be passed to dependent projects
100+
## INCLUDE_DIRS: uncomment this if your package contains header files
101+
## LIBRARIES: libraries you create in this project that dependent projects also need
102+
## CATKIN_DEPENDS: catkin_packages dependent projects also need
103+
## DEPENDS: system dependencies of this project that dependent projects also need
104+
catkin_package(
105+
# INCLUDE_DIRS include
106+
# LIBRARIES rosviz
107+
# CATKIN_DEPENDS roscpp rospy
108+
# DEPENDS system_lib
109+
)
110+
111+
###########
112+
## Build ##
113+
###########
114+
115+
## Specify additional locations of header files
116+
## Your package locations should be listed before other locations
117+
include_directories(
118+
# include
119+
${catkin_INCLUDE_DIRS}
120+
)
121+
122+
## Declare a C++ library
123+
# add_library(${PROJECT_NAME}
124+
# src/${PROJECT_NAME}/rosviz.cpp
125+
# )
126+
127+
## Add cmake target dependencies of the library
128+
## as an example, code may need to be generated before libraries
129+
## either from message generation or dynamic reconfigure
130+
# add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS})
131+
132+
## Declare a C++ executable
133+
## With catkin_make all packages are built within a single CMake context
134+
## The recommended prefix ensures that target names across packages don't collide
135+
# add_executable(${PROJECT_NAME}_node src/rosviz_node.cpp)
136+
137+
## Rename C++ executable without prefix
138+
## The above recommended prefix causes long target names, the following renames the
139+
## target back to the shorter version for ease of user use
140+
## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node"
141+
# set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "")
142+
143+
## Add cmake target dependencies of the executable
144+
## same as for the library above
145+
# add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS})
146+
147+
## Specify libraries to link a library or executable target against
148+
# target_link_libraries(${PROJECT_NAME}_node
149+
# ${catkin_LIBRARIES}
150+
# )
151+
152+
#############
153+
## Install ##
154+
#############
155+
156+
# all install targets should use catkin DESTINATION variables
157+
# See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html
158+
159+
## Mark executable scripts (Python etc.) for installation
160+
## in contrast to setup.py, you can choose the destination
161+
# catkin_install_python(PROGRAMS
162+
# scripts/my_python_script
163+
# DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
164+
# )
165+
166+
## Mark executables for installation
167+
## See http://docs.ros.org/melodic/api/catkin/html/howto/format1/building_executables.html
168+
# install(TARGETS ${PROJECT_NAME}_node
169+
# RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
170+
# )
171+
172+
## Mark libraries for installation
173+
## See http://docs.ros.org/melodic/api/catkin/html/howto/format1/building_libraries.html
174+
# install(TARGETS ${PROJECT_NAME}
175+
# ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION}
176+
# LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION}
177+
# RUNTIME DESTINATION ${CATKIN_GLOBAL_BIN_DESTINATION}
178+
# )
179+
180+
## Mark cpp header files for installation
181+
# install(DIRECTORY include/${PROJECT_NAME}/
182+
# DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION}
183+
# FILES_MATCHING PATTERN "*.h"
184+
# PATTERN ".svn" EXCLUDE
185+
# )
186+
187+
## Mark other files for installation (e.g. launch and bag files, etc.)
188+
# install(FILES
189+
# # myfile1
190+
# # myfile2
191+
# DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION}
192+
# )
193+
194+
#############
195+
## Testing ##
196+
#############
197+
198+
## Add gtest based cpp test target and link libraries
199+
# catkin_add_gtest(${PROJECT_NAME}-test test/test_rosviz.cpp)
200+
# if(TARGET ${PROJECT_NAME}-test)
201+
# target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME})
202+
# endif()
203+
204+
## Add folders to be run by python nosetests
205+
# catkin_add_nosetests(test)

0 commit comments

Comments
 (0)