soilclassification / image_processing.py
manny1313's picture
First Commit
e06826c
# image_processing.py
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from skimage.feature import graycomatrix, graycoprops, local_binary_pattern
import joblib
import warnings
warnings.filterwarnings("ignore")
# Load model and preprocessing tools
model = joblib.load("model.pkl")
label_encoder = joblib.load("label_encoder.pkl")
scaler = joblib.load("scaler.pkl")
# Preprocessing: Resize, apply CLAHE, sharpen
def preprocessing(single_image, count=1):
single_image = cv2.resize(single_image, (256,256))
rgb_image = cv2.cvtColor(single_image, cv2.COLOR_BGR2RGB)
r, g, b = cv2.split(rgb_image)
clahe = cv2.createCLAHE(clipLimit=0.4, tileGridSize=(8, 8))
r_clahe = clahe.apply(r)
g_clahe = clahe.apply(g)
b_clahe = clahe.apply(b)
clahe_image = cv2.merge((r_clahe, g_clahe, b_clahe))
clahe_bgr = cv2.cvtColor(clahe_image, cv2.COLOR_RGB2BGR)
blurred = cv2.GaussianBlur(clahe_bgr, (5, 5), 1.5)
sharp = cv2.addWeighted(clahe_bgr, 1.5, blurred, -0.5, 0)
return rgb_image, clahe_image, cv2.cvtColor(sharp, cv2.COLOR_BGR2RGB)
# RGB histogram plotting
def plot_rgb_histogram(image):
color = ('b', 'g', 'r')
fig, ax = plt.subplots()
for i, col in enumerate(color):
hist = cv2.calcHist([image], [i], None, [256], [0,256])
ax.plot(hist, color=col)
ax.set_title("RGB Histogram")
ax.set_xlim([0, 256])
return fig
# Extract features: GLCM, LBP, color, edge, etc.
def feature_extraction(image, return_df=True):
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
glcm = graycomatrix(gray, distances=[1], angles=[0], levels=256, symmetric=True, normed=True)
contrast = graycoprops(glcm, 'contrast')[0, 0]
correlation = graycoprops(glcm, 'correlation')[0, 0]
energy = graycoprops(glcm, 'energy')[0, 0]
homogeneity = graycoprops(glcm, 'homogeneity')[0, 0]
lbp = local_binary_pattern(gray, P=8, R=1, method='uniform')
lbp_mean = np.mean(lbp)
mean_r = np.mean(image[:, :, 0])
mean_g = np.mean(image[:, :, 1])
mean_b = np.mean(image[:, :, 2])
diff_black = ((1-(mean_r-255)/255) + (1-(mean_g-255)/255) + (1-(mean_b-255)/255))/3
sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=5)
sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=5)
edge_count = np.sum(cv2.magnitude(sobelx, sobely) > 0)
features = [contrast, correlation, energy, homogeneity, lbp_mean,
mean_r, mean_g, mean_b, edge_count, diff_black]
if return_df:
df = pd.DataFrame([features], columns=[
"Contrast", "Correlation", "Energy", "Homogeneity", "LBP_Mean",
"Mean_R", "Mean_G", "Mean_B", "Edge_Count", "Black"
])
return df
else:
return np.array([features])
# Predict class and confidence from image
def predict_image_class_with_features(image):
_, _, sharp = preprocessing(image, count=0)
features_df = feature_extraction(sharp)
features_scaled = scaler.transform(features_df)
prediction = model.predict(features_scaled)
predicted_class = label_encoder.inverse_transform(prediction)[0]
confidence = np.max(model.predict_proba(features_scaled))
return features_df, predicted_class, confidence
# Segment image and classify each region
def segment_and_classify_regions(image, k_clusters=2):
rgb_img, _, sharp_img = preprocessing(image)
reshaped = sharp_img.reshape((-1, 3)).astype(np.float32)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
_, labels, centers = cv2.kmeans(reshaped, k_clusters, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
segmented = labels.flatten().reshape(sharp_img.shape[:2])
output = image.copy()
region_predictions = []
for i in range(k_clusters):
mask = (segmented == i).astype(np.uint8) * 255
region = cv2.bitwise_and(sharp_img, sharp_img, mask=mask)
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
continue
x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
region_crop = region[y:y+h, x:x+w]
if region_crop.size == 0:
continue
features_df = feature_extraction(region_crop)
features_scaled = scaler.transform(features_df)
prediction = model.predict(features_scaled)
predicted_class = label_encoder.inverse_transform(prediction)[0]
confidence = np.max(model.predict_proba(features_scaled))
region_predictions.append({
"class": predicted_class,
"confidence": confidence,
"bbox": (x, y, w, h)
})
cv2.rectangle(output, (x, y), (x+w, y+h), (0,255,0), 2)
cv2.putText(output, f"{predicted_class} ({confidence*100:.1f}%)", (x, y-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
return output, region_predictions