import sys
import numpy as np
import cv2
import os
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, accuracy_score

# Function to load images and labels
def load_images_from_folder(folder, label):
    images = []
    labels = []
    for filename in os.listdir(folder):
        img = cv2.imread(os.path.join(folder, filename))
        if img is not None:
            img = cv2.resize(img, (64, 64))
            img = img.flatten()
            images.append(img)
            labels.append(label)
    return images, labels

# Load good and bad welding images
good_images, good_labels = load_images_from_folder('good_welding_image/', 0)  # 0 for good
bad_images, bad_labels = load_images_from_folder('bad_welding_image/', 1)     # 1 for bad

# Combine the images and labels
images = np.array(good_images + bad_images)
labels = np.array(good_labels + bad_labels)

# Split the data into training and testing sets with stratified sampling
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, stratify=labels, random_state=42)

# Create and train the KNN classifier
knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_train, y_train)

# Function to predict a single image
def predict_welding_image(image_path, model):
    img = cv2.imread(image_path)
    img = cv2.resize(img, (64, 64))
    img = img.flatten()
    img = np.array([img])
    prediction = model.predict(img)
    return 'Good' if prediction == 0 else 'Bad'

# Get image path from command line argument
image_path = sys.argv[1]

# Predict the image
result = predict_welding_image(image_path, knn) 
print(result, end='')