Spaces:
Runtime error
Runtime error
File size: 3,710 Bytes
fe66586 b4a6bc0 fe66586 5adf67c fe66586 b4a6bc0 fe66586 2607de5 fe66586 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import numpy as np
import cv2
import pandas as pd
from sklearn.model_selection import train_test_split
from src import data, features, model, augmentation
def process_dataset_with_augmentation(df, is_training=False):
#Loops through the dataframe.
#If is_training=True, it augments minority classes to balance the data.
X = []
y = []
# 1. Calculate Statistics for Balancing (Only needed for training)
if is_training:
class_counts = df['target'].value_counts().to_dict()
max_count = max(class_counts.values())
print(f" [Augmentation] Balancing classes to match majority count: {max_count}")
total = len(df)
for idx, row in df.iterrows():
if idx % 100 == 0: print(f" Processing image {idx}/{total}...")
# Load Original Image
img = cv2.imread(row['path'])
if img is None: continue
# A. Extract Features for Original Image
feats = features.extract_all_features_pipeline(img)
if feats is not None:
X.append(feats)
y.append(row['label_idx'])
# B. Augmentation Logic (Training Only)
if is_training:
# Check how many extra copies we need
class_name = row['target']
# Calculate factor. e.g., if factor is 5, we generate 4 NEW images
# so total = 1 original + 4 augmented = 5
factor = augmentation.get_augmentation_factor(class_name, class_counts, max_count)
num_new_images = factor - 1
if num_new_images > 0:
# Generate variations
aug_imgs = augmentation.generate_augmented_images(img, count=num_new_images)
# Extract features for every augmented variation
for aug_img in aug_imgs:
aug_feats = features.extract_all_features_pipeline(aug_img)
if aug_feats is not None:
X.append(aug_feats)
y.append(row['label_idx'])
return np.array(X), np.array(y)
def main():
#This is the starting point of the training.
#It produces -between others- plots and a model saved in serialized format (classes.pkl)
# 1. Load Data (Metadata only)
df, classes = data.load_metadata(limit=None) # Adjust limit as needed
print("-" * 50)
print("STEP 1: Splitting Data (Train/Test) on File Paths")
print("-" * 50)
# Split DataFrame FIRST to avoid data leakage
df_train, df_test = train_test_split(
df, test_size=0.1, stratify=df['label_idx']
)
print(f"Training Samples (Files): {len(df_train)}")
print(f"Test Samples (Files): {len(df_test)}")
# 2. Process Test Data (No Augmentation, just feature extraction)
print("\n" + "-" * 50)
print("STEP 2: Extracting Test Features (Standard)")
print("-" * 50)
X_test, y_test = process_dataset_with_augmentation(df_test, is_training=False)
# 3. Process Training Data (WITH Augmentation)
print("\n" + "-" * 50)
print("STEP 3: Extracting Training Features (With Keras Augmentation)")
print("-" * 50)
X_train, y_train = process_dataset_with_augmentation(df_train, is_training=True)
print(f"\nFinal Feature Matrix Shapes:")
print(f"X_train: {X_train.shape}, y_train: {y_train.shape}")
print(f"X_test: {X_test.shape}, y_test: {y_test.shape}")
# 4. Train & Evaluate
# We pass the pre-split arrays directly to a modified train function
if len(X_train) > 0 and len(X_test) > 0:
model.train_and_evaluate_split(X_train, y_train, X_test, y_test, classes)
else:
print("Error: Feature extraction failed.")
if __name__ == "__main__":
main() |