Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import torch | |
| import numpy as np | |
| import time | |
| from threading import Thread | |
| import queue | |
| # Load the YOLO model on the CPU | |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) | |
| # Global variables for camera control | |
| camera_queue = queue.Queue(maxsize=1) | |
| stop_camera = False | |
| def move_robot(command): | |
| """ | |
| Simulate sending a command to the robot to move. | |
| Replace this with actual communication code to send commands to the robot. | |
| """ | |
| commands = { | |
| "forward": "Robot moving forward", | |
| "backward": "Robot moving backward", | |
| "left": "Robot turning left", | |
| "right": "Robot turning right", | |
| "stop": "Robot stopped" | |
| } | |
| print(commands.get(command, "Invalid command")) | |
| return f"Command sent: {command}" | |
| def camera_loop(): | |
| """Background thread function to continuously capture frames""" | |
| global stop_camera | |
| cap = cv2.VideoCapture(0) | |
| if not cap.isOpened(): | |
| print("Error: Could not open camera") | |
| return | |
| while not stop_camera: | |
| ret, frame = cap.read() | |
| if ret: | |
| # Run YOLO detection | |
| results = model(frame) | |
| processed_frame = results.render()[0] | |
| # Convert BGR to RGB for display | |
| processed_frame = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB) | |
| # Update the queue with the latest frame | |
| try: | |
| if camera_queue.full(): | |
| camera_queue.get_nowait() # Remove old frame | |
| camera_queue.put_nowait(processed_frame) | |
| except queue.Full: | |
| pass | |
| time.sleep(0.1) # Prevent excessive CPU usage | |
| cap.release() | |
| def start_camera(): | |
| """Start the camera feed""" | |
| global stop_camera | |
| stop_camera = False | |
| Thread(target=camera_loop, daemon=True).start() | |
| return "Camera started" | |
| def stop_camera_feed(): | |
| """Stop the camera feed""" | |
| global stop_camera | |
| stop_camera = True | |
| empty_frame = np.zeros((480, 640, 3), dtype=np.uint8) | |
| return "Camera stopped", empty_frame | |
| def get_current_frame(): | |
| """Get the latest frame from the camera queue""" | |
| try: | |
| frame = camera_queue.get_nowait() | |
| return frame | |
| except queue.Empty: | |
| return np.zeros((480, 640, 3), dtype=np.uint8) | |
| def robot_interface(): | |
| with gr.Blocks() as interface: | |
| with gr.Row(): | |
| # Left column for controls | |
| with gr.Column(scale=1): | |
| # Control buttons | |
| forward_btn = gr.Button("Forward", variant="primary") | |
| backward_btn = gr.Button("Backward", variant="primary") | |
| left_btn = gr.Button("Left", variant="primary") | |
| right_btn = gr.Button("Right", variant="primary") | |
| stop_btn = gr.Button("Stop", variant="secondary") | |
| status_output = gr.Textbox(label="Robot Status") | |
| # Right column for camera feed | |
| with gr.Column(scale=2): | |
| # Camera feed output | |
| camera_output = gr.Image( | |
| label="Robot Camera Feed with Object Detection", | |
| type="numpy", | |
| height=480, | |
| streaming=True | |
| ) | |
| # Start/Stop buttons for camera feed | |
| start_camera_btn = gr.Button("Start Camera", variant="primary") | |
| stop_camera_btn = gr.Button("Stop Camera", variant="secondary") | |
| # Set up button click events for robot movement | |
| forward_btn.click(fn=move_robot, inputs=[gr.Textbox(value="forward", visible=False)], outputs=status_output) | |
| backward_btn.click(fn=move_robot, inputs=[gr.Textbox(value="backward", visible=False)], outputs=status_output) | |
| left_btn.click(fn=move_robot, inputs=[gr.Textbox(value="left", visible=False)], outputs=status_output) | |
| right_btn.click(fn=move_robot, inputs=[gr.Textbox(value="right", visible=False)], outputs=status_output) | |
| stop_btn.click(fn=move_robot, inputs=[gr.Textbox(value="stop", visible=False)], outputs=status_output) | |
| # Set up button click events for camera control | |
| start_camera_btn.click(fn=start_camera, inputs=None, outputs=status_output) | |
| stop_camera_btn.click(fn=stop_camera_feed, inputs=None, outputs=[status_output, camera_output]) | |
| # Stream the camera feed | |
| camera_output.stream(fn=get_current_frame) | |
| return interface | |
| # Launch the Gradio interface | |
| if __name__ == "__main__": | |
| robot_interface().launch(share=True, debug=True) | |