import cv2
import torch
import os
import time
import pandas as pd
import numpy as np
import platform
import pathlib
pause = 0.2 # 每次扫描间隔时间
yolov5_path = '/home/yahboom/yolov5'
weight_path = "/home/yahboom/Dofbot/3.ctrl_Arm/best.pt"
def save_to_csv(coordinates, filename='coordinates.csv'):
"""保存坐
标到CSV文件
coordinates - 坐标列表，np.array格式"""
df = pd.DataFrame(coordinates)
df.to_csv(filename, index=False, encoding='utf-8-sig')
def output(results):
coordinates = []
maxA = 640*480
for _, row in results.iterrows():
u = (row['xmin'] + row['xmax']) / 2
v = (row['ymin'] + row['ymax']) / 2
size = (row['xmax'] - row['xmin']) * (row['ymax'] - row['ymin'])
if row['confidence'] > 0.4 and size < maxA*0.8:
coordinates.append((u,v,row['confidence']))
filename = 'coordinates.csv'
save_to_csv(coordinates, filename)
print(f"已保存 {len(coordinates)} 组坐标至 {filename}")
def scan(yolo_path, weight_path, pause):
# Initialize the camera and the YOLOv5 model
model = torch.hub.load(yolo_path, "custom", path=weight_path, source="local") # local repo
cap = cv2.VideoCapture(0) # 0 is usually the default camera, 1 is the us camera connected
if not cap.isOpened():
print("Error: Could not open camera.")
return NotImplementedError
while True:
print("================================================================")
# Capture a frame
ret, frame = cap.read()
if not ret:
print("Error: Could not read frame.")
return NotImplementedError
# Save the frame
frame_filename = 'temp_frame.jpg'
cv2.imwrite(frame_filename, frame)
# use yolov5 to detect the object and draw the bounding box
resultimg = model(frame_filename)
result = resultimg.pandas().xyxy[0].sort_values("xmin")
output(result)
try:
os.remove('temp_frame.jpg')
except:
pass
time.sleep(pause)
if name == "main":
scan(yolov5_path, weight_path, pause)