andito HF Staff commited on
Commit
9eede3b
·
verified ·
1 Parent(s): 869f80d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -17
app.py CHANGED
@@ -12,13 +12,18 @@ import cv2
12
  import gradio as gr
13
  import numpy as np
14
  from fastapi import FastAPI, WebSocket
 
15
 
16
  from reachy_mini.utils import create_head_pose
17
 
18
  robot_ws = None # type: Optional[WebSocket]
19
  robot_loop = None # type: Optional[asyncio.AbstractEventLoop]
20
  # Add after imports
21
- latest_test_frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
 
 
 
 
22
 
23
  @dataclass
24
  class Movement:
@@ -63,7 +68,6 @@ class ReachyController:
63
  self.daemon_process = None
64
  self.reachy_mini = None
65
  self.running = False
66
- self.current_frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
67
 
68
  self.movement_queue: List[Movement] = []
69
  self.is_playing = False
@@ -120,21 +124,20 @@ class ReachyController:
120
  frame = cv2.imdecode(arr, cv2.IMREAD_COLOR) # returns BGR np.ndarray
121
  if frame is None:
122
  # Decoding failed, show black frame
123
- frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
124
  else:
125
- # Convert BGR->RGB for Gradio and resize to 1080x1080
126
  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
127
- frame = cv2.resize(frame, (1080, 1080), interpolation=cv2.INTER_LINEAR)
128
  else:
129
  # Unexpected type, show black frame
130
- frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
131
  decode_time = time.perf_counter() - decode_start
132
  except Exception as e:
133
  print("[stream_frames] JPEG decode error:", e)
134
- frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
135
  decode_time = 0
136
 
137
- frame_time = time.perf_counter() - frame_start_time
138
  frame_count += 1
139
 
140
  # Calculate and print FPS every 30 frames (~1 second at 30fps)
@@ -149,7 +152,20 @@ class ReachyController:
149
 
150
  yield frame
151
  time.sleep(0.04)
 
 
 
 
 
 
 
 
 
 
 
152
 
 
 
153
  def auto_start(self):
154
  """Auto-start daemon and robot connection"""
155
  status_msgs = []
@@ -492,14 +508,22 @@ with gr.Blocks(title="Reachy Controller", theme=gr.themes.Soft()) as demo:
492
 
493
  # Right panel - Simulation view (larger and more square)
494
  with gr.Column(scale=3):
495
- sim_view = gr.Image(
496
- label="🎬 Robot Simulation",
497
- type="numpy",
498
- height=1080,
499
- width=1080,
500
- show_label=True,
501
- streaming=True
502
- )
 
 
 
 
 
 
 
 
503
 
504
 
505
  # Movement builder section below
@@ -552,7 +576,7 @@ with gr.Blocks(title="Reachy Controller", theme=gr.themes.Soft()) as demo:
552
  add_custom_btn = gr.Button("➕ Add to Queue", variant="primary")
553
 
554
  # Stream video
555
- demo.load(fn=manager.stream_frames, outputs=sim_view)
556
 
557
  # Connect events - Connection settings
558
  apply_connection_btn.click(
@@ -674,6 +698,13 @@ async def robot_socket(ws: WebSocket):
674
  except Exception as e:
675
  print("[Space] MuJoCo stream disconnected", e)
676
 
 
 
 
 
 
 
 
677
  @app.websocket("/robot")
678
  async def robot_socket(ws: WebSocket):
679
  global robot_ws, robot_loop
 
12
  import gradio as gr
13
  import numpy as np
14
  from fastapi import FastAPI, WebSocket
15
+ from fastapi.responses import StreamingResponse
16
 
17
  from reachy_mini.utils import create_head_pose
18
 
19
  robot_ws = None # type: Optional[WebSocket]
20
  robot_loop = None # type: Optional[asyncio.AbstractEventLoop]
21
  # Add after imports
22
+ _black_frame = np.zeros((640, 640, 3), dtype=np.uint8)
23
+ # Encode it to JPEG immediately
24
+ _, _buffer = cv2.imencode('.jpg', _black_frame)
25
+ # Set the global variable to bytes
26
+ latest_test_frame = _buffer.tobytes()
27
 
28
  @dataclass
29
  class Movement:
 
68
  self.daemon_process = None
69
  self.reachy_mini = None
70
  self.running = False
 
71
 
72
  self.movement_queue: List[Movement] = []
73
  self.is_playing = False
 
124
  frame = cv2.imdecode(arr, cv2.IMREAD_COLOR) # returns BGR np.ndarray
125
  if frame is None:
126
  # Decoding failed, show black frame
127
+ frame = _black_frame
128
  else:
129
+ # Convert BGR->RGB for Gradio and resize to 1080x1080dd
130
  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
131
+ # frame = cv2.resize(frame, (1080, 1080), interpolation=cv2.INTER_LINEAR)
132
  else:
133
  # Unexpected type, show black frame
134
+ frame = _black_frame
135
  decode_time = time.perf_counter() - decode_start
136
  except Exception as e:
137
  print("[stream_frames] JPEG decode error:", e)
138
+ frame = _black_frame
139
  decode_time = 0
140
 
 
141
  frame_count += 1
142
 
143
  # Calculate and print FPS every 30 frames (~1 second at 30fps)
 
152
 
153
  yield frame
154
  time.sleep(0.04)
155
+
156
+ def generate_mjpeg_stream(self):
157
+ global latest_test_frame
158
+ while True:
159
+ # Get the raw bytes (which you said are already JPEG)
160
+ frame_bytes = latest_test_frame
161
+
162
+ if frame_bytes is not None and len(frame_bytes) > 0:
163
+ # Create the MJPEG boundary format
164
+ yield (b'--frame\r\n'
165
+ b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
166
 
167
+ time.sleep(0.04) # Cap at ~25 FPS
168
+
169
  def auto_start(self):
170
  """Auto-start daemon and robot connection"""
171
  status_msgs = []
 
508
 
509
  # Right panel - Simulation view (larger and more square)
510
  with gr.Column(scale=3):
511
+ # sim_view = gr.Image(
512
+ # label="🎬 Robot Simulation",
513
+ # type="numpy",
514
+ # height=1080,
515
+ # width=1080,
516
+ # show_label=True,
517
+ # streaming=True
518
+ # )
519
+ html_code = """
520
+ <html>
521
+ <body>
522
+ <img src="/video_feed" style="width: 100%; max-width: 1080px; border-radius: 8px;">
523
+ </body>
524
+ </html>
525
+ """
526
+ sim_view = gr.HTML(value=html_code, label="🎬 Robot Simulation")
527
 
528
 
529
  # Movement builder section below
 
576
  add_custom_btn = gr.Button("➕ Add to Queue", variant="primary")
577
 
578
  # Stream video
579
+ # demo.load(fn=manager.stream_frames, outputs=sim_view)
580
 
581
  # Connect events - Connection settings
582
  apply_connection_btn.click(
 
698
  except Exception as e:
699
  print("[Space] MuJoCo stream disconnected", e)
700
 
701
+ @app.get("/video_feed")
702
+ def video_feed():
703
+ return StreamingResponse(
704
+ manager.generate_mjpeg_stream(),
705
+ media_type="multipart/x-mixed-replace; boundary=frame"
706
+ )
707
+
708
  @app.websocket("/robot")
709
  async def robot_socket(ws: WebSocket):
710
  global robot_ws, robot_loop