Merge pull request #294 from dclong/dev

Merge dev into main

Author: dclong

dsutil/__init__.py

@@ -3,4 +3,4 @@
 from . import git
 from . import poetry
 
-__version__ = "0.64.0"
+__version__ = "0.65.0"

dsutil/cv.py

@@ -3,28 +3,41 @@
 from __future__ import annotations
 from typing import Union, Iterable
 from pathlib import Path
-from tqdm import tqdm
+from tqdm import tqdm, trange
 import numpy as np
+import pandas as pd
 from PIL import Image
+import skimage
 import cv2
 
 
-def video_to_image(video, frames_per_image: int = 60, output: str = "frame_{:0>7}.png"):
-    """Extract images from a video.
+def video_to_image(
+    file: str,
+    step: int,
+    bbox: Union[None, tuple[int, int, int, int]] = None,
+    output: str = "frame_{:0>7}.png"
+):
+    """Extract images from a video file.
 
-    :param video: The path to a video file.
-    :param frames_per_image: Extract 1 image every frames_per_image.
-    :param output: The pattern of the output files for the extracted images.
+    :param file: The path to video file.
+    :param bbox: A bounding box. 
+        If specified, crop images using the bounding box.
+    :param output_dir: The directory to save extracted images.
+    :param step: Keep 1 image every step frames.
     """
-    vidcap = cv2.VideoCapture(video)
-    count = 0
-    while True:
-        success, image = vidcap.read()
+    Path(output.format(0)).parent.mkdir(parents=True, exist_ok=True)
+    vidcap = cv2.VideoCapture(file)
+    total = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    for idx in trange(total):
+        success, arr = vidcap.read()
         if not success:
             break
-        if count % frames_per_image == 0:
-            cv2.imwrite(output.format(count), image)
-        count += 1
+        if idx % step == 0:
+            img = Image.fromarray(np.flip(arr, 2))
+            if bbox:
+                img = img.crop(bbox)
+            img.save(output.format(idx))
+    vidcap.release()
 
 
 def resize_image(
@@ -201,56 +214,6 @@ def deshade_3(img, threshold=0.4, cutoff=30) -> Image.Image:
     return Image.fromarray(arr)
 
 
-def highlight_frame(
-    rgb: tuple[int, int, int],
-    shape: tuple[int, int],
-    thickness: int = 3
-) -> Image.Image:
-    """Generate a rectangle frame with the specified color and thickness.
-
-    :param rgb: The color in RGB (as a tuple) to use for the frame.
-    :param shape: The shape of the frame.
-    :param thickness: The thickness of the frame.
-    :return: A PIL image presenting the frame.
-    """
-    nrow = shape[0]
-    ncol = shape[1]
-    arr = np.zeros((nrow, ncol, 3), np.uint8)
-    arr[0:thickness, :, 0] = rgb[0]
-    arr[0:thickness, :, 1] = rgb[1]
-    arr[0:thickness, :, 2] = rgb[2]
-    arr[(nrow - thickness):nrow, :, 0] = rgb[0]
-    arr[(nrow - thickness):nrow, :, 1] = rgb[1]
-    arr[(nrow - thickness):nrow, :, 2] = rgb[2]
-    arr[:, 0:thickness, 0] = rgb[0]
-    arr[:, 0:thickness, 1] = rgb[1]
-    arr[:, 0:thickness, 2] = rgb[2]
-    arr[:, (ncol - thickness):ncol, 0] = rgb[0]
-    arr[:, (ncol - thickness):ncol, 1] = rgb[1]
-    arr[:, (ncol - thickness):ncol, 2] = rgb[2]
-    return Image.fromarray(arr)
-
-
-def frame_image(
-    img: Image.Image, rgb: tuple[int, int, int], thickness: int = 3
-) -> Image.Image:
-    """Add a highlight frame to an image.
-
-    :param img: A PIL image.
-    :param rgb: The color in RGB (as a tuple) to use for the frame.
-    :param thickness: The thickness of the frame.
-    :return: A new image with the frame added.
-    """
-    shape = img.size
-    shape = (shape[1], shape[0])
-    frame = highlight_frame(rgb, shape=shape, thickness=thickness)
-    mask = highlight_frame((255, 255, 255), shape=shape,
-                           thickness=thickness).convert("1")
-    img = img.copy()
-    img.paste(frame, mask=mask)
-    return img
-
-
 def add_frames(
     arr: Union[np.ndarray, Image.Image],
     bboxes: list[tuple[int, int, int, int]],
@@ -275,3 +238,63 @@ def add_frames(
         arr[y1:y2, x1, :] = rgb
         arr[y1:y2, x2, :] = rgb
     return arr
+
+
+def duplicate_image(
+    path: Union[str, Path],
+    copies: int,
+    des_dir: Union[str, Path, None] = None,
+    noise_amount: float = 0.05
+):
+    """Duplicate an image with some noises added.
+
+    :param path: The path to the image to be duplicated.
+    :param copies: The number of copies to duplicate.
+    :param noise_amount: Proportion of image pixels to replace with noise on range [0, 1].
+    """
+    if isinstance(path, str):
+        path = Path(path)
+    if isinstance(des_dir, str):
+        des_dir = Path(des_dir)
+    if des_dir is None:
+        des_dir = path.parent
+    des_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(copies):
+        file_i = des_dir / f"{path.stem}_copy{i}.png"
+        noise = skimage.util.random_noise(
+            np.array(Image.open(path)), mode="s&p", amount=noise_amount
+        )
+        Image.fromarray(np.array(noise * 255, dtype=np.uint8)).save(file_i)
+
+
+def structural_similarity(im1, im2) -> float:
+    """Extend skimage.metrics.structural_similarity 
+        to calculate the similarity of (any) two images.
+
+    :param im1: A PIL image.
+    :param im2: Another PIL image.
+    """
+    size = im1.size
+    if im2.size != size:
+        im2 = im2.resize(size)
+    return skimage.metrics.structural_similarity(
+        np.array(im1), np.array(im2), multichannel=True
+    )
+
+
+def calc_image_similarities(img: Union[Image.Image, str, Path], dir_: Union[str, Path]):
+    """Calculate the similarities between an image and all images in a directory.
+
+    :param img: A PIL image or the path to an image file.
+    :param dir_: A directory containing images.
+    """
+    if isinstance(img, (str, Path)):
+        img = Image.open(img)
+    if isinstance(dir_, str):
+        dir_ = Path(dir_)
+    paths = list(dir_.glob("*.png"))
+    sims = [structural_similarity(img, Image.open(p)) for p in tqdm(paths)]
+    return pd.DataFrame({
+        "path": paths,
+        "similarity": sims,
+    })

dsutil/docker/builder.py

@@ -54,10 +54,18 @@ def _push_image_timing(repo: str, tag: str) -> DockerActionResult:
     #    print()
     time_begin = time.perf_counter_ns()
     try:
-        retry(lambda: sp.run(f"docker push {repo}:{tag}", shell=True, check=True), times=3)
-        return DockerActionResult(True, "", repo, tag, "push", (time.perf_counter_ns() - time_begin) / 1E9)
+        retry(
+            lambda: sp.run(f"docker push {repo}:{tag}", shell=True, check=True),
+            times=3
+        )
+        return DockerActionResult(
+            True, "", repo, tag, "push", (time.perf_counter_ns() - time_begin) / 1E9
+        )
     except Exception as err:
-        return DockerActionResult(False, str(err), repo, tag, "push", (time.perf_counter_ns() - time_begin) / 1E9)
+        return DockerActionResult(
+            False, str(err), repo, tag, "push",
+            (time.perf_counter_ns() - time_begin) / 1E9
+        )
 
 
 #def _is_image_pushed(msg: dict[str, Any]):

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dsutil"
-version = "0.64.0"
+version = "0.65.0"
 description = "A utils Python package for data scientists."
 authors = ["Benjamin Du <[email protected]>"]
 
@@ -42,6 +42,7 @@ nbformat = { version = ">=5.0.7", optional = true }
 nbconvert = { version = ">=5.6.1", optional = true }
 yapf = { version = ">=0.30.0", optional = true}
 dulwich = ">=0.20.24"
+scikit-image = ">=0.18.3"
 
 [tool.poetry.extras]
 cv = ["opencv-python", "pillow"]