Spaces:
Sleeping
Sleeping
File size: 2,446 Bytes
a36604b | 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 | import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import silhouette_score
import pickle
import os
# Create directories if they don't exist
os.makedirs('model', exist_ok=True)
os.makedirs('outputs', exist_ok=True)
def train_model():
print("Loading dataset...")
df = pd.read_csv('data/Mall_Customers.csv')
# 1. Data Preprocessing
print("Preprocessing data...")
# Drop CustomerID
df_processed = df.drop('CustomerID', axis=1)
# Encode Gender
le = LabelEncoder()
df_processed['Gender'] = le.fit_transform(df_processed['Gender'])
# Scale features
scaler = StandardScaler()
features = ['Gender', 'Age', 'Annual Income (k$)', 'Spending Score (1-100)']
df_scaled = scaler.fit_transform(df_processed[features])
# 2. Elbow Method
print("Implementing Elbow Method...")
wcss = []
for i in range(1, 11):
kmeans = KMeans(n_clusters=i, init='k-means++', random_state=42, n_init=10)
kmeans.fit(df_scaled)
wcss.append(kmeans.inertia_)
plt.figure(figsize=(10, 6))
plt.plot(range(1, 11), wcss, marker='o', linestyle='--')
plt.title('Elbow Method')
plt.xlabel('Number of clusters')
plt.ylabel('WCSS')
plt.savefig('outputs/elbow_plot.png')
print("Elbow plot saved to outputs/elbow_plot.png")
# 3. Train with optimal k (default 5 as requested)
k = 5
print(f"Training K-Means with k={k}...")
kmeans = KMeans(n_clusters=k, init='k-means++', random_state=42, n_init=10)
clusters = kmeans.fit_predict(df_scaled)
# 4. Compute Silhouette Score
score = silhouette_score(df_scaled, clusters)
print(f"Silhouette Score for k={k}: {score:.4f}")
# 5. Save model, scaler and encoder
print("Saving model and tools...")
with open('model/kmeans.pkl', 'wb') as f:
pickle.dump(kmeans, f)
with open('model/scaler.pkl', 'wb') as f:
pickle.dump(scaler, f)
with open('model/label_encoder.pkl', 'wb') as f:
pickle.dump(le, f)
# Save the scaled data for visualization in app
df_result = df_processed.copy()
df_result['Cluster'] = clusters
df_result.to_csv('data/clustered_customers.csv', index=False)
print("Training complete!")
if __name__ == "__main__":
train_model()
|