fix camera.py

rafcabezas · Apr 21, 2020 · b9f5882 · b9f5882
1 parent dd22cfc
commit b9f5882
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/tools/replay/camera.py b/tools/replay/camera.py
@@ -3,7 +3,7 @@
 
 from common.basedir import BASEDIR
 os.environ['BASEDIR'] = BASEDIR
-SCALE = 3
+SCALE = 1
 
 import argparse
 import zmq
@@ -18,6 +18,7 @@
 from tools.lib.lazy_property import lazy_property
 from cereal.messaging import sub_sock, recv_one_or_none, recv_one
 from cereal.services import service_list
+import cereal.messaging as messaging
 
 _BB_OFFSET = 0, 0
 _BB_TO_FULL_FRAME = np.asarray([[1., 0., _BB_OFFSET[0]], [0., 1., _BB_OFFSET[1]],
@@ -44,9 +45,7 @@ def ui_thread(addr, frame_address):
 
   camera_surface = pygame.surface.Surface((_FULL_FRAME_SIZE[0] * SCALE, _FULL_FRAME_SIZE[1] * SCALE), 0, 24).convert()
 
-  frame = context.socket(zmq.SUB)
-  frame.connect(frame_address or "tcp://%s:%d" % (addr, service_list['frame'].port))
-  frame.setsockopt(zmq.SUBSCRIBE, "")
+  frame = messaging.sub_sock('frame', conflate=True)
 
   img = np.zeros((_FULL_FRAME_SIZE[1], _FULL_FRAME_SIZE[0], 3), dtype='uint8')
   imgff = np.zeros((_FULL_FRAME_SIZE[1], _FULL_FRAME_SIZE[0], 3), dtype=np.uint8)
@@ -56,7 +55,7 @@ def ui_thread(addr, frame_address):
     screen.fill((64, 64, 64))
 
     # ***** frame *****
-    fpkt = recv_one(frame)
+    fpkt = messaging.recv_one(frame)
     yuv_img = fpkt.frame.image
 
     if fpkt.frame.transform:
@@ -76,7 +75,9 @@ def ui_thread(addr, frame_address):
         dst=img,
         flags=cv2.WARP_INVERSE_MAP)
     else:
-      img.fill(0)
+      # actually RGB
+      img = np.frombuffer(yuv_img, dtype=np.uint8).reshape((_FULL_FRAME_SIZE[1], _FULL_FRAME_SIZE[0], 3))
+      img = img[:, :, ::-1] # Convert BGR to RGB
 
     height, width = img.shape[:2]
     img_resized = cv2.resize(