Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import cv2 | |
| from shapely import Polygon | |
| from shapely.wkt import dumps | |
| from shapely.geometry import box | |
| import pandas as pd | |
| def predict_segments(model, image_path, conf = 0.5, classes = 0): | |
| """ Predict segments using YOLOv8 segmentation model for floorplan drawings. | |
| Args: | |
| - model: Local stored segmentation model | |
| - image_path: path to uploaded image | |
| - conf: confidence threshold for the model | |
| - classes: 0 (BRA, area inside exterior walls) or 1 (area including exterior walls) | |
| """ | |
| results = model.predict(image_path, conf=conf, classes = classes, retina_masks=True) | |
| return results | |
| def segmentation_to_binary(results): | |
| """ | |
| Converts | |
| """ | |
| binary_masks = [] | |
| for result in results: | |
| if result.masks is not None: | |
| for mask in result.masks.data: | |
| mask_np = mask.cpu().numpy().astype(np.uint8) * 255 | |
| binary_masks.append(mask_np) | |
| return binary_masks | |
| def segmask_to_pandas(seg_results): | |
| """ | |
| Stores the segmentation mask as polygons in dataframe with id | |
| Args: | |
| - seg_results (tensor object): segmentation results from YOLO | |
| Returns: | |
| - Pandas Dataframe: 'mask_id' column, 'polygon_coord' column with polygon coord. | |
| """ | |
| mask_count = 0 | |
| masks_list = [] | |
| for result in seg_results: | |
| boxes = result.boxes.xyxy.cpu().numpy() | |
| for mask,box in zip(result.masks.xy, boxes): | |
| box = box.astype(int) | |
| mask_pts = mask.astype(int).tolist() | |
| masks_list.append({"mask_id": mask_count,"polygon": mask_pts, 'bboxes': box}) | |
| mask_count +=1 | |
| df = pd.DataFrame(masks_list) | |
| return df | |
| def fill_segments(mask,bboxes): | |
| mask = np.uint8(mask>0)*255 | |
| bbox_mask = np.zeros_like(mask) | |
| for x1,y1,x2,y2 in bboxes: | |
| bbox_mask[y1:y2, x1:x2] | |
| bbox_inv = cv2.bitwise_not(bbox_mask) | |
| kernel = np.ones((3,3),np.uint8) | |
| expanded_mask = cv2.dilate(mask, kernel, iterations=3) | |
| final_mask = cv2.bitwise_and(expanded_mask,bbox_inv) | |
| floodfill_mask = final_mask.copy() | |
| h,w = floodfill_mask.shape | |
| mask_floodfill = np.zeros((h+2,w+2), np.uint8) | |
| cv2.floodFill(floodfill_mask, mask_floodfill, (0,0), 255) | |
| floodfill_mask = cv2.bitwise_not(floodfill_mask) | |
| final_filled_mask = cv2.bitwise_or(floodfill_mask, final_mask) | |
| return final_filled_mask | |
| def find_text_in_segments(masks, ocr_results): | |
| pass |