Skip to content

Commit

Permalink
first commit
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniParr authored and cbrxyz committed Mar 7, 2024
0 parents commit b3ef334
Show file tree
Hide file tree
Showing 6 changed files with 407 additions and 0 deletions.
44 changes: 44 additions & 0 deletions layers/rgbChannelExtraction.py
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
91 changes: 91 additions & 0 deletions ml/converter/ptToTflite.py
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)
1 change: 1 addition & 0 deletions ml/weights/README.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
##Weights will be placed here
1 change: 1 addition & 0 deletions ml/yolov7
Submodule yolov7 added at a20784
205 changes: 205 additions & 0 deletions ros/rosviz/CMakeLists.txt
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)
Loading

0 comments on commit b3ef334

Please sign in to comment.