Spaces:
Sleeping
Sleeping
| """ | |
| model_train.py | |
| Generates synthetic Indian crop data, trains a Random Forest model, | |
| and saves model + encoders to disk. | |
| Run this once before starting app.py | |
| """ | |
| import numpy as np | |
| import pandas as pd | |
| from sklearn.ensemble import RandomForestRegressor | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.metrics import r2_score, mean_absolute_error | |
| from sklearn.preprocessing import LabelEncoder | |
| import joblib | |
| import os | |
| np.random.seed(42) | |
| N = 5000 | |
| # ── Categorical options ─────────────────────────────────────────────────────── | |
| CROPS = ["Wheat", "Rice", "Maize", "Soybean", "Cotton"] | |
| SOILS = ["Loamy", "Clayey", "Sandy", "Black", "Red"] | |
| SEASONS = ["Kharif", "Rabi", "Zaid"] | |
| DISEASE = ["None", "Mild", "Moderate", "Severe"] | |
| IRRIGATE = ["Rainfed", "Canal", "Drip", "Sprinkler"] | |
| # ── Base yields (kg/acre) per crop ──────────────────────────────────────────── | |
| BASE_YIELD = { | |
| "Wheat": 950, | |
| "Rice": 1100, | |
| "Maize": 800, | |
| "Soybean": 600, | |
| "Cotton": 500, | |
| } | |
| # ── Synthetic data generation ───────────────────────────────────────────────── | |
| def generate_data(n=N): | |
| crop = np.random.choice(CROPS, n) | |
| soil = np.random.choice(SOILS, n) | |
| season = np.random.choice(SEASONS, n) | |
| disease = np.random.choice(DISEASE, n) | |
| irrigate = np.random.choice(IRRIGATE, n) | |
| rainfall = np.random.uniform(100, 1200, n) | |
| temperature = np.random.uniform(10, 45, n) | |
| fertilizer = np.random.uniform(0, 150, n) | |
| area = np.random.uniform(1, 50, n) | |
| # ── Yield simulation ────────────────────────────────────────────────────── | |
| base = np.array([BASE_YIELD[c] for c in crop], dtype=float) | |
| # rainfall: sweet-spot around 600 mm | |
| rain_factor = 1 - 0.0003 * (rainfall - 600) ** 2 / 1000 | |
| rain_factor = np.clip(rain_factor, 0.5, 1.1) | |
| # temperature: optimal ~28 °C | |
| temp_factor = 1 - 0.002 * (temperature - 28) ** 2 | |
| temp_factor = np.clip(temp_factor, 0.6, 1.1) | |
| # fertilizer: log benefit | |
| fert_factor = 1 + 0.003 * fertilizer - 0.00001 * fertilizer ** 2 | |
| fert_factor = np.clip(fert_factor, 0.8, 1.3) | |
| # disease penalty | |
| disease_penalty = {"None": 1.0, "Mild": 0.85, "Moderate": 0.65, "Severe": 0.40} | |
| dis_factor = np.array([disease_penalty[d] for d in disease]) | |
| # irrigation boost | |
| irr_boost = {"Rainfed": 1.0, "Canal": 1.1, "Drip": 1.2, "Sprinkler": 1.15} | |
| irr_factor = np.array([irr_boost[i] for i in irrigate]) | |
| # soil-crop compatibility bonus | |
| compat = { | |
| ("Wheat", "Loamy"): 1.1, ("Wheat", "Black"): 1.05, | |
| ("Rice", "Clayey"):1.15,("Rice", "Black"): 1.1, | |
| ("Maize", "Loamy"): 1.1, ("Soybean","Black"): 1.1, | |
| ("Cotton", "Black"): 1.2, ("Cotton", "Red"): 1.0, | |
| } | |
| soil_factor = np.array([compat.get((c, s), 1.0) for c, s in zip(crop, soil)]) | |
| yield_kg = (base | |
| * rain_factor | |
| * temp_factor | |
| * fert_factor | |
| * dis_factor | |
| * irr_factor | |
| * soil_factor | |
| + np.random.normal(0, 40, n)) # noise | |
| yield_kg = np.clip(yield_kg, 50, 2500) | |
| return pd.DataFrame({ | |
| "crop": crop, "soil": soil, "season": season, | |
| "disease": disease, "irrigation": irrigate, | |
| "rainfall": rainfall, "temperature": temperature, | |
| "fertilizer": fertilizer, "area": area, | |
| "yield_kg_per_acre": yield_kg, | |
| }) | |
| # ── Encode + train ──────────────────────────────────────────────────────────── | |
| def train(): | |
| print("Generating synthetic data...") | |
| df = generate_data() | |
| cat_cols = ["crop", "soil", "season", "disease", "irrigation"] | |
| encoders = {} | |
| for col in cat_cols: | |
| le = LabelEncoder() | |
| df[col + "_enc"] = le.fit_transform(df[col]) | |
| encoders[col] = le | |
| features = [c + "_enc" for c in cat_cols] + [ | |
| "rainfall", "temperature", "fertilizer", "area" | |
| ] | |
| X = df[features] | |
| y = df["yield_kg_per_acre"] | |
| X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) | |
| print("Training Random Forest (200 trees)...") | |
| model = RandomForestRegressor( | |
| n_estimators=200, | |
| max_depth=12, | |
| min_samples_split=4, | |
| random_state=42, | |
| n_jobs=-1, | |
| ) | |
| model.fit(X_train, y_train) | |
| preds = model.predict(X_test) | |
| r2 = r2_score(y_test, preds) | |
| mae = mean_absolute_error(y_test, preds) | |
| print(f"\n✅ R² = {r2:.3f} MAE = {mae:.1f} kg/acre") | |
| os.makedirs("model", exist_ok=True) | |
| joblib.dump(model, "model/yield_model.pkl") | |
| joblib.dump(encoders, "model/encoders.pkl") | |
| joblib.dump(features, "model/features.pkl") | |
| print("Model saved → model/yield_model.pkl") | |
| return r2, mae | |
| if __name__ == "__main__": | |
| train() | |