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() |