Spaces:
Sleeping
Sleeping
File size: 3,581 Bytes
50d6ddc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
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
|