Update app.py

app.py

@@ -3,16 +3,14 @@ import numpy as np
 import rasterio
 import cv2
 import matplotlib.pyplot as plt
-import joblib
 from sklearn.ensemble import RandomForestRegressor
 from sklearn.model_selection import train_test_split
-from PIL import Image
-import io
+from sklearn.metrics import mean_squared_error
+import os
 
 # Function to process a single TIFF file
-def process_tiff(file):
-    # Read file from BytesIO object
-    with rasterio.open(io.BytesIO(file.read())) as src:
+def process_tiff(file_path):
+    with rasterio.open(file_path) as src:
         tiff_data = src.read()
     B2_image = tiff_data[1, :, :]  # Assuming B2 is the second band
     target_size = (50, 50)
@@ -20,20 +18,36 @@ def process_tiff(file):
     return B2_resized.reshape(-1, 1)  # Reshape for the model input
 
 # Function to train the RandomForestRegressor model
-def train_random_forest_model():
+def train_random_forest_model(data_dir):
     X_list = []
     y_list = []
 
-    # Placeholder for actual file paths
-    # Modify to load and process multiple TIFF files if needed
-    # For now, using a single uploaded file for training
-    return
+    # Load all TIFF files and preprocess data
+    for root, dirs, files in os.walk(data_dir):
+        for file in files:
+            if file.endswith('.tiff'):
+                file_path = os.path.join(root, file)
+                X_list.append(process_tiff(file_path))
+                # Generate synthetic target data for demonstration (replace with actual targets)
+                y_list.append(np.random.rand(2500))  # Assuming target_size is (50, 50)
 
-# Function to make predictions using the RandomForestRegressor model
-def predict_crop_yield(file, model_name):
+    X = np.vstack(X_list)
+    y = np.hstack(y_list)
+
+    # Split the data into training and testing sets
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    # Train the model
+    model = RandomForestRegressor(n_estimators=100, random_state=42)
+    model.fit(X_train, y_train)
+
+    return model
+
+# Function to make predictions using the trained model
+def predict_crop_yield(file, model_name, data_dir):
     if model_name == 'Random Forest':
-        model = joblib.load('crop_yield_model.joblib')  # Load pre-trained model
-        processed_image = process_tiff(file)
+        model = train_random_forest_model(data_dir)
+        processed_image = process_tiff(file.name)
         prediction = model.predict(processed_image)
         prediction_reshaped = prediction.reshape((50, 50))
         plt.imshow(prediction_reshaped, cmap='viridis')
@@ -44,6 +58,8 @@ def predict_crop_yield(file, model_name):
     else:
         return "Model not found"
 
+data_dir = 'Data'  # Path to the folder containing TIFF files
+
 inputs = [
     gr.File(label='Upload TIFF File'),
     gr.Dropdown(choices=['Random Forest'], label='Model')
@@ -51,7 +67,7 @@ inputs = [
 outputs = gr.Image(type='filepath', label='Predicted Yield Visualization')
 
 demo = gr.Interface(
-    fn=predict_crop_yield,
+    fn=lambda file, model_name: predict_crop_yield(file, model_name, data_dir),
     inputs=inputs,
     outputs=outputs,
     title="Crop Yield Prediction using Random Forest",