CoolFace
Datasetpublic

skywalker290/Watermeter

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes70downloads
extract.py101 linesDownload Raw Back to root
1import os2import cv23import numpy as np4from tqdm import tqdm5 6def extract_masked_region(img: np.ndarray,7                          msk: np.ndarray,8                          threshold: int = 127,9                          crop: bool = True) -> np.ndarray:10    # Binarize mask11    _, bin_mask = cv2.threshold(msk, threshold, 255, cv2.THRESH_BINARY)12    # Apply mask13    result = cv2.bitwise_and(img, img, mask=bin_mask)14    if crop:15        ys, xs = np.where(bin_mask == 255)16        if ys.size and xs.size:17            y1, y2 = ys.min(), ys.max()18            x1, x2 = xs.min(), xs.max()19            result = result[y1:y2+1, x1:x2+1]20    return result21 22def batch_process(collage_dir: str,23                  mask_dir:    str,24                  out_dir:     str):25    os.makedirs(out_dir, exist_ok=True)26 27    files = [f for f in os.listdir(collage_dir)28             if f.lower().endswith(('.png','.jpg','.jpeg','.bmp','tif','tiff'))]29    for fname in tqdm(files, desc="Processing images"):30        img_path  = os.path.join(collage_dir, fname)31        mask_path = os.path.join(mask_dir,    fname)32        out_path  = os.path.join(out_dir,     fname)33 34        if not os.path.isfile(mask_path):35            tqdm.write(f"⚠️  mask not found for {fname}, skipping")36            continue37 38        img = cv2.imread(img_path, cv2.IMREAD_COLOR)39        msk = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)40        if img is None or msk is None:41            tqdm.write(f"⚠️  failed to load {fname} or its mask, skipping")42            continue43 44        cropped = extract_masked_region(img, msk, threshold=127, crop=True)45        cv2.imwrite(out_path, cropped)46 47    tqdm.write("✅ All done!")48 49if __name__ == "__main__":50    # adjust these paths as needed:51    collage_folder = "/home/darth/#/WaterMeters/images"52    mask_folder    = "/home/darth/#/WaterMeters/masks"53    output_folder  = "/home/darth/#/WaterMeters/cropped"54 55    batch_process(collage_folder, mask_folder, output_folder)56 57 58# import cv259# import numpy as np60 61# def extract_masked_region(image_path: str,62#                           mask_path: str,63#                           threshold: int = 127,64#                           crop: bool = True):65#     """66#     Loads an image and its mask, applies the mask, and returns the resulting image.67#     If crop=True, it also crops to the bounding box of the white region in the mask.68 69#     :param image_path:   Path to the original BGR image.70#     :param mask_path:    Path to the grayscale mask (white=keep, black=discard).71#     :param threshold:    Grayscale threshold to binarize the mask (0–255).72#     :param crop:         If True, crop to the mask's bounding box.73#     :return:             A BGR image with only the masked region (and optionally cropped).74#     """75#     # 1. Load images76#     img  = cv2.imread(image_path, cv2.IMREAD_COLOR)77#     msk  = cv2.imread(mask_path,  cv2.IMREAD_GRAYSCALE)78#     if img is None or msk is None:79#         raise FileNotFoundError("Could not load image or mask. Check your paths.")80 81#     # 2. Binarize mask82#     _, bin_mask = cv2.threshold(msk, threshold, 255, cv2.THRESH_BINARY)83 84#     # 3. Apply mask85#     result = cv2.bitwise_and(img, img, mask=bin_mask)86 87#     if crop:88#         # 4. Find bounding box of white region89#         ys, xs = np.where(bin_mask == 255)90#         if ys.size and xs.size:91#             y1, y2 = ys.min(), ys.max()92#             x1, x2 = xs.min(), xs.max()93#             result = result[y1:y2+1, x1:x2+1]94 95#     return result96 97# if __name__ == "__main__":98#     out = extract_masked_region("/home/darth/#/WaterMeters/collage/id_1_value_13_116.jpg", "/home/darth/#/WaterMeters/masks/id_1_value_13_116.jpg")99#     cv2.imwrite("meter_extracted.png", out)100#     print("Saved extracted region to meter_extracted.png")101