107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
from datetime import datetime
|
|
from ultralytics import YOLO
|
|
import numpy as np
|
|
import cv2
|
|
import csv
|
|
import os
|
|
|
|
def min_max_normalize(matrix):
|
|
min_vals = np.min(matrix, axis=0)
|
|
max_vals = np.max(matrix, axis=0)
|
|
range_vals = max_vals - min_vals
|
|
range_vals[range_vals == 0] = 1
|
|
normalized_matrix = (matrix - min_vals) / range_vals
|
|
return normalized_matrix
|
|
|
|
|
|
|
|
def load_dataset(csv_file,unlabeled_ratio=0.15, test_ratio=0.4):
|
|
|
|
data = np.genfromtxt(csv_file, delimiter=",", dtype=str, skip_header=1)
|
|
class_names = np.unique(data[:, -1])
|
|
print(f"classes: {class_names[0]} / {class_names[1]}")
|
|
print(f"dataset samples: {data.shape[0]} / features: {data.shape[1] - 1}")
|
|
if class_names[0] in np.unique(data[:, -1]) or class_names[1] in np.unique(data[:, -1]):
|
|
data[:, -1] = np.where(data[:, -1] == class_names[0], 1, -1)
|
|
|
|
data = data.astype(np.float32)
|
|
|
|
features = min_max_normalize(data[:, :-1])
|
|
|
|
|
|
np.random.seed(10000)
|
|
indices = np.random.permutation(len(features))
|
|
|
|
split_idx = int(len(features) * (1 - unlabeled_ratio))
|
|
labeled_test_features = features[indices[:split_idx]]
|
|
labeled_test_labels = data[indices[:split_idx]][:, -1]
|
|
U = features[indices[split_idx:]]
|
|
|
|
test_split_idx = int(len(labeled_test_features) * (1 - test_ratio))
|
|
X = labeled_test_features[:test_split_idx]
|
|
y = labeled_test_labels[:test_split_idx]
|
|
X_test = labeled_test_features[test_split_idx:]
|
|
y_test = labeled_test_labels[test_split_idx:]
|
|
y = y.reshape(y.shape[0], 1)
|
|
y_test = y_test.reshape(y_test.shape[0], 1)
|
|
|
|
|
|
return X, y, X_test, y_test, U
|
|
|
|
|
|
def save_result(model, dataset, accuracy, params, results_file):
|
|
file_exists = os.path.isfile(results_file)
|
|
with open(results_file, mode="a", newline="") as f:
|
|
writer = csv.DictWriter(f, fieldnames=["timestamp", "model", "dataset", "parameters", "accuracy"])
|
|
if not file_exists:
|
|
writer.writeheader()
|
|
writer.writerow({
|
|
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|
"model": model,
|
|
"dataset": dataset,
|
|
"parameters": params,
|
|
"accuracy": accuracy
|
|
})
|
|
|
|
def load_results(limit, results_file):
|
|
if not os.path.isfile(results_file):
|
|
return []
|
|
with open(results_file, mode="r") as f:
|
|
reader = list(csv.DictReader(f))
|
|
return reader[::-1][:limit]
|
|
|
|
|
|
|
|
|
|
def predict_yolo(image_path, confidence):
|
|
print("Predicting:", image_path)
|
|
|
|
# Load YOLO model
|
|
model = YOLO("templates/static/public/files/repair/weights/14_class_best.pt")
|
|
|
|
# Run prediction on the input image
|
|
results = model.predict(
|
|
source=f"templates/static/public/files/repair/images/{image_path}",
|
|
save=False,
|
|
conf=confidence,
|
|
device = "cpu",
|
|
batch=4,
|
|
imgsz=320
|
|
)
|
|
|
|
|
|
predicted_img = results[0].plot() # OpenCV image with boxes drawn
|
|
|
|
# Save location (always overwrite the same file)
|
|
output_dir = "templates/static/public/files/repair/predicted"
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_filename = "predicted_image.jpg" # fixed name
|
|
output_path = os.path.join(output_dir, output_filename)
|
|
print(f"image {predicted_img} written in : ", output_path)
|
|
|
|
# Write image
|
|
cv2.imwrite(output_path, predicted_img)
|
|
|
|
# Return only the filename so template can use it
|
|
return output_filename
|