OpenCV 5 Python Bindings Is Here

1sam_result = sam_segment(img)
2if sam_result is None:
3    print("  SAM failed, aborting")
4    return
5mask, point_coords = sam_result
6
7inpainted = lama_inpaint(img, mask)

 1
 2def sam_segment(img, point_coords=None):
 3    """Segment object using SAM2.1 with point prompts."""
 4    encoder_path = MODELS_DIR / "sam_encoder_v2.onnx"
 5    decoder_path = MODELS_DIR / "sam_decoder_v2.onnx"
 6    if not encoder_path.exists() or not decoder_path.exists():
 7        print(f"  [!] SAM models not found in {MODELS_DIR}")
 8        return None
 9
10    t0 = time.time()
11    h, w = img.shape[:2]
12
13    scale = SAM_IMAGE_SIZE / max(h, w)
14    new_h, new_w = int(h * scale), int(w * scale)
15    resized = cv.resize(img, (new_w, new_h), interpolation=cv.INTER_LINEAR)
16    padded = np.zeros((SAM_IMAGE_SIZE, SAM_IMAGE_SIZE, 3), dtype=np.uint8)
17    padded[:new_h, :new_w] = resized
18
19    pixel_values = padded.astype(np.float32) / 255.0
20    pixel_values = (pixel_values - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
21    pixel_values = pixel_values.transpose(2, 0, 1)[np.newaxis]
22
23    encoder = cv.dnn.readNetFromONNX(str(encoder_path), engine=cv.dnn.ENGINE_NEW)
24    encoder.setInput(pixel_values, "image")
25    image_embed = encoder.forward("image_embed")
26    high_res_0 = encoder.forward("high_res_feats_0")
27    high_res_1 = encoder.forward("high_res_feats_1")
28
29    if point_coords is None:
30        point_coords = [(w // 2, h // 2)]
31    point_labels = [1] * len(point_coords)
32    print(f"  Point prompts: {point_coords}")
33
34    scale_x, scale_y = new_w / w, new_h / h
35    pts = np.array([[[int(x * scale_x), int(y * scale_y)] for x, y in point_coords]], dtype=np.float32)
36    labels = np.array([point_labels], dtype=np.float32)
37
38    mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
39    has_mask_input = np.array([0.0], dtype=np.float32)
40
41    decoder = cv.dnn.readNetFromONNX(str(decoder_path), engine=cv.dnn.ENGINE_AUTO)
42    decoder.setInput(image_embed, "image_embed")
43    decoder.setInput(high_res_0, "high_res_feats_0")
44    decoder.setInput(high_res_1, "high_res_feats_1")
45    decoder.setInput(pts, "point_coords")
46    decoder.setInput(labels, "point_labels")
47    decoder.setInput(mask_input, "mask_input")
48    decoder.setInput(has_mask_input, "has_mask_input")
49
50    masks = decoder.forward("masks")
51    iou_predictions = decoder.forward("iou_predictions")
52
53    best_idx, best_score = 0, -1
54    for i in range(masks.shape[1]):
55        sig = 1.0 / (1.0 + np.exp(-masks[0, i]))
56        area_ratio = sig.mean()
57        if 0.02 < area_ratio < 0.80:
58            score = iou_predictions[0, i]
59            if score > best_score:
60                best_score = score
61                best_idx = i
62    if best_score < 0:
63        for i in range(masks.shape[1]):
64            sig = 1.0 / (1.0 + np.exp(-masks[0, i]))
65            area_ratio = sig.mean()
66            if area_ratio < 0.95:
67                score = iou_predictions[0, i]
68                if score > best_score:
69                    best_score = score
70                    best_idx = i
71    if best_score < 0:
72        best_idx = int(np.argmax(iou_predictions[0]))
73
74    mask_logits = masks[0, best_idx]
75    mask = (1.0 / (1.0 + np.exp(-mask_logits)) > 0.5).astype(np.uint8) * 255
76    mask = cv.resize(mask, (w, h), interpolation=cv.INTER_NEAREST)
77
78    kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (15, 15))
79    mask = cv.dilate(mask, kernel, iterations=3)
80
81    elapsed = time.time() - t0
82    coverage = (mask > 127).mean() * 100
83    print(f"  SAM: IoU={iou_predictions[0, best_idx]:.3f}, coverage={coverage:.1f}%  [{elapsed:.2f}s]")
84    return mask, point_coords

 1
 2def lama_inpaint(img, mask):
 3    """Remove object using LaMa inpainting model."""
 4    model_path = MODELS_DIR / "lama.onnx"
 5    if not model_path.exists():
 6        print(f"  [!] LaMa model not found: {model_path}")
 7        return None
 8
 9    t0 = time.time()
10    h, w = img.shape[:2]
11
12    resized_img = cv.resize(img, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE))
13    resized_mask = cv.resize(mask, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE))
14    _, resized_mask = cv.threshold(resized_mask, 127, 255, cv.THRESH_BINARY)
15
16    image_blob = cv.dnn.blobFromImage(resized_img, 1 / 255.0, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE))
17    mask_blob = cv.dnn.blobFromImage(resized_mask, scalefactor=1.0, size=(LAMA_MODEL_SIZE, LAMA_MODEL_SIZE), mean=(0,), swapRB=False, crop=False)
18    mask_blob = (mask_blob > 0).astype(np.float32)
19
20    net = cv.dnn.readNetFromONNX(str(model_path), engine=cv.dnn.ENGINE_AUTO)
21    net.setInput(image_blob, "image")
22    net.setInput(mask_blob, "mask")
23    output = net.forward()
24
25    result = output[0]
26    result = np.transpose(result, (1, 2, 0))
27    result = np.clip(result, 0, 255).astype(np.uint8)
28    result = cv.resize(result, (w, h))
29
30    elapsed = time.time() - t0
31    print(f"  LaMa: inpainting complete  [{elapsed:.2f}s]")
32    return result