Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| import gradio as gr | |
| from PIL import Image | |
| import torch | |
| import cv2 | |
| from utils.image_utils import preprocess_image, count_instance | |
| from utils.data_utils import get_all_file_paths, generate_uuid | |
| import numpy as np | |
| import os | |
| def load_model(): | |
| """ | |
| Load the YOLO model with pre-trained weights. | |
| Returns: | |
| model: Loaded YOLO model. | |
| """ | |
| return YOLO('./weights/best.pt') | |
| def detect_pattern(image_path): | |
| """ | |
| Detect concrete cracks in the binary image. | |
| Parameters: | |
| image_path (str): Path to the binary image. | |
| Returns: | |
| tuple: Principal orientation and orientation category. | |
| """ | |
| image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) | |
| skeleton = cv2.erode(image, np.ones((3, 3), dtype=np.uint8), iterations=1) | |
| contours, _ = cv2.findContours(skeleton, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| data_pts = np.vstack([contour.squeeze() for contour in contours]) | |
| mean, eigenvectors = cv2.PCACompute(data_pts.astype(np.float32), mean=None) | |
| principal_orientation = np.arctan2(eigenvectors[0, 1], eigenvectors[0, 0]) | |
| if -0.05 <= principal_orientation <= 0.05: | |
| orientation_category = "Horizontal" | |
| elif 1 <= principal_orientation <= 1.8: | |
| orientation_category = "Vertical" | |
| elif -0.99 <= principal_orientation <= 0.99: | |
| orientation_category = "Diagonal" | |
| else: | |
| orientation_category = "Other" | |
| return principal_orientation, orientation_category | |
| def predict_segmentation(image, conf): | |
| """ | |
| Perform segmentation prediction on a list of images. | |
| Parameters: | |
| image (list): List of images for segmentation. | |
| conf (float): Confidence score for prediction. | |
| Returns: | |
| tuple: Paths of the processed images, CSV file, DataFrame, and Markdown. | |
| """ | |
| uuid = generate_uuid() | |
| image_list = [preprocess_image(Image.open(file.name)) for file in image] | |
| filenames = [file.name for file in image] | |
| conf= conf * 0.01 | |
| model = load_model() | |
| results = model.predict(image_list, conf=conf, save=True, project='output', name=uuid, stream=True) | |
| processed_image_paths = [] | |
| annotated_image_paths = [] | |
| for i, r in enumerate(results): | |
| for m in r: | |
| masks = r.masks.data | |
| boxes = r.boxes.data | |
| clss = boxes[:, 5] | |
| people_indices = torch.where(clss == 0) | |
| people_masks = masks[people_indices] | |
| people_mask = torch.any(people_masks, dim=0).int() * 255 | |
| processed_image_path = str(model.predictor.save_dir / f'binarize{i}.jpg') | |
| cv2.imwrite(processed_image_path, people_mask.cpu().numpy()) | |
| processed_image_paths.append(processed_image_path) | |
| crack_image_path = processed_image_path | |
| principal_orientation, orientation_category = detect_pattern(crack_image_path) | |
| # Print the results if needed | |
| print(f"Crack Detection Results for {crack_image_path}:") | |
| print("Principal Component Analysis Orientation:", principal_orientation) | |
| print("Orientation Category:", orientation_category) | |
| csv, df = count_instance(results, filenames, uuid) | |
| csv = gr.File(value=csv, visible=True) | |
| df = gr.DataFrame(value=df, visible=True) | |
| md = gr.Markdown(visible=True) | |
| # # Delete binarized images after processing | |
| # for path in processed_image_paths: | |
| # if os.path.exists(path): | |
| # os.remove(path) | |
| return get_all_file_paths(f'output/{uuid}'), csv, df, md | |