sketchkit.vectorization.LineDrawer.CNNVE.utils package¶
Submodules¶
sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils module¶
- class sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.Edge(start, end=None, pixels=None)[source]¶
Bases:
objectRepresents an edge in a graph, connecting two Vertex objects.
- pixels¶
A list of (y, x) coordinate points that make up the path of this edge.
- Type:
list
- visited¶
A flag used in graph traversal algorithms.
- Type:
bool
- class sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.Vertex(point, degree=0, edges=None)[source]¶
Bases:
objectRepresents a vertex (node) in a graph, typically a pixel in a skeleton.
- point¶
The (y, x) coordinates of the vertex.
- Type:
np.ndarray
- degree¶
The degree of the vertex (number of connected edges).
- Type:
int
- edges¶
A list of Edge objects connected to this vertex.
- Type:
list
- visited¶
A flag used in graph traversal algorithms.
- Type:
bool
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.buildTree(img, start=None)[source]¶
Builds a graph representation from a skeleton image.
Traverses the skeleton image using BFS to construct a graph where pixels are vertices and 8-way connections are edges. If the image contains disconnected components, it will generate a list of graphs.
- Parameters:
img (np.ndarray) – A binary/boolean skeleton image (True or 255 is foreground). The image will be modified in place.
start (np.ndarray, optional) – The (y, x) starting coordinate for the traversal. If None, it finds the first available foreground pixel.
- Returns:
list[nx.Graph]: A list of generated NetworkX graphs.
np.ndarray: The modified image with visited pixels set to False.
- Return type:
tuple
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.draw_simpleGraphs(simpleGraphs, timgSk, outfname)[source]¶
Visualizes the simplified graphs on top of their skeleton.
Draws the skeleton, marks end nodes (degree 1) in red, and cross nodes (degree > 2) in blue.
- Parameters:
simpleGraphs (list[dict]) – A list of graph dictionaries, as generated by simple_graph_gene.
timgSk (np.ndarray) – The boolean skeleton image.
outfname (str) – The file path to save the visualization.
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.getEndNodes(g)[source]¶
Gets key nodes (endpoints and junctions) from a graph.
Key nodes are defined as nodes with degree 1 (endpoints) or degree > 2 (junctions). Includes fallback logic if no such nodes are found.
- Parameters:
g (nx.Graph) – The input NetworkX graph.
- Returns:
A list of key Vertex objects.
- Return type:
list[Vertex]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.getEndNodes_sk(g)[source]¶
Get all nodes with a degree of 1 (end nodes) from a graph.
- Parameters:
g (nx.Graph) – The input graph.
- Returns:
A list of end nodes.
- Return type:
list[Vertex]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.mergeEdges(graph)[source]¶
Simplifies a graph by merging edges connected by degree-2 nodes.
This function iterates through the graph and “dissolves” nodes that have a degree of 2 (i.e., are part of a simple path). The two edges connected to such a node are merged into a single edge, and their pixel lists are concatenated.
- Parameters:
graph (nx.Graph) – The input graph, typically from buildTree.
- Returns:
- A simplified graph where paths of degree-2 nodes have
been contracted into single edges.
- Return type:
nx.Graph
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.merge_list(L)[source]¶
(Deprecated/Legacy) Merges sets in a list that share common elements.
- Parameters:
L (list[set]) – A list of sets.
- Returns:
A list of merged sets.
- Return type:
list[set]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.merge_list_G(lst)[source]¶
Merges lists of integers that share common elements (transitive closure).
Uses a NetworkX graph to find connected components. For a list of pairs [[1, 2], [3, 4], [2, 3]], this will find that 1, 2, 3, and 4 are all connected and return [[1, 2, 3, 4]].
- Parameters:
lst (list[list[int]]) – A list of lists, where each inner list represents a group or a pair of connected items.
- Returns:
A list of the final merged component groups.
- Return type:
list[list[int]]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.my_print(*args, **kwargs)[source]¶
Custom print function that can be toggled by a boolean flag.
- Parameters:
*args – Variable length argument list.
**kwargs – Arbitrary keyword arguments.
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.my_skeletonize(gray_img_blackbg, m_th=128)[source]¶
Performs skeletonization on a grayscale image with a black background.
The image is first binarized using the threshold m_th. It then applies OpenCV’s thinning algorithm followed by skimage’s medial axis transform.
- Parameters:
gray_img_blackbg (np.ndarray) – The input grayscale image. Assumes the object is light and the background is dark.
m_th (int, optional) – The threshold value for binarization. Defaults to 128.
- Returns:
A boolean array representing the skeleton (medial axis).
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.graph_utils.simple_graph_gene(single_stroke_mask, outpath='', is_blackbg=True, get_graph=True)[source]¶
Generates a simplified graph from a single stroke mask.
This function crops the mask to its bounding box, skeletonizes it, builds a detailed pixel-graph, and then simplifies that graph by merging nodes of degree 2.
- Parameters:
single_stroke_mask (np.ndarray) – A grayscale image of a single stroke.
outpath (str, optional) – A path for (currently disabled) caching. Defaults to ‘’.
is_blackbg (bool, optional) – True if the mask has a black background. Defaults to True.
get_graph (bool, optional) – If False, stops after skeletonization and returns no graph. Defaults to True.
- Returns:
list[dict] or None: A list of simplified graph dictionaries. Each dict contains the “graph”, “endNodes”, and “crossNodes”.
np.ndarray or None: The boolean skeleton image.
tuple or None: The bounding box (xmin, xmax, ymin, ymax).
- Return type:
tuple
sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils module¶
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.checkMergedTopo_Components(m1, m2, img_bool_sk, ii, jj, iter_id)[source]¶
Checks if merging two masks preserves their topology.
This heuristic checks if the non-overlapping parts of the masks (m1 - intersection, m2 - intersection) break into more components than they had originally. This is used to prevent invalid merges.
- Parameters:
m1 (dict) – Mask dictionary for the first stroke, must contain ‘update_mask’.
m2 (dict) – Mask dictionary for the second stroke, must contain ‘update_mask’.
img_bool_sk (np.ndarray) – The boolean skeleton of the entire image.
ii (int) – Index of the first mask (for debugging).
jj (int) – Index of the second mask (for debugging).
iter_id (int) – Iteration ID (for debugging).
- Returns:
True if the merge is topologically “safe”, False otherwise.
- Return type:
bool
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.checkMergedTopo_no_increate_cross(m1, m2, img_bool_sk, ii, jj, iter_id)[source]¶
Checks if merging masks would increase cross-nodes (currently stubbed).
Note: This function is hardcoded to always return True. The intended logic (commented out in the source) appears to have been to generate graphs for m1, m2, and their union, then compare cross-node counts.
- Parameters:
m1 (dict) – (Unused) Mask dictionary for the first stroke.
m2 (dict) – (Unused) Mask dictionary for the second stroke.
img_bool_sk (np.ndarray) – (Unused) The boolean skeleton of the entire image.
ii (int) – (Unused) Index of the first mask.
jj (int) – (Unused) Index of the second mask.
iter_id (int) – (Unused) Iteration ID.
- Returns:
Always returns True.
- Return type:
bool
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.color_gene(N=100)[source]¶
Generates a list of N random colors.
- Parameters:
N (int, optional) – The number of colors to generate. Defaults to 100.
- Returns:
A list of N random RGB tuples.
- Return type:
list
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.criterion_mask_overlap(m_s, m_d, merge_th, i, j)[source]¶
Calculates the overlap between two masks to check for merging.
Calculates the intersection (I) and union (U) and pixel counts for each mask (c, d). The “rate” is max(I/c, I/d). A merge is considered valid only if the intersection a is above a minimum pixel threshold I_th.
- Parameters:
m_s (np.ndarray) – Source mask (binary).
m_d (np.ndarray) – Destination mask (binary).
merge_th (float) – (Unused) The overlap threshold.
i (int) – (Unused) Index of first mask.
j (int) – (Unused) Index of second mask.
- Returns:
bool: True if the intersection is non-zero and above I_th.
float or None: The overlap rate, or None if the overlap is too small or zero.
- Return type:
tuple (bool, float or None)
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.get_connectedComponents(im_c1)[source]¶
Counts the number of significant connected components in a binary image.
Components with 3 or fewer pixels are ignored.
- Parameters:
im_c1 (np.ndarray) – A single-channel, 8-bit binary image.
- Returns:
The number of components with more than 3 pixels.
- Return type:
int
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.get_rect_p0p1(cvimg)[source]¶
Finds the bounding box of non-white pixels in a white-background image.
- Parameters:
cvimg (np.ndarray) – A grayscale image, assumed to have a white (or >127) background.
- Returns:
- The bounding box (xmin, xmax, ymin, ymax).
Returns (None, None, None, None) if no non-white pixels are found.
- Return type:
tuple (int, int, int, int)
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.get_rect_p0p1_blackbg(cvimg)[source]¶
Finds the bounding box of non-black pixels in a black-background image.
This is a wrapper around get_rect_p0p1 that inverts the image first.
- Parameters:
cvimg (np.ndarray) – A grayscale image with a black background.
- Returns:
The bounding box (xmin, xmax, ymin, ymax).
- Return type:
tuple (int, int, int, int)
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.is_background_black(image, threshold=0.8)[source]¶
Checks if an image’s background is predominantly black.
Checks if the ratio of pure black pixels (value 0) exceeds the threshold.
- Parameters:
image (np.ndarray) – The input image.
threshold (float, optional) – The ratio of black pixels. Defaults to 0.8.
- Returns:
True if the background is predominantly black, False otherwise.
- Return type:
bool
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.is_background_black_undopainst(image, threshold=0.6)[source]¶
Checks if an image’s background is predominantly black.
Note: This function modifies the input image by binarizing it. It checks if the ratio of black pixels (<= 127) exceeds the threshold.
- Parameters:
image (np.ndarray) – The input image (will be modified).
threshold (float, optional) – The ratio of black pixels to consider the background black. Defaults to 0.6.
- Returns:
True if the background is predominantly black, False otherwise.
- Return type:
bool
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.loadImage(f)[source]¶
Loads an image from a file path as a float array.
- Parameters:
f (str) – The file path.
- Returns:
The image as a float array (values 0.0 to 1.0).
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.rect_gene(point, rect_w)[source]¶
Generates a square bounding box (rectangle) centered at a point.
- Parameters:
point (tuple[int, int]) – The center point (x, y).
rect_w (int) – The “radius” of the box; the box will have width and height of 2 * rect_w.
- Returns:
- The rectangle (x, y, w, h), where (x, y)
is the top-left corner.
- Return type:
tuple[int, int, int, int]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.revise_mask_with_connectedComponents(im_c1, pixel_lst)[source]¶
Cleans a mask by keeping only the connected component touching the skeleton.
Given a predicted mask (im_c1) and the list of skeleton pixels (pixel_lst) that generated it, this function finds all connected components in the mask. It then returns only the component that is “touched” by the first pixel in pixel_lst.
- Parameters:
im_c1 (np.ndarray) – The single-channel mask to revise (8-bit).
pixel_lst (list[tuple]) – The list of (y, x) skeleton pixels.
- Returns:
- The revised 8-bit mask, containing only the selected
connected component.
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.image_utils.save_gif(gif_path, images, duration=0.05)[source]¶
Saves a list of images as a GIF.
- Parameters:
gif_path (str) – The file path to save the GIF to.
images (list[np.ndarray]) – A list of image frames.
duration (float, optional) – The duration (in seconds) for each frame. Defaults to 0.05.
sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils module¶
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.dist_L2(points)[source]¶
Calculates the total L2 (Euclidean) distance of a path.
- Parameters:
points (list or np.ndarray) – A list of (x, y) points.
- Returns:
The total path length.
- Return type:
float
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.gene_list_index_map(model, simpleGraphs, im_ori_c1, outdir, th=0.2, img_bool_sk=None, iter_id=0)[source]¶
Generates stroke masks for each edge in the simplified graphs.
Iterates through all edges in simpleGraphs. For each edge: 1. Gets the skeleton pixel list. 2. If it’s a simple component, gets the mask from connectedComponents. 3. Otherwise, generates a mask using the ML model (update_mask_edge). 4. Cleans the generated mask (revise_mask_with_connectedComponents). 5. Stores the mask and related info in a dictionary mdic indexed
by a counter.
Creates a lookup table edic mapping edge names to the counter index.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
simpleGraphs (list[dict]) – List of simplified graph dictionaries.
im_ori_c1 (np.ndarray) – The original binarized image (black bg).
outdir (str) – (Unused) Output directory.
th (float, optional) – (Unused) Threshold. Defaults to 0.2.
img_bool_sk (np.ndarray, optional) – (Unused) Boolean skeleton image.
iter_id (int, optional) – (Unused) Iteration ID.
- Returns:
- A dictionary containing:
’EdgeDic’ (edic): Maps edge names to mask IDs.
’SplitMaskDic’ (mdic): Maps mask IDs to mask data.
’MergeList’: A list of all mask IDs, unmerged.
- Return type:
dict
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.line_extraction()[source]¶
Initializes and loads the line extraction deep learning model.
Loads the stNet_double model from CNNVE.pm and loads pre-trained weights (“500.pth”) onto the available device (CUDA or CPU).
- Returns:
The initialized and loaded PyTorch model.
- Return type:
pm.stNet_double
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.line_extraction_batch_with_mask(model, img_c1, point_list_xy, mask_list, rect_w=256)[source]¶
Performs batch line extraction using a model on specified masks.
Takes a list of masks and their corresponding center points, generates bounding boxes, and runs them through the model’s get_result_batch_s2 method in batches.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
img_c1 (np.ndarray) – The original single-channel image.
point_list_xy (list[list[tuple]]) – A list where each item is a list of (x, y) points. The first point is used as the center.
mask_list (list[np.ndarray]) – A list of mask images corresponding to point_list_xy.
rect_w (int, optional) – The “radius” for the bounding box. Defaults to 256.
- Returns:
tuple(np.ndarray, np.ndarray, np.ndarray): Model results (m1, m2, m0) concatenated from all batches.
list[tuple]: The list of generated rectangles (rect_lst).
- Return type:
tuple
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.merge_list_1st(mask_dic, th, img_bool_sk=None, iter_id=0)[source]¶
Performs the first pass of stroke merging.
Iterates through all junctions (nodes with degree > 2) in the simplified graphs. For each pair of edges connected to that junction, it checks if their corresponding masks overlap significantly (criterion_mask_overlap) and pass a topological check (checkMergedTopo_Components and checkMergedTopo_no_increate_cross).
If they pass, they are added to a component list (c_list) for merging. This function updates the mask_dic with the merge results.
- Parameters:
mask_dic (dict) – The main dictionary containing ‘simpleGraphs’, ‘SplitMaskDic’, and ‘EdgeDic’. This dict is modified in place.
th (float) – The overlap rate threshold for merging.
img_bool_sk (np.ndarray, optional) – The boolean skeleton image.
iter_id (int, optional) – The current iteration ID.
- Returns:
The list of merged pairs (c_list).
- Return type:
list[list[int]]
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.stroke_extracion_multi(model, points, im_ori, b_with_remove=False)[source]¶
Extracts strokes from an image using a model, given multiple center points.
For each point in points, it creates a small circular mask, then runs the batch prediction model.get_result_batch. Finally, it combines all the resulting masks using bitwise OR.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
points (list[tuple]) – A list of (x, y) center points.
im_ori (np.ndarray) – The original single-channel image.
b_with_remove (bool, optional) – (Unused) Defaults to False.
- Returns:
A (H, W, 1) mask combining all extracted strokes.
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.stroke_extracion_multi_segments(model, points, im_ori, b_with_remove=False)[source]¶
Extracts strokes, optimized for long segments of points.
If the number of points is small, it falls back to stroke_extracion_multi. If large, it breaks the point list into overlapping segments, creates a polyline mask for each segment, and runs the batch prediction model.get_result_batch_s2.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
points (list[tuple]) – A list of (x, y) points defining the stroke.
im_ori (np.ndarray) – The original single-channel image.
b_with_remove (bool, optional) – (Unused) Defaults to False.
- Returns:
A (H, W, 1) mask combining all extracted strokes.
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.truncate(points)[source]¶
(Deprecated/Stub) Truncates the start and end of a point list.
Note: This function is hardcoded to return the original list. The commented-out logic suggests an intent to remove a small number of points from the beginning and end of the list.
- Parameters:
points (list) – A list of points.
- Returns:
The original, unmodified list of points.
- Return type:
list
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.update_mask(model, ori_img, mask_curr, pixel_lst)[source]¶
(Deprecated/Heuristic) Updates a mask if it seems to mismatch the skeleton.
Checks if a large number of skeleton pixels (pixel_lst) fall outside the current mask (mask_curr). If the mismatch is large, it re-generates the mask using stroke_extracion_multi_segments.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
ori_img (np.ndarray) – The original single-channel image.
mask_curr (np.ndarray) – The existing (H, W, 1) mask.
pixel_lst (list[tuple]) – A list of (y, x) skeleton pixels.
- Returns:
The original mask or a newly generated one.
- Return type:
np.ndarray
- sketchkit.vectorization.LineDrawer.CNNVE.utils.vectorization_utils.update_mask_edge(model, ori_img, pixel_lst)[source]¶
Generates a stroke mask from a list of skeleton pixels.
This function converts skeleton (y, x) pixels to (x, y) points, calculates the path length, and decides on an extraction strategy. - If the path is very short, it uses the midpoint. - Otherwise, it uses the (potentially truncated) full list of points. It calls stroke_extracion_multi_segments to get the final mask.
- Parameters:
model (pm.stNet_double) – The loaded line extraction model.
ori_img (np.ndarray) – The original single-channel image.
pixel_lst (list[tuple]) – A list of (y, x) skeleton pixels.
- Returns:
The (H, W, 1) generated stroke mask.
- Return type:
np.ndarray