Spaces:
Sleeping
Sleeping
| import os | |
| import random | |
| import string | |
| import io | |
| import joblib | |
| import cv2 | |
| import numpy as np | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| from tensorflow.keras import layers, models | |
| from tensorflow.keras.applications import ResNet50, EfficientNetB0 | |
| from tensorflow.keras.applications.resnet import preprocess_input | |
| from tensorflow.keras.applications.efficientnet import preprocess_input | |
| from tensorflow.keras.layers import Lambda # Đảm bảo nhập Lambda từ tensorflow.keras.layers | |
| from keras.applications import ResNet50 | |
| from sklearn.preprocessing import StandardScaler | |
| from sklearn.metrics import ( | |
| confusion_matrix, | |
| ConfusionMatrixDisplay, | |
| accuracy_score, | |
| precision_score, | |
| recall_score, | |
| f1_score | |
| ) | |
| import matplotlib.pyplot as plt | |
| from skimage.feature import graycomatrix, graycoprops | |
| from PIL import Image | |
| import streamlit as st | |
| import cloudinary | |
| import cloudinary.uploader | |
| from cloudinary.utils import cloudinary_url | |
| import torch | |
| # Cloudinary Configuration | |
| cloudinary.config( | |
| cloud_name = os.getenv("CLOUD"), | |
| api_key = os.getenv("API"), | |
| api_secret = os.getenv("SECRET"), | |
| secure=True | |
| ) | |
| # Set page config | |
| st.set_page_config( | |
| page_title="Stone Detection & Classification", | |
| page_icon="🪨", | |
| layout="wide" | |
| ) | |
| def generate_random_filename(extension="png"): | |
| random_string = ''.join(random.choices(string.ascii_letters + string.digits, k=8)) | |
| return f"temp_image_{random_string}.{extension}" | |
| # Custom CSS to improve the appearance | |
| st.markdown(""" | |
| <style> | |
| .main { | |
| padding: 2rem; | |
| } | |
| .stButton>button { | |
| width: 100%; | |
| margin-top: 1rem; | |
| } | |
| .upload-text { | |
| text-align: center; | |
| padding: 2rem; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| def upload_to_cloudinary(file_path, label): | |
| """ | |
| Upload file to Cloudinary with specified label as folder | |
| """ | |
| try: | |
| # Upload to Cloudinary | |
| upload_result = cloudinary.uploader.upload( | |
| file_path, | |
| folder=label, | |
| public_id=f"{label}_{os.path.basename(file_path)}" | |
| ) | |
| # Generate optimized URLs | |
| optimize_url, _ = cloudinary_url( | |
| upload_result['public_id'], | |
| fetch_format="auto", | |
| quality="auto" | |
| ) | |
| auto_crop_url, _ = cloudinary_url( | |
| upload_result['public_id'], | |
| width=500, | |
| height=500, | |
| crop="auto", | |
| gravity="auto" | |
| ) | |
| return { | |
| "upload_result": upload_result, | |
| "optimize_url": optimize_url, | |
| "auto_crop_url": auto_crop_url | |
| } | |
| except Exception as e: | |
| return f"Error uploading to Cloudinary: {str(e)}" | |
| def resize_to_square(image): | |
| """Resize image to square while maintaining aspect ratio""" | |
| size = max(image.shape[0], image.shape[1]) | |
| new_img = np.zeros((size, size, 3), dtype=np.uint8) | |
| # Calculate position to paste original image | |
| x_center = (size - image.shape[1]) // 2 | |
| y_center = (size - image.shape[0]) // 2 | |
| # Copy the image into center of result image | |
| new_img[y_center:y_center+image.shape[0], | |
| x_center:x_center+image.shape[1]] = image | |
| return new_img | |
| def color_histogram(image, bins=16): | |
| """ | |
| Tính histogram màu cho ảnh RGB | |
| Args: | |
| image (np.ndarray): Ảnh đầu vào | |
| bins (int): Số lượng bins của histogram | |
| Returns: | |
| np.ndarray: Histogram màu được chuẩn hóa | |
| """ | |
| # Kiểm tra và chuyển đổi ảnh | |
| if image is None or image.size == 0: | |
| raise ValueError("Ảnh không hợp lệ") | |
| # Đảm bảo ảnh ở dạng uint8 | |
| if image.dtype != np.uint8: | |
| image = (image * 255).astype(np.uint8) | |
| # Tính histogram cho từng kênh màu | |
| hist_r = cv2.calcHist([image], [0], None, [bins], [0, 256]).flatten() | |
| hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten() | |
| hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten() | |
| # Chuẩn hóa histogram | |
| hist_r = hist_r / np.sum(hist_r) if np.sum(hist_r) > 0 else hist_r | |
| hist_g = hist_g / np.sum(hist_g) if np.sum(hist_g) > 0 else hist_g | |
| hist_b = hist_b / np.sum(hist_b) if np.sum(hist_b) > 0 else hist_b | |
| return np.concatenate([hist_r, hist_g, hist_b]) | |
| def color_moments(image): | |
| """ | |
| Tính các moment màu cho ảnh | |
| Args: | |
| image (np.ndarray): Ảnh đầu vào | |
| Returns: | |
| np.ndarray: Các moment màu | |
| """ | |
| # Kiểm tra và chuyển đổi ảnh | |
| if image is None or image.size == 0: | |
| raise ValueError("Ảnh không hợp lệ") | |
| # Đảm bảo ảnh ở dạng float và chuẩn hóa | |
| img = image.astype(np.float32) / 255.0 if image.max() > 1 else image.astype(np.float32) | |
| moments = [] | |
| for i in range(3): # Cho mỗi kênh màu | |
| channel = img[:,:,i] | |
| # Tính các moment | |
| mean = np.mean(channel) | |
| std = np.std(channel) | |
| skewness = np.mean(((channel - mean) / (std + 1e-8)) ** 3) | |
| moments.extend([mean, std, skewness]) | |
| return np.array(moments) | |
| def dominant_color_descriptor(image, k=3): | |
| """ | |
| Xác định các màu chính thống trị trong ảnh | |
| Args: | |
| image (np.ndarray): Ảnh đầu vào | |
| k (int): Số lượng màu chủ đạo | |
| Returns: | |
| np.ndarray: Các màu chủ đạo và tỷ lệ | |
| """ | |
| # Kiểm tra và chuyển đổi ảnh | |
| if image is None or image.size == 0: | |
| raise ValueError("Ảnh không hợp lệ") | |
| # Đảm bảo ảnh ở dạng uint8 | |
| if image.dtype != np.uint8: | |
| image = (image * 255).astype(np.uint8) | |
| # Reshape ảnh thành mảng pixel | |
| pixels = image.reshape(-1, 3) | |
| # Các tham số cho K-means | |
| criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2) | |
| flags = cv2.KMEANS_RANDOM_CENTERS | |
| try: | |
| # Thực hiện phân cụm K-means | |
| _, labels, centers = cv2.kmeans( | |
| pixels.astype(np.float32), k, None, criteria, 10, flags | |
| ) | |
| # Tính toán số lượng và tỷ lệ của từng cụm | |
| unique, counts = np.unique(labels, return_counts=True) | |
| percentages = counts / len(labels) | |
| # Kết hợp các màu và tỷ lệ | |
| dominant_colors = centers.flatten() | |
| color_percentages = percentages | |
| return np.concatenate([dominant_colors, color_percentages]) | |
| except Exception: | |
| # Trả về mảng 0 nếu có lỗi | |
| return np.zeros(2 * k) | |
| def color_coherence_vector(image, k=3): | |
| """ | |
| Tính vector liên kết màu | |
| Args: | |
| image (np.ndarray): Ảnh đầu vào | |
| k (int): Số lượng vùng | |
| Returns: | |
| np.ndarray: Vector liên kết màu | |
| """ | |
| # Kiểm tra và chuyển đổi ảnh | |
| if image is None or image.size == 0: | |
| raise ValueError("Ảnh không hợp lệ") | |
| # Chuyển sang ảnh xám | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| # Đảm bảo ảnh ở dạng uint8 | |
| if gray.dtype != np.uint8: | |
| gray = np.uint8(gray) | |
| # Áp dụng Otsu's thresholding | |
| _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| # Phân tích thành phần liên thông | |
| num_labels, labels = cv2.connectedComponents(binary) | |
| ccv = [] | |
| for i in range(1, min(k+1, num_labels)): | |
| region_mask = (labels == i) | |
| total_pixels = np.sum(region_mask) | |
| coherent_pixels = total_pixels | |
| ccv.extend([coherent_pixels, total_pixels]) | |
| # Đảm bảo độ dài vector | |
| while len(ccv) < 2 * k: | |
| ccv.append(0) | |
| return np.array(ccv) | |
| def edge_features(image, bins=16): | |
| """ | |
| Trích xuất đặc trưng cạnh từ ảnh | |
| Args: | |
| image (np.ndarray): Ảnh đầu vào | |
| bins (int): Số lượng bins của histogram | |
| Returns: | |
| np.ndarray: Đặc trưng cạnh | |
| """ | |
| # Kiểm tra và chuyển đổi ảnh | |
| if image is None or image.size == 0: | |
| raise ValueError("Ảnh không hợp lệ") | |
| # Chuyển sang ảnh xám | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| # Đảm bảo ảnh ở dạng uint8 | |
| if gray.dtype != np.uint8: | |
| gray = np.uint8(gray) | |
| # Tính Sobel edges | |
| sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3) | |
| sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3) | |
| sobel_mag = np.sqrt(sobel_x**2 + sobel_y**2) | |
| # Chuẩn hóa độ lớn Sobel | |
| sobel_mag = np.uint8(255 * sobel_mag / np.max(sobel_mag)) | |
| # Tính histogram của Sobel magnitude | |
| sobel_hist = cv2.calcHist([sobel_mag], [0], None, [bins], [0, 256]).flatten() | |
| sobel_hist = sobel_hist / np.sum(sobel_hist) if np.sum(sobel_hist) > 0 else sobel_hist | |
| # Tính mật độ cạnh bằng Canny | |
| canny_edges = cv2.Canny(gray, 100, 200) | |
| edge_density = np.sum(canny_edges) / (gray.shape[0] * gray.shape[1]) | |
| return np.concatenate([sobel_hist, [edge_density]]) | |
| import pywt # Thư viện xử lý wavelet | |
| def histogram_in_color_space(image, color_space='HSV', bins=16): | |
| """ | |
| Tính histogram của ảnh trong một không gian màu mới. | |
| """ | |
| if color_space == 'HSV': | |
| converted = cv2.cvtColor(image, cv2.COLOR_RGB2HSV) | |
| elif color_space == 'LAB': | |
| converted = cv2.cvtColor(image, cv2.COLOR_RGB2Lab) | |
| else: | |
| raise ValueError("Unsupported color space") | |
| histograms = [] | |
| for i in range(3): # 3 kênh màu | |
| hist = cv2.calcHist([converted], [i], None, [bins], [0, 256]).flatten() | |
| hist = hist / np.sum(hist) | |
| histograms.append(hist) | |
| return np.concatenate(histograms) | |
| def glcm_features(image, distances=[1, 2, 3], angles=[0, np.pi/4, np.pi/2, 3*np.pi/4], levels=256): | |
| """ | |
| Tính các đặc trưng GLCM của ảnh grayscale. | |
| """ | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| # Đảm bảo ảnh ở dạng uint8 | |
| if gray.dtype != np.uint8: | |
| gray = (gray * 255).astype(np.uint8) | |
| glcm = graycomatrix(gray, distances=distances, angles=angles, levels=levels, symmetric=True, normed=True) | |
| features = [] | |
| # Các thuộc tính phổ biến: contrast, homogeneity, energy, correlation | |
| for prop in ['contrast', 'homogeneity', 'energy', 'correlation']: | |
| features.extend(graycoprops(glcm, prop).flatten()) | |
| return np.array(features) | |
| def gabor_features(image, kernels=None): | |
| """ | |
| Tính các đặc trưng từ bộ lọc Gabor. | |
| """ | |
| if kernels is None: | |
| kernels = [] | |
| for theta in np.arange(0, np.pi, np.pi / 4): # Các góc từ 0 đến 180 độ | |
| for sigma in [1, 3]: # Các giá trị sigma | |
| for frequency in [0.1, 0.5]: # Các tần số | |
| kernel = cv2.getGaborKernel((9, 9), sigma, theta, 1/frequency, gamma=0.5, ktype=cv2.CV_32F) | |
| kernels.append(kernel) | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| features = [] | |
| for kernel in kernels: | |
| filtered = cv2.filter2D(gray, cv2.CV_32F, kernel) | |
| features.append(filtered.mean()) | |
| features.append(filtered.var()) | |
| return np.array(features) | |
| def wavelet_features(image, wavelet='db1', level=3): | |
| """ | |
| Trích xuất các hệ số wavelet từ ảnh grayscale. | |
| """ | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| coeffs = pywt.wavedec2(gray, wavelet, level=level) | |
| features = [] | |
| for coeff in coeffs: | |
| if isinstance(coeff, tuple): # Chi tiết (LH, HL, HH) | |
| for subband in coeff: | |
| features.append(subband.mean()) | |
| features.append(subband.var()) | |
| else: # Xấp xỉ (LL) | |
| features.append(coeff.mean()) | |
| features.append(coeff.var()) | |
| return np.array(features) | |
| from skimage.feature import local_binary_pattern | |
| from skimage.color import rgb2gray | |
| from skimage.measure import shannon_entropy | |
| from skimage.feature import hog | |
| def illumination_features(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| mean_brightness = np.mean(gray) | |
| contrast = np.std(gray) | |
| return np.array([mean_brightness, contrast]) | |
| def saturation_index(image): | |
| hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV) | |
| s_channel = hsv[:, :, 1] | |
| mean_saturation = np.mean(s_channel) | |
| std_saturation = np.std(s_channel) | |
| return np.array([mean_saturation, std_saturation]) | |
| def local_binary_pattern_features(image, num_points=24, radius=3): | |
| gray = rgb2gray(image) | |
| lbp = local_binary_pattern(gray, num_points, radius, method="uniform") | |
| hist, _ = np.histogram(lbp.ravel(), bins=np.arange(0, num_points + 3), range=(0, num_points + 2)) | |
| hist = hist / np.sum(hist) | |
| return hist | |
| def fourier_transform_features(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| f_transform = np.fft.fft2(gray) | |
| f_shift = np.fft.fftshift(f_transform) | |
| magnitude_spectrum = 20 * np.log(np.abs(f_shift) + 1) | |
| mean_frequency = np.mean(magnitude_spectrum) | |
| std_frequency = np.std(magnitude_spectrum) | |
| return np.array([mean_frequency, std_frequency]) | |
| def fractal_dimension(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| size = gray.shape[0] * gray.shape[1] | |
| edges = cv2.Canny(gray, 100, 200) | |
| count = np.sum(edges > 0) | |
| fractal_dim = np.log(count + 1) / np.log(size) | |
| return np.array([fractal_dim]) | |
| def glossiness_index(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| glossiness = np.mean(gray[binary == 255]) | |
| return np.array([glossiness]) | |
| def histogram_oriented_gradients(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| features, _ = hog(gray, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2), visualize=True) | |
| return features | |
| def color_entropy(image): | |
| entropy = shannon_entropy(image) | |
| return np.array([entropy]) | |
| def spatial_color_distribution(image, grid_size=4): | |
| h, w, _ = image.shape | |
| features = [] | |
| for i in range(grid_size): | |
| for j in range(grid_size): | |
| x_start = i * h // grid_size | |
| x_end = (i + 1) * h // grid_size | |
| y_start = j * w // grid_size | |
| y_end = (j + 1) * w // grid_size | |
| patch = image[x_start:x_end, y_start:y_end] | |
| hist = cv2.calcHist([patch], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256]).flatten() | |
| hist = hist / np.sum(hist) | |
| features.extend(hist) | |
| return np.array(features) | |
| def uniform_region_features(image): | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| num_labels, labels = cv2.connectedComponents(gray) | |
| unique, counts = np.unique(labels, return_counts=True) | |
| uniformity = np.sum((counts / np.sum(counts)) ** 2) | |
| return np.array([uniformity]) | |
| def color_space_features(image): | |
| ycbcr = cv2.cvtColor(image, cv2.COLOR_RGB2YCrCb) | |
| ycbcr_hist = cv2.calcHist([ycbcr], [1, 2], None, [16, 16], [0, 256, 0, 256]).flatten() | |
| lab = cv2.cvtColor(image, cv2.COLOR_RGB2Lab) | |
| lab_hist = cv2.calcHist([lab], [1, 2], None, [16, 16], [0, 256, 0, 256]).flatten() | |
| ycbcr_hist = ycbcr_hist / np.sum(ycbcr_hist) | |
| lab_hist = lab_hist / np.sum(lab_hist) | |
| return np.concatenate([ycbcr_hist, lab_hist]) | |
| def extract_features(image): | |
| color_hist = color_histogram(image) | |
| color_mom = color_moments(image) | |
| dom_color = dominant_color_descriptor(image) | |
| ccv = color_coherence_vector(image) | |
| edges = edge_features(image) | |
| hsv_hist = histogram_in_color_space(image, color_space='HSV') | |
| glcm = glcm_features(image) | |
| gabor = gabor_features(image) | |
| wavelet = wavelet_features(image) | |
| illumination = illumination_features(image) | |
| saturation = saturation_index(image) | |
| lbp = local_binary_pattern_features(image) | |
| fourier = fourier_transform_features(image) | |
| fractal = fractal_dimension(image) | |
| return np.concatenate([ | |
| color_hist, | |
| color_mom, | |
| dom_color, | |
| ccv, | |
| edges, | |
| hsv_hist, | |
| glcm, | |
| gabor, | |
| wavelet, | |
| illumination, | |
| saturation, | |
| lbp, | |
| fourier, | |
| fractal, | |
| ]) | |
| def load_models(): | |
| """Load both object detection and classification models""" | |
| # Load object detection model | |
| device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') | |
| object_detection_model = torch.load("fasterrcnn_resnet50_fpn_090824.pth", map_location=device) | |
| object_detection_model.to(device) | |
| object_detection_model.eval() | |
| # Load classification model | |
| classification_model = tf.keras.models.load_model('mlp_model.h5') | |
| return object_detection_model, classification_model, device | |
| def create_efficientnetb0_feature_extractor(input_shape=(256, 256, 3), num_classes=None): | |
| # Xây dựng mô hình EfficientNetB0 đã huấn luyện sẵn từ TensorFlow | |
| inputs = layers.Input(shape=input_shape) | |
| # Thêm lớp Lambda để tiền xử lý ảnh | |
| x = Lambda(preprocess_input, output_shape=input_shape)(inputs) # Xử lý ảnh đầu vào | |
| # Sử dụng mô hình EfficientNetB0 đã được huấn luyện sẵn | |
| efficientnetb0_model = EfficientNetB0(include_top=False, weights='imagenet', input_tensor=x) | |
| # Trích xuất đặc trưng từ mô hình EfficientNetB0 | |
| x = layers.GlobalAveragePooling2D()(efficientnetb0_model.output) | |
| if num_classes: | |
| x = layers.Dense(num_classes, activation='softmax')(x) # Thêm lớp phân loại (nếu có) | |
| return models.Model(inputs=inputs, outputs=x) | |
| def perform_object_detection(image, model, device): | |
| original_size = image.size | |
| target_size = (256, 256) | |
| frame_resized = cv2.resize(np.array(image), dsize=target_size, interpolation=cv2.INTER_AREA) | |
| frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_RGB2BGR).astype(np.float32) | |
| frame_rgb /= 255.0 | |
| frame_rgb = frame_rgb.transpose(2, 0, 1) | |
| frame_rgb = torch.from_numpy(frame_rgb).float().unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| outputs = model(frame_rgb) | |
| boxes = outputs[0]['boxes'].cpu().detach().numpy().astype(np.int32) | |
| labels = outputs[0]['labels'].cpu().detach().numpy().astype(np.int32) | |
| scores = outputs[0]['scores'].cpu().detach().numpy() | |
| result_image = frame_resized.copy() | |
| cropped_images = [] | |
| detected_boxes = [] | |
| for i in range(len(boxes)): | |
| if scores[i] >= 0.75: | |
| x1, y1, x2, y2 = boxes[i] | |
| if (int(labels[i])-1) == 1 or (int(labels[i])-1) == 0: | |
| color = (0, 255, 0) # Green bounding box | |
| label_text = f'Region {i}' | |
| # Scale coordinates to original image size | |
| original_h, original_w = original_size[::-1] | |
| scale_h, scale_w = original_h / target_size[0], original_w / target_size[1] | |
| x1_orig, y1_orig = int(x1 * scale_w), int(y1 * scale_h) | |
| x2_orig, y2_orig = int(x2 * scale_w), int(y2 * scale_h) | |
| # Crop and process detected region | |
| cropped_image = np.array(image)[y1_orig:y2_orig, x1_orig:x2_orig] | |
| # Check if image has 4 channels (RGBA), convert to RGB | |
| if cropped_image.shape[-1] == 4: | |
| cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_RGBA2RGB) | |
| else: | |
| cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB) | |
| # Resize cropped image | |
| resized_crop = resize_to_square(cropped_image) | |
| cropped_images.append([i,resized_crop]) | |
| detected_boxes.append((x1, y1, x2, y2)) | |
| # Draw bounding box | |
| cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 1) | |
| cv2.putText(result_image, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1, cv2.LINE_AA) # Yellow text, smaller font | |
| return Image.fromarray(result_image), cropped_images, detected_boxes | |
| def preprocess_image(image): | |
| """Preprocess the image for classification""" | |
| # Convert image to numpy array and resize | |
| img_array = np.array(image) | |
| img_array = cv2.resize(img_array, (256, 256)) | |
| # Extract custom features (ensure this returns a 1D array) | |
| features = extract_features(img_array) | |
| features = features.flatten() # Ensure 1D | |
| # Extract EfficientNet features | |
| model_extractor = create_efficientnetb0_feature_extractor() | |
| model_features = model_extractor.predict(np.expand_dims(img_array, axis=0)) | |
| model_features = model_features.flatten() # Convert to 1D array | |
| # Combine features | |
| features_combined = np.concatenate([features, model_features]) | |
| features_combined = features_combined.reshape(1, -1) # Reshape to 2D for scaler | |
| # Load and apply scaler | |
| scaler = joblib.load('scaler.pkl') | |
| processed_image = scaler.transform(features_combined) | |
| return processed_image | |
| def get_top_predictions(prediction, class_names, top_k=5): | |
| """Get top k predictions with their probabilities""" | |
| top_indices = prediction.argsort()[0][-top_k:][::-1] | |
| top_predictions = [ | |
| (class_names[i], float(prediction[0][i]) * 100) | |
| for i in top_indices | |
| ] | |
| return top_predictions | |
| def main(): | |
| st.title("🪨 Stone Detection & Classification") | |
| st.write("Upload an image to detect and classify stone surfaces") | |
| if 'predictions' not in st.session_state: | |
| st.session_state.predictions = None | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.subheader("Upload Image") | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| with st.spinner('Processing image...'): | |
| try: | |
| # Load both models | |
| object_detection_model, classification_model, device = load_models() | |
| # Perform object detection | |
| result_image, cropped_images, detected_boxes = perform_object_detection( | |
| image, object_detection_model, device | |
| ) | |
| if not cropped_images: | |
| st.warning("No stone surfaces detected in the image") | |
| return | |
| # Display detection results | |
| st.subheader("Detection Results") | |
| st.image(result_image, caption="Detected Stone Surfaces", use_column_width=True) | |
| # Process each detected region | |
| class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7'] | |
| all_predictions = [] | |
| all_image=[] | |
| for idx, cropped_image in cropped_images: | |
| processed_image = preprocess_image(cropped_image) | |
| prediction = classification_model.predict(processed_image) | |
| top_predictions = get_top_predictions(prediction, class_names) | |
| all_predictions.append([idx,top_predictions]) | |
| all_image.append(cropped_image) | |
| # Store in session state | |
| st.session_state.predictions = all_predictions | |
| st.session_state.image = all_image | |
| except Exception as e: | |
| st.error(f"Error during processing: {str(e)}") | |
| with col2: | |
| st.subheader("Classification Results") | |
| if st.session_state.predictions is not None: | |
| for idx, predictions in st.session_state.predictions: | |
| st.markdown(f"### Region {idx}") | |
| st.image(st.session_state.image[idx], use_column_width=True) | |
| # Display main prediction | |
| top_class, top_confidence = predictions[0] | |
| st.markdown(f"**Primary Prediction: Grade {top_class}**") | |
| st.markdown(f"**Confidence: {top_confidence:.2f}%**") | |
| st.progress(top_confidence / 100) | |
| # Display all predictions for this region | |
| st.markdown("**Top 5 Predictions**") | |
| for class_name, confidence in predictions: | |
| col_label, col_bar, col_value = st.columns([2, 6, 2]) | |
| with col_label: | |
| st.write(f"Grade {class_name}") | |
| with col_bar: | |
| st.progress(confidence / 100) | |
| with col_value: | |
| st.write(f"{confidence:.2f}%") | |
| st.markdown("---") | |
| st.markdown("</div>", unsafe_allow_html=True) | |
| # User Confirmation Section | |
| st.markdown("### Xác nhận độ chính xác của mô hình") | |
| st.write("Giúp chúng tôi cải thiện mô hình bằng cách xác nhận độ chính xác của dự đoán.") | |
| # Accuracy Radio Button | |
| accuracy_option = st.radio( | |
| "Dự đoán có chính xác không?", | |
| ["Chọn", "Chính xác", "Không chính xác"], | |
| index=0, | |
| key=f"accuracy_radio_{idx}" | |
| ) | |
| if accuracy_option == "Không chính xác": | |
| # Input for correct grade | |
| correct_grade = st.selectbox( | |
| "Chọn màu đá đúng:", | |
| ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7'], | |
| index=None, | |
| placeholder="Chọn màu đúng", | |
| key=f"selectbox_correct_grade_{idx}" | |
| ) | |
| # Kiểm tra xem đã tải lên hay chưa | |
| if f"uploaded_{idx}" not in st.session_state: | |
| st.session_state[f"uploaded_{idx}"] = False | |
| # Chỉ thực hiện khi người dùng đã chọn giá trị và chưa tải lên | |
| if correct_grade and not st.session_state[f"uploaded_{idx}"]: | |
| st.info(f"Đã chọn màu đúng: {correct_grade}") | |
| # Resize hình ảnh xuống 256x256 | |
| resized_image = Image.fromarray(st.session_state.image[idx]).resize((256, 256)) | |
| temp_image_path = generate_random_filename() | |
| # Lưu tệp resize tạm thời | |
| resized_image.save(temp_image_path) | |
| # Tải ảnh lên Cloudinary | |
| cloudinary_result = upload_to_cloudinary(temp_image_path, correct_grade) | |
| if isinstance(cloudinary_result, dict): | |
| st.success(f"Hình ảnh đã được tải lên thành công cho màu {correct_grade}") | |
| st.write(f"URL công khai: {cloudinary_result['upload_result']['secure_url']}") | |
| # Đánh dấu trạng thái đã tải lên | |
| st.session_state[f"uploaded_{idx}"] = True | |
| else: | |
| st.error(cloudinary_result) | |
| else: | |
| st.info("Upload an image to see detection and classification results") | |
| if __name__ == "__main__": | |
| main() |