added initial version of yolo standalone controller

This commit is contained in:
Nickiel12 2024-04-13 20:28:46 -07:00
parent 3b7bc8cc8f
commit 22ca17977e
2 changed files with 82 additions and 1 deletions

81
Client/yolo_client.py Normal file
View file

@ -0,0 +1,81 @@
import cv2
from websocket import create_connection
import argparse
import math
from ultralytics import YOLO
# object classes
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
"baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
"carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
"diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
"teddy bear", "hair drier", "toothbrush"
]
def main(websocket_addr):
# Connect to websocket server
# ws = create_connection(websocket_addr)
# Load YOLO model
model = YOLO('yolov8s.pt') # Load an official Detect model
# Open webcam
cap = cv2.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)
frame_count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Perform object detection
results = model.predict(frame)
for r in results:
boxes = r.boxes
for box in boxes:
if box.cls[0].item() == 0:
# bounding box
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values
# put box in cam
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)
# class name
cls = int(box.cls[0])
# object details
org = [x1, y1]
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 1
color = (255, 0, 0)
thickness = 2
cv2.putText(frame, classNames[cls], org, font, fontScale, color, thickness)
print(f"{frame_count}, {x1}:{y1}, {x2}:{y2}")
frame_count += 1
cv2.imshow('Webcam', frame)
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Detect people from webcam feed and send data to websocket server')
parser.add_argument('websocket_addr', type=str, help='Websocket server address')
args = parser.parse_args()
main(args.websocket_addr)

@ -1 +1 @@
Subproject commit 852e23dc7c92f80bf59f679803fb6344942bc138
Subproject commit 4530cf45a81cbd605d3fae4a371847239b47cb6e