代码拉取完成，页面将自动刷新
# encoding: utf-8
__author__ = 'Gui'
'''
@Time:2022年6月5日 18:15
@Auth:19级机器人工程 (按拼音排序) 桂源泽 苏琦 颜欢
@File:fira_Collect_offline.py
@IDE:fira项目数据收集程序
@Software: PyCharm
'''
import rospy
import numpy as np
import cv2
from robot import Robot
from geometry_msgs.msg import Twist
import time
import os
rate = 10 # 向ros发布命令频率
v = 0.1 # 前进距离
angular = 0.1 # 旋转角度
threshold_r = 50 # 红二值化阈值 110
threshold_b = 40 # 蓝二值化阈值 110
threshold_g = 30 # 绿二值化阈值 110
class CollectTrainingData(object):
"""
input:
commands and video
k: cmd: control:
1 0 0 0 0 w:前进 u i o 左前 前进 右前
0 1 0 0 0 a:左前 j k l 左转 停止 右转
0 0 1 0 0 d:右前 m , . 左后 后退 右后
0 0 0 1 0 s:停止
0 0 0 0 1 t:冲刺
output:
带有标签的灰度图像集，标签（0, 1, 2, 3 , 4）分别代表（前进， 左转，右转，停止, 冲刺）
每种标签数量上限1000张，像素为H*W = 480×180
"""
def __init__(self):
self.raw_height = 480 # 原始视频高度
self.raw_width = 640 # 原始视频宽度
self.video_width = 480 # 截取图像宽度
self.video_width_save = 480 # 保存图像宽度*3
self.video_height = 180 # 截取图像高度
self.channels = 1 # 通道数量 1
self.NUM = 5 # 分类数量：0, 1, 2, 3, 4
self.range = 300 # 每个分类的图片数
self.data_path = "dataset"
self.saved_file_name = 'labeled_img_data_' + str(int(time.time()))
#控制底盘
self.robot = Robot()
#self.mv = Movement()
# 发布话题相关参数
self.rate_run = rospy.Rate(rate)
self.twist = Twist()
# 创建标签列表
self.k = np.zeros((self.NUM, self.NUM), 'float')
for i in range(self.NUM):
self.k[i, i] = 1
self.collect_image()#开始收集图片
def collect_image(self):
# 初始化数数
total_images_collected = 0
num_list = [0, 0, 0, 0, 0, 0, 0] # 当前各标签存储图片数量
#cap = cv2.VideoCapture(0) # 开启摄像头
images = np.zeros((1, self.video_height * self.video_width), dtype=float)
labels = np.zeros((1, self.NUM), dtype=float)
# Send an action to begin program.
# # 示意准备完毕
print("prepar to continue.............")
while True:
frame = self.robot.get_image()
cv2.imshow("orgin", frame)
resized_height = int(self.video_width * 0.75)
# 计算缩放比例
frame = cv2.resize(frame, (self.video_width, resized_height))
# frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # 灰度化
# frame = cv2.GaussianBlur(frame, (5, 5), 0) # 高斯模糊
frame = cv2.medianBlur(frame, 3) # 中值滤波
# slice the lower part of a frame
res = frame[90:270, :]#剪切画面
cv2.imshow("review", res)
[aisle_b, aisle_g, aisle_r] = cv2.split(res)
# aisle_b = cv2.adaptiveThreshold(aisle_b,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,3,5)
_, aisle_g = cv2.threshold(aisle_g, threshold_g, 255, cv2.THRESH_BINARY) # 统二值化
# aisle_b = cv2.Canny(aisle_b,40,140)
aisle_b = cv2.adaptiveThreshold(aisle_b,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,3,5)
# aisle_r = cv2.adaptiveThreshold(aisle_r,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,3,5)
_, aisle_r = cv2.threshold(aisle_r, threshold_r, 255, cv2.THRESH_BINARY) # 统二值化
res_tmp = cv2.vconcat([aisle_b,aisle_g,aisle_r])
res = cv2.resize(res_tmp, (self.video_width, self.video_height), interpolation=cv2.INTER_AREA)
# _, res = cv2.threshold(res, threshold_yuzhi, 255, cv2.THRESH_BINARY) # 统一二值化
cv2.imshow("review2", res)
command = cv2.waitKey(100) & 0xFF # 等待输入按键取后八位
if command == 255:
self.twist.linear.x = 0
self.twist.angular.z = 0
continue
elif command == ord('q'):
print("..............quiting.............")
break
# forward -- 0
elif command == ord('w'):
if num_list[0] < self.range: #小于最大图片数量
num_list[0] += 1
total_images_collected += 1
# self.twist.linear.x = v
# self.twist.angular.z = 0
res = np.reshape(res, [1, -1])
images=np.vstack((images, res)) # 将当前画面按垂直方向堆叠
labels = np.vstack((labels, self.k[0])) #将当前标签按垂直方向堆叠
print("Forward image collect: ", num_list[0])
else:
print("list full!!!")
# self.twist.linear.x = v
# self.twist.angular.z = 0
continue
# forward-left -- 1
elif command == ord('a'):
if num_list[1] < self.range:
num_list[1] += 1
total_images_collected += 1
# self.twist.linear.x = v
# self.twist.angular.z = angular
res = np.reshape(res, [1, -1])
images=np.vstack((images, res))
labels = np.vstack((labels, self.k[1]))
print("Left image collect: ", num_list[1])
else:
print("list full!!!")
# self.twist.linear.x = 0
# self.twist.angular.z = 0
continue
# forward-right -- 2
elif command == ord('d'):
if num_list[2] < self.range:
num_list[2] += 1
total_images_collected += 1
# self.twist.linear.x = v
# self.twist.angular.z = -angular
res = np.reshape(res, [1, -1])
images=np.vstack((images, res))
labels = np.vstack((labels, self.k[2]))
print("Right image collect: ", num_list[2])
else:
print("list full!!!")
# self.twist.linear.x = 0
# self.twist.angular.z = 0
continue
# stop-sign -- 3
elif command == ord('s'):
if num_list[3] < self.range:
num_list[3] += 1
total_images_collected += 1
# self.twist.linear.x = 0
# self.twist.angular.z = 0
res = np.reshape(res, [1, -1])
images=np.vstack((images, res))
labels = np.vstack((labels, self.k[3]))
print("Stop image collect: ", num_list[3])
else:
print("list full!!!")
# self.twist.linear.x = 0
# self.twist.angular.z = 0
continue
# road banner front
elif command == ord('t'):
if num_list[4] < self.range:
num_list[4] += 1
total_images_collected += 1
# self.twist.linear.x = v
# self.twist.angular.z = 0
res = np.reshape(res, [1, -1])
images = np.vstack((images, res))
labels = np.vstack((labels, self.k[4]))
print("rush image collect: ", num_list[4])
else:
print("list full!!!")
# self.twist.linear.x = 0
# self.twist.angular.z = 0
continue
## 控制小车运动辅助拍照
# forward
elif command == ord('i'):
self.twist.linear.x = v
self.twist.angular.z = 0
print("前进")
# forward-left
elif command == ord('u'):
self.twist.linear.x = v
self.twist.angular.z = angular
print("左转")
# forward-right
elif command == ord('o'):
self.twist.linear.x = v
self.twist.angular.z = -angular
print("右转")
# back
elif command == ord(','):
self.twist.linear.x = -v
self.twist.angular.z = 0
print("后退")
# self-left
elif command == ord('j'):
self.twist.linear.x = 0.1*v
self.twist.angular.z = angular
print("左自转")
# self-right
elif command == ord('l'):
self.twist.linear.x = 0.1*v
self.twist.angular.z = -angular
print("右自转")
# stop
elif command == ord('k'):
self.twist.linear.x = 0
self.twist.angular.z = 0
print("停止")
elif command == ord('m'):
self.twist.linear.x = -v
self.twist.angular.z = -angular
print("左后")
# forward-left
elif command == ord('.'):
self.twist.linear.x = -v
self.twist.angular.z = angular
print("右后")
elif num_list[0] == self.range and num_list[1] == self.range and num_list[2] == self.range and num_list[3] == self.range and num_list[4] == self.range and num_list[5] == self.range and num_list[6] == self.range:
#elf.mv.wave_hands()
print("---------All list full!!!----------")
break
# exit(0)
# 向机器人底盘发布数据
self.robot.publish_twist(self.twist)
self.rate_run.sleep()
img = images[1:, :]
lbl = labels[1:, :]
print("image shape:", img.shape)
print("label shape:", lbl.shape)
print("\n")
print("forward images num:", num_list[0])
print("forward left images num:", num_list[1])
print("forward right images num:", num_list[2])
print("stop sign images num:", num_list[3])
# 保存数据
if not os.path.exists(self.data_path):
os.mkdir(self.data_path)
try:
# 保存文件
print(".................saving file...............")
name = self.data_path + '/' + self.saved_file_name + '.npz'
np.savez(name, train=img, train_labels=lbl, num_list=num_list)
print("saving file:", name)
except IOError as e:
print(e)
cv2.destroyAllWindows()
if __name__ == '__main__':
try:
CollectTrainingData()
except KeyboardInterrupt:
pass
此处可能存在不合适展示的内容，页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容，可点击提交进行申诉，我们将尽快为您处理。