Add cv2.imread and cv2.imwrite

.github/workflows/test.yml

@@ -20,7 +20,7 @@ jobs:
     - name: Deps
       run: bash ./scripts/deps_ci.sh
     - name: Download Test Video
-      run: curl -O https://f.dominik.win/data/dve2/tos_720p.mp4
+      run: curl -O https://f.dominik.win/data/dve2/tos_720p.mp4 && curl -O https://f.dominik.win/data/dve2/apollo.jpg
     - name: Build Debug
       run: cargo build
     - name: Build Release
@@ -31,7 +31,7 @@ jobs:
     - name: Install vidformer-py
       run: pip3 install ./vidformer-py
     - name: snake-pit pytest
-      run: ln -s ../tos_720p.mp4 snake-pit/tos_720p.mp4 && pushd snake-pit && VIDFORMER_BIN='../target/debug/vidformer-cli' pytest . --verbose && popd
+      run: ln -s ../tos_720p.mp4 snake-pit/tos_720p.mp4 && cp apollo.jpg snake-pit/apollo.jpg && pushd snake-pit && VIDFORMER_BIN='../target/debug/vidformer-cli' pytest . --verbose && popd
       timeout-minutes: 10
     - name: Valgrind Tests
       run: bash ./scripts/valgrind_test.sh

.gitignore

@@ -3,6 +3,9 @@
 /data
 Cargo.lock
 *.mp4
+*.mov
+*.jpg
+*.jpeg
 *.mkv
 *.srt
 *.json

scripts/deps_devcontainer.sh

@@ -11,3 +11,4 @@ pip3 install -e ./vidformer-py
 
 # Download test video if it doesn't exist
 [ -f tos_720p.mp4 ] || curl -O https://f.dominik.win/data/dve2/tos_720p.mp4
+[-f apollo.jpg ] || curl -O https://f.dominik.win/data/dve2/apollo.jpg

snake-pit/test_cv2.py

@@ -47,6 +47,11 @@ def rw(cv2):
         ret, frame = cap.read()
         if not ret or count > 100:
             break
+
+        assert frame.shape[0] == height
+        assert frame.shape[1] == width
+        assert frame.shape[2] == 3
+
         out.write(frame)
         count += 1
 
@@ -426,3 +431,79 @@ def test_addWeighted_ocv():
 
 def test_addWeighted_vf():
     addWeighted(vf_cv2)
+
+
+def test_imread():
+    import vidformer.cv2 as vf_cv2
+
+    img = vf_cv2.imread("apollo.jpg")
+
+    assert img._fmt["width"] == 3912
+    assert img._fmt["height"] == 3936
+    assert img._fmt["pix_fmt"] == "yuvj444p"
+
+
+def imread(cv2):
+    img = cv2.imread("apollo.jpg")
+
+    assert img.shape[0] == 3936
+    assert img.shape[1] == 3912
+    assert img.shape[2] == 3
+
+
+def test_imread_ocv():
+    imread(ocv_cv2)
+
+
+def test_imread_vf():
+    imread(vf_cv2)
+
+
+def imwrite(cv2):
+    # from apollo.jpg
+    img = cv2.imread("apollo.jpg")
+
+    # jpg
+    cv2.imwrite("apollo2.jpg", img)
+    assert os.path.exists("apollo2.jpg")
+    os.remove("apollo2.jpg")
+
+    # jpeg
+    cv2.imwrite("apollo2.jpeg", img)
+    assert os.path.exists("apollo2.jpeg")
+    os.remove("apollo2.jpeg")
+
+    # png
+    cv2.imwrite("apollo2.png", img)
+    assert os.path.exists("apollo2.png")
+    os.remove("apollo2.png")
+
+    # from 1000th frame of tos_720p.mp4
+    cap = cv2.VideoCapture(VID_PATH)
+    assert cap.isOpened()
+
+    cap.set(cv2.CAP_PROP_POS_FRAMES, 1000)
+    ret, frame = cap.read()
+    assert ret
+
+    # jpg
+    cv2.imwrite("tos_720p_1000.jpg", frame)
+    assert os.path.exists("tos_720p_1000.jpg")
+    os.remove("tos_720p_1000.jpg")
+
+    # jpeg
+    cv2.imwrite("tos_720p_1000.jpeg", frame)
+    assert os.path.exists("tos_720p_1000.jpeg")
+    os.remove("tos_720p_1000.jpeg")
+
+    # png
+    cv2.imwrite("tos_720p_1000.png", frame)
+    assert os.path.exists("tos_720p_1000.png")
+
+
+def test_imwrite_ocv():
+    imwrite(ocv_cv2)
+
+
+def test_imwrite_vf():
+    imwrite(vf_cv2)

vidformer-py/vidformer/cv2/vf_cv2.py

@@ -66,8 +66,10 @@ def set_cv2_server(server):
 
 
 class _Frame:
-    def __init__(self, f):
+    def __init__(self, f, fmt):
         self._f = f
+        self._fmt = fmt
+        self.shape = (fmt["height"], fmt["width"], 3)
 
         # denotes that the frame has not yet been modified
         # when a frame is modified, it is converted to rgb24 first
@@ -115,7 +117,7 @@ def read(self):
             return False, None
         frame = self._source.iloc[self._next_frame_idx]
         self._next_frame_idx += 1
-        frame = _Frame(frame)
+        frame = _Frame(frame, self._source.fmt())
         return True, frame
 
     def release(self):
@@ -163,6 +165,55 @@ def __init__(self, *args):
         self._args = args
 
 
+def imread(path, *args):
+    if len(args) > 0:
+        raise NotImplementedError("imread does not support additional arguments")
+
+    assert path.lower().endswith((".jpg", ".jpeg", ".png"))
+    server = _server()
+    source = vf.Source(server, str(uuid.uuid4()), path, 0)
+    frame = _Frame(source.iloc[0], source.fmt())
+    return frame
+
+
+def imwrite(path, img, *args):
+    if len(args) > 0:
+        raise NotImplementedError("imwrite does not support additional arguments")
+
+    if not isinstance(img, _Frame):
+        raise Exception("img must be a _Frame object")
+
+    fmt = img._fmt.copy()
+    width = fmt["width"]
+    height = fmt["height"]
+    f = img._f
+
+    domain = [Fraction(0, 1)]
+
+    if path.lower().endswith(".png"):
+        img._mut()  # Make sure it's in rgb24
+        spec = vf.Spec(
+            domain,
+            lambda t, i: img._f,
+            {"width": width, "height": height, "pix_fmt": "rgb24"},
+        )
+        spec.save(_server(), path, encoder="png")
+    elif path.lower().endswith((".jpg", ".jpeg")):
+        if img._modified:
+            # it's rgb24, we need to convert to something jpeg can handle
+            f = _filter_scale(img._f, pix_fmt="yuv420p")
+            fmt["pix_fmt"] = "yuv420p"
+        else:
+            if fmt["pix_fmt"] not in ["yuvj420p", "yuvj422p", "yuvj444p"]:
+                f = _filter_scale(img._f, pix_fmt="yuvj420p")
+                fmt["pix_fmt"] = "yuvj420p"
+
+        spec = vf.Spec(domain, lambda t, i: f, fmt)
+        spec.save(_server(), path, encoder="mjpeg")
+    else:
+        raise Exception("Unsupported image format")
+
+
 def rectangle(img, pt1, pt2, color, thickness=None, lineType=None, shift=None):
     """
     cv.rectangle(	img, pt1, pt2, color[, thickness[, lineType[, shift]]]	)
@@ -372,7 +423,7 @@ def addWeighted(src1, alpha, src2, beta, gamma, dst=None, dtype=-1):
     src2._mut()
 
     if dst is None:
-        dst = _Frame(src1._f)
+        dst = _Frame(src1._f, src1._fmt.copy())
     else:
         assert isinstance(dst, _Frame)
     dst._mut()