{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":11171890,"sourceType":"datasetVersion","datasetId":6972203},{"sourceId":11193440,"sourceType":"datasetVersion","datasetId":6987935},{"sourceId":11193880,"sourceType":"datasetVersion","datasetId":6988268},{"sourceId":9975392,"sourceType":"datasetVersion","datasetId":6132251},{"sourceId":3848,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":2749,"modelId":324}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Segment Anything Model (SAM)\nResearch by Meta AI\n\nhttps://segment-anything.com/\n> SAM is a promptable segmentation system with zero-shot generalization to unfamiliar objects and images, without the need for additional training.\n\nhttps://github.com/facebookresearch/segment-anything\n\n> The Segment Anything Model (SAM) produces high quality object masks from input prompts such as points or boxes, and it can be used to generate masks for all objects in an image. It has been trained on a dataset of 11 million images and 1.1 billion masks, and has strong zero-shot performance on a variety of segmentation tasks.","metadata":{}},{"cell_type":"code","source":"!python --version","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:09:45.851336Z","iopub.execute_input":"2025-03-28T10:09:45.851665Z","iopub.status.idle":"2025-03-28T10:09:45.976651Z","shell.execute_reply.started":"2025-03-28T10:09:45.851634Z","shell.execute_reply":"2025-03-28T10:09:45.975652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:09:45.977970Z","iopub.execute_input":"2025-03-28T10:09:45.978220Z","iopub.status.idle":"2025-03-28T10:09:51.375119Z","shell.execute_reply.started":"2025-03-28T10:09:45.978197Z","shell.execute_reply":"2025-03-28T10:09:51.374049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nimport os\nfrom matplotlib import pyplot as plt\nimport torch\nimport cv2\nfrom ultralytics import YOLO\nfrom PIL import Image\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:09:51.376902Z","iopub.execute_input":"2025-03-28T10:09:51.377163Z","iopub.status.idle":"2025-03-28T10:09:55.812908Z","shell.execute_reply.started":"2025-03-28T10:09:51.377140Z","shell.execute_reply":"2025-03-28T10:09:55.812242Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install git+https://github.com/facebookresearch/segment-anything.git","metadata":{"execution":{"iopub.status.busy":"2025-03-28T10:09:55.813928Z","iopub.execute_input":"2025-03-28T10:09:55.814357Z","iopub.status.idle":"2025-03-28T10:10:02.175674Z","shell.execute_reply.started":"2025-03-28T10:09:55.814321Z","shell.execute_reply":"2025-03-28T10:10:02.174570Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ModelLocalization:\n    def __init__(self, model_path):\n        self.model = YOLO(model_path)\n        self.classes = [\n            \"door\",\n            \"hinge\",\n            \"knob\",\n            \"lever\",\n            \"window\",\n            \"pocket_handles\",\n            \"opening\",\n        ]\n        self.threshold_door_localization = {\n            \"door\": 0.367936372756958,\n            \"hinge\": 0.25,\n            \"knob\": 0.359599769115448,\n            \"lever\": 0.2923234701156616,\n            \"window\": 0.27619636058807373,\n            \"pocket_handles\": 0.25,\n            \"opening\": 0.25,\n        }\n\n    def get_door_localization_predict(self, im):\n        results = self.model.predict(source=im, conf=0.1, iou=0.5, imgsz=640)\n        bb = results[0]\n\n        box = bb.boxes.xyxy\n        cls = bb.boxes.cls\n        conf = bb.boxes.conf\n\n        return (box.tolist(), cls.tolist(), conf.tolist())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:02.176857Z","iopub.execute_input":"2025-03-28T10:10:02.177235Z","iopub.status.idle":"2025-03-28T10:10:02.183834Z","shell.execute_reply.started":"2025-03-28T10:10:02.177198Z","shell.execute_reply":"2025-03-28T10:10:02.182989Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def show_anns(anns, axes=None):\n    if len(anns) == 0:\n        return\n    if axes:\n        ax = axes\n    else:\n        ax = plt.gca()\n        ax.set_autoscale_on(False)\n    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)\n    polygons = []\n    color = []\n    for ann in sorted_anns:\n        m = ann['segmentation']\n        img = np.ones((m.shape[0], m.shape[1], 3))\n        color_mask = np.random.random((1, 3)).tolist()[0]\n        for i in range(3):\n            img[:,:,i] = color_mask[i]\n        ax.imshow(np.dstack((img, m*0.5)))\n\ndef show_mask(mask, ax, random_color=False):\n    if random_color:\n        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n    else:\n        color = np.array([30/255, 144/255, 255/255, 0.6])\n    h, w = mask.shape[-2:]\n    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n    ax.imshow(mask_image)\n\n    \ndef show_points(coords, labels, ax, marker_size=375):\n    pos_points = coords[labels==1]\n    neg_points = coords[labels==0]\n    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n\n    \ndef show_box(box, ax):\n    x0, y0 = box[0], box[1]\n    w, h = box[2] - box[0], box[3] - box[1]\n    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))    ","metadata":{"execution":{"iopub.status.busy":"2025-03-28T10:10:02.184844Z","iopub.execute_input":"2025-03-28T10:10:02.185105Z","iopub.status.idle":"2025-03-28T10:10:02.212807Z","shell.execute_reply.started":"2025-03-28T10:10:02.185076Z","shell.execute_reply":"2025-03-28T10:10:02.211931Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Automatically Generating Object Masks with SAM\n\nhttps://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb","metadata":{}},{"cell_type":"code","source":"from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor\n\nsam_checkpoint = \"/kaggle/input/segment-anything/pytorch/vit-b/1/model.pth\"\nmodel_type = \"vit_b\"\n\ndevice = \"cuda\"\n\nsam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\nsam.to(device=device)\n\nmask_generator = SamAutomaticMaskGenerator(sam, points_per_batch=16)","metadata":{"execution":{"iopub.status.busy":"2025-03-28T10:10:02.213546Z","iopub.execute_input":"2025-03-28T10:10:02.213767Z","iopub.status.idle":"2025-03-28T10:10:08.658295Z","shell.execute_reply.started":"2025-03-28T10:10:02.213748Z","shell.execute_reply":"2025-03-28T10:10:08.657595Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Object masks from prompts with SAM\n\nhttps://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb","metadata":{}},{"cell_type":"code","source":"from segment_anything import sam_model_registry, SamPredictor\n\npredictor = SamPredictor(sam)","metadata":{"execution":{"iopub.status.busy":"2025-03-28T10:10:08.660209Z","iopub.execute_input":"2025-03-28T10:10:08.660431Z","iopub.status.idle":"2025-03-28T10:10:08.663995Z","shell.execute_reply.started":"2025-03-28T10:10:08.660411Z","shell.execute_reply":"2025-03-28T10:10:08.663176Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Selecting objects with SAM\n","metadata":{}},{"cell_type":"markdown","source":"## Specifying a specific object with a box\n","metadata":{}},{"cell_type":"code","source":"def expand_box(box, img_shape, margin=20):\n    x1, y1, x2, y2 = box\n    x1 = max(0, x1 - margin)\n    y1 = max(0, y1 - margin)\n    x2 = min(img_shape[1] - 1, x2 + margin)\n    y2 = min(img_shape[0] - 1, y2 + margin)\n    return np.array([x1, y1, x2, y2])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:08.664854Z","iopub.execute_input":"2025-03-28T10:10:08.665050Z","iopub.status.idle":"2025-03-28T10:10:08.681353Z","shell.execute_reply.started":"2025-03-28T10:10:08.665023Z","shell.execute_reply":"2025-03-28T10:10:08.680587Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_four_corners(image_path, predictor, detector):\n    img = cv2.imread(image_path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    predictor.set_image(img)\n    boxes, cls, _ = detector.get_door_localization_predict(img)\n    if 0 not in cls:\n        return\n    for i, box_sample in enumerate(boxes):\n        if cls[i]==0:\n            box=box_sample\n    box = np.array(box)\n    #box = expand_box(box, img.shape)\n    masks, _, _ = predictor.predict(\n        point_coords=None,\n        point_labels=None,\n        box=box[None, :],\n        multimask_output=False,\n    )\n    mask = masks[0]\n    binary_mask = (mask.astype(np.uint8))*255\n    # --- Visualize Bounding Box ---\n    cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color=(255, 255, 0), thickness=3)\n\n    # --- Visualize Mask Overlay ---\n    color_mask = np.zeros_like(img)\n    color_mask[mask.astype(bool)] = [0, 255, 0]  # green overlay for mask\n    img = cv2.addWeighted(img, 1.0, color_mask, 0.4, 0)  # Blend\n    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n    contour = max(contours, key=cv2.contourArea)\n    epsilon = 0.01 * cv2.arcLength(contour, True)\n    approx = cv2.approxPolyDP(contour, epsilon, True)\n    \n    if len(approx) == 4:\n        print(True)\n        corners = approx.reshape(4, 2)\n    \n        # Sort: top-left, top-right, bottom-right, bottom-left\n        def sort_corners(pts):\n            pts = sorted(pts, key=lambda x: x[1])  # sort by y\n            top = sorted(pts[:2], key=lambda x: x[0])\n            bottom = sorted(pts[2:], key=lambda x: x[0], reverse=True)\n            return np.array([top[0], top[1], bottom[0], bottom[1]])\n    \n        sorted_corners = sort_corners(corners)\n        \n        # Print corners\n        print(\"Detected door corners (top-left, top-right, bottom-right, bottom-left):\")\n        for i, pt in enumerate(sorted_corners):\n            print(f\"Corner {i+1}: {pt}\")\n        for pt in sorted_corners:\n            cv2.circle(img, tuple(pt), 5, (255, 0, 0), -1)\n            \n        output_folder = \"/kaggle/working/output_images\"\n        os.makedirs(output_folder, exist_ok=True)  # Create folder if it doesn't exist\n        \n        # Save image with a custom name\n        filename = os.path.basename(image_path).split('.')[0] + \"_corners.png\"\n        output_path = os.path.join(output_folder, filename)\n        \n        # Convert RGB back to BGR for saving with OpenCV\n        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n        cv2.imwrite(output_path, img_bgr)\n        \n        print(f\"Saved image with corners to: {output_path}\")\n    else:\n        print(f\"Expected 4 corners, but got {len(approx)}. Adjust `epsilon` or clean up the mask.\") ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:08.682192Z","iopub.execute_input":"2025-03-28T10:10:08.682486Z","iopub.status.idle":"2025-03-28T10:10:08.702015Z","shell.execute_reply.started":"2025-03-28T10:10:08.682457Z","shell.execute_reply":"2025-03-28T10:10:08.701412Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_folder = '/kaggle/input/door-window-img-level-v4-1/door_window_img_level_v4.1/door_window_img_level_v4.1/door_window_img_level_v4.1/test'\nlist_folder = [os.path.join(data_folder, f) for f in os.listdir(data_folder) if os.path.isdir(os.path.join(data_folder, f))]\nlist_data = []\nfor folder in list_folder:\n    list_img_paths = os.listdir(folder)\n    list_img_paths = [os.path.join(folder, filename) for filename in list_img_paths]\n    list_data.extend(list_img_paths)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:08.702718Z","iopub.execute_input":"2025-03-28T10:10:08.702991Z","iopub.status.idle":"2025-03-28T10:10:08.961869Z","shell.execute_reply.started":"2025-03-28T10:10:08.702965Z","shell.execute_reply":"2025-03-28T10:10:08.961142Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"detector = ModelLocalization('/kaggle/input/roomsketcher-yolo-checkpoint/localization_weight_26feb2025.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:08.962561Z","iopub.execute_input":"2025-03-28T10:10:08.962780Z","iopub.status.idle":"2025-03-28T10:10:09.226349Z","shell.execute_reply.started":"2025-03-28T10:10:08.962752Z","shell.execute_reply":"2025-03-28T10:10:09.225394Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for path in list_data:\n    extract_four_corners(path, predictor, detector)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:10:09.227268Z","iopub.execute_input":"2025-03-28T10:10:09.227561Z","iopub.status.idle":"2025-03-28T10:12:40.364723Z","shell.execute_reply.started":"2025-03-28T10:10:09.227527Z","shell.execute_reply":"2025-03-28T10:12:40.364045Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Follow this notebook for additional prompting options:\n\nhttps://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb","metadata":{}},{"cell_type":"code","source":"import shutil\n\n# Folder you want to zip\nfolder_to_zip = \"/kaggle/working/output_images\"\n\n# Output zip file (without extension for make_archive)\nshutil.make_archive(\"/kaggle/working/output_images\", 'zip', folder_to_zip)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:12:40.365637Z","iopub.execute_input":"2025-03-28T10:12:40.365906Z","iopub.status.idle":"2025-03-28T10:12:43.421064Z","shell.execute_reply.started":"2025-03-28T10:12:40.365875Z","shell.execute_reply":"2025-03-28T10:12:43.420177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import supervision as sv\n\n# import cv2\n\n# from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor\n\n# import torch\n\n\n\n# DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n# MODEL_TYPE = \"vit_h\"\n\n\n\n# sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)\n\n\n\n# mask_generator = SamAutomaticMaskGenerator(sam)\n\n\n\n# image_bgr = cv2.imread(IMAGE_PATH)\n# image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n\n\n\n# sam_result = mask_generator.generate(image_rgb)\n\n\n\n# mask_annotator = sv.MaskAnnotator()\n\n\n\n# detections = sv.Detections.from_sam(sam_result=sam_result)\n\n\n\n# annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)\n\n\n\n# sv.plot_images_grid(\n#     images=[image_bgr, annotated_image],\n#     grid_size=(1, 2),\n#     titles=['source image', 'segmented image']\n# )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T10:12:43.422117Z","iopub.execute_input":"2025-03-28T10:12:43.422433Z","iopub.status.idle":"2025-03-28T10:12:43.426158Z","shell.execute_reply.started":"2025-03-28T10:12:43.422401Z","shell.execute_reply":"2025-03-28T10:12:43.425378Z"}},"outputs":[],"execution_count":null}]}