Fix small bug for jpeg encoding

dora-rs · Oct 10, 2024 · 4dcf2cf · 4dcf2cf
1 parent 254e74d
commit 4dcf2cf
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 19 deletions.
diff --git a/node-hub/dora-qwenvl/dora_qwenvl/main.py b/node-hub/dora-qwenvl/dora_qwenvl/main.py
@@ -116,23 +116,29 @@ def main():
  elif encoding == "rgb8":
  channels = 3
  storage_type = np.uint8
+ elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
+ channels = 3
+ storage_type = np.uint8
  else:
  raise RuntimeError(f"Unsupported image encoding: {encoding}")
 
- frame = (
- storage.to_numpy()
- .astype(storage_type)
- .reshape((height, width, channels))
- )
  if encoding == "bgr8":
+ frame = (
+ storage.to_numpy()
+ .astype(storage_type)
+ .reshape((height, width, channels))
+ )
  frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB)
  elif encoding == "rgb8":
- pass
+ frame = (
+ storage.to_numpy()
+ .astype(storage_type)
+ .reshape((height, width, channels))
+ )
  elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
- channels = 3
- storage_type = np.uint8
  storage = storage.to_numpy()
  frame = cv2.imdecode(storage, cv2.IMREAD_COLOR)
+ frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB)
  else:
  raise RuntimeError(f"Unsupported image encoding: {encoding}")
  frames[event_id] = Image.fromarray(frame)

diff --git a/node-hub/llama-factory-recorder/llama_factory_recorder/main.py b/node-hub/llama-factory-recorder/llama_factory_recorder/main.py
@@ -147,23 +147,31 @@ def main():
  elif encoding == "rgb8":
  channels = 3
  storage_type = np.uint8
+ elif encoding == "jpeg":
+ channels = 3
+ storage_type = np.uint8
  else:
  raise RuntimeError(f"Unsupported image encoding: {encoding}")
 
- frame = (
- storage.to_numpy()
- .astype(storage_type)
- .reshape((height, width, channels))
- )
  if encoding == "bgr8":
+ frame = (
+ storage.to_numpy()
+ .astype(storage_type)
+ .reshape((height, width, channels))
+ )
  frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB)
  elif encoding == "rgb8":
- pass
+ frame = (
+ storage.to_numpy()
+ .astype(storage_type)
+ .reshape((height, width, channels))
+ )
  elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
  channels = 3
  storage_type = np.uint8
  storage = storage.to_numpy()
  frame = cv2.imdecode(storage, cv2.IMREAD_COLOR)
+ frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB)
  else:
  raise RuntimeError(f"Unsupported image encoding: {encoding}")
 

diff --git a/node-hub/opencv-video-capture/opencv_video_capture/main.py b/node-hub/opencv-video-capture/opencv_video_capture/main.py
@@ -10,6 +10,8 @@
 
 RUNNER_CI = True if os.getenv("CI") == "true" else False
 
+FLIP = os.getenv("FLIP", "")
+
 
 def main():
  # Handle dynamic nodes, ask for the name of the node in the dataflow, and the same values as the ENV variables.
@@ -101,10 +103,12 @@ def main():
  1,
  )
 
- metadata = event["metadata"]
- metadata["encoding"] = encoding
- metadata["width"] = int(frame.shape[1])
- metadata["height"] = int(frame.shape[0])
+ if FLIP == "VERTICAL":
+ frame = cv2.flip(frame, 0)
+ elif FLIP == "HORIZONTAL":
+ frame = cv2.flip(frame, 1)
+ elif FLIP == "BOTH":
+ frame = cv2.flip(frame, -1)
 
  # resize the frame
  if (
@@ -116,13 +120,18 @@ def main():
  ):
  frame = cv2.resize(frame, (image_width, image_height))
 
+ metadata = event["metadata"]
+ metadata["encoding"] = encoding
+ metadata["width"] = int(frame.shape[1])
+ metadata["height"] = int(frame.shape[0])
+
  # Get the right encoding
  if encoding == "rgb8":
  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
  elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
  ret, frame = cv2.imencode("." + encoding, frame)
  if not ret:
- print("Could not encode image...")
+ print("Error encoding image...")
  continue
 
  storage = pa.array(frame.ravel())