Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from typing import List | |
| import pandas as pd | |
| def detect_and_scale(image_path, model) -> List[tuple]: | |
| """ | |
| Detect drawingtype (plantegning, fasade, snitt or situasjonskart). | |
| Filter by only plantegning, and scale bboxes to original image size. | |
| Args: | |
| - image_path: path to uploaded image | |
| - model: Local stored object detection model | |
| Returns: | |
| List of tuples with bbox coordinates (x1,y1,x2,y2) | |
| """ | |
| original_image = cv2.imread(image_path) | |
| orig_h, orig_w = original_image.shape[:2] | |
| results = model(image_path) | |
| detected_boxes = [] | |
| for result in results: | |
| boxes = result.boxes | |
| conf_boxes = boxes.conf | |
| class_labels = result.names | |
| yolo_h, yolo_w = result.orig_shape[:2] | |
| scale_x = orig_w / yolo_w | |
| scale_y = orig_h / yolo_h | |
| for box, cls in zip(boxes.xyxy, boxes.cls): | |
| label = class_labels[int(cls)] | |
| if label == "plantegning": | |
| x1, y1, x2, y2 = box.tolist() | |
| x1 = int(x1 * scale_x) | |
| y1 = int(y1 * scale_y) | |
| x2 = int(x2 * scale_x) | |
| y2 = int(y2 * scale_y) | |
| detected_boxes.append((x1, y1, x2, y2)) | |
| #cv2.rectangle(original_image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| return detected_boxes, original_image | |
| def detect_and_plot_plantegning(model, image): | |
| results = model.predict(image, classes=1) | |
| annotated_img = results[0].plot() | |
| return annotated_img | |
| def obj_to_pandas(objdet_results): | |
| detected_obj_list = [] | |
| obj_count = 0 | |
| for box in objdet_results: | |
| obj_count +=1 | |
| detected_obj_list.append({"mask_id": obj_count, 'bboxes': box}) | |
| return pd.DataFrame(detected_obj_list) | |
| # ------------ AREAL MÅLEREN -------------------------- | |
| def detect_resize_walls(image_path, model): | |
| """ Detect and plot YOLO bounding boxes on the original image size """ | |
| original_image = cv2.imread(image_path) | |
| orig_h, orig_w = original_image.shape[:2] | |
| results = model.predict(image_path) | |
| detected_boxes = [] | |
| for result in results: | |
| boxes = result.boxes | |
| class_labels = result.names | |
| yolo_h, yolo_w = result.orig_shape[:2] | |
| scale_x = orig_w / yolo_w | |
| scale_y = orig_h / yolo_h | |
| for box, cls in zip(boxes.xyxy, boxes.cls): | |
| label = class_labels[int(cls)] | |
| if label.lower() == "wall": | |
| x1, y1, x2, y2 = box.tolist() | |
| x1 = int(x1 * scale_x) | |
| y1 = int(y1 * scale_y) | |
| x2 = int(x2 * scale_x) | |
| y2 = int(y2 * scale_y) | |
| detected_boxes.append((x1, y1, x2, y2)) | |
| cv2.rectangle(original_image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| return detected_boxes | |
| def remove_interior_walls(masks, bboxes): | |
| filtered_bboxes = [] | |
| for bbox in bboxes: | |
| x1, y1, x2, y2 = bbox | |
| inside_mask = False | |
| for mask in masks: | |
| mask_region = mask[y1:y2, x1:x2] | |
| mask_coverage = np.sum(mask_region > 0) / (mask_region.shape[0] * mask_region.shape[1]) | |
| if mask_coverage > 0.8: # If 80% or more of the bbox is covered by a mask, remove it | |
| inside_mask = True | |
| break | |
| if not inside_mask: | |
| filtered_bboxes.append(bbox) | |
| return filtered_bboxes | |
| def reference_object(image_path, model, object_type="wall"): | |
| img = cv2.imread(image_path) | |
| results = model.predict(image_path, conf=0.5) | |
| longest_x = 0 | |
| bbox = None | |
| longest_y = 0 | |
| for result in results: | |
| boxes = result.boxes | |
| class_labels = result.names | |
| for box, cls in zip(boxes.xyxy, boxes.cls): | |
| label = class_labels[int(cls)] | |
| if label.lower() == object_type: | |
| x1, y1, x2, y2 = map(int, box.tolist()) | |
| length_x = x2 - x1 | |
| length_y = y2 - y1 | |
| if length_x > longest_x: | |
| longest_x = length_x | |
| bbox = (x1, y1, x2, y2) | |
| if length_y > longest_y: | |
| longest_y = length_y | |
| if bbox: | |
| cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2) | |
| return img, longest_x, longest_y | |