Update app.py

app.py

@@ -15,6 +15,8 @@ import rasterio
 import matplotlib.pyplot as plt
 from tensorflow.keras.applications import ResNet50
 from tensorflow.keras.models import Model
+import cv2
+import joblib
 
 # Load crop data
 def load_data():
@@ -55,7 +57,64 @@ def predict_traditional(model_name, year, state, crop, yield_):
     else:
         return "Model not found"
 
-# Load pre-trained deep learning models
+# Train RandomForestRegressor model for deep learning model
+def train_random_forest_model():
+    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)
+        B2_resized = cv2.resize(B2_image, target_size, interpolation=cv2.INTER_NEAREST)
+        return B2_resized.reshape(-1, 1)
+
+    data_dir = 'Data'
+    X_list = []
+    y_list = []
+
+    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))
+                y_list.append(np.random.rand(2500))  # Replace with actual target data
+
+    X = np.vstack(X_list)
+    y = np.hstack(y_list)
+
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    model = RandomForestRegressor(n_estimators=100, random_state=42)
+    model.fit(X_train, y_train)
+
+    return model
+
+rf_model = train_random_forest_model()
+
+def predict_random_forest(file):
+    if file is not None:
+        def process_tiff(file_path):
+            with rasterio.open(file_path) as src:
+                tiff_data = src.read()
+            B2_image = tiff_data[1, :, :]
+            target_size = (50, 50)
+            B2_resized = cv2.resize(B2_image, target_size, interpolation=cv2.INTER_NEAREST)
+            return B2_resized.reshape(-1, 1)
+
+        tiff_processed = process_tiff(file.name)
+        prediction = rf_model.predict(tiff_processed)
+        prediction_reshaped = prediction.reshape((50, 50))
+
+        plt.figure(figsize=(10, 10))
+        plt.imshow(prediction_reshaped, cmap='viridis')
+        plt.colorbar()
+        plt.title('Yield Prediction for Single TIFF File')
+        plt.savefig('/tmp/rf_prediction_overlay.png')
+
+        return '/tmp/rf_prediction_overlay.png'
+    else:
+        return "No file uploaded"
+
+# Load deep learning models
 def load_deep_learning_model(model_name):
     base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(128, 128, 3))
     base_model.trainable = False
@@ -123,9 +182,9 @@ def predict_deep_learning(model_name, file):
             plt.colorbar()
 
             # Save the plot to a file
-            plt.savefig('/tmp/prediction_overlay.png')
+            plt.savefig('/tmp/dl_prediction_overlay.png')
 
-            return '/tmp/prediction_overlay.png'
+            return '/tmp/dl_prediction_overlay.png'
         else:
             return "No file uploaded"
     else:
@@ -141,7 +200,7 @@ inputs_traditional = [
 outputs_traditional = gr.Textbox(label='Predicted Profit')
 
 inputs_deep_learning = [
-    gr.Dropdown(choices=list(deep_learning_models.keys()), label='Model'),
+    gr.Dropdown(choices=list(deep_learning_models.keys()) + ['Random Forest'], label='Model'),
     gr.File(label='Upload TIFF File')
 ]
 outputs_deep_learning = gr.Image(label='Prediction Overlay')
@@ -157,10 +216,10 @@ with gr.Blocks() as demo:
         
     with gr.Tab("Deep Learning Models"):
         gr.Interface(
-            fn=predict_deep_learning,
+            fn=lambda model_name, file: predict_deep_learning(model_name, file) if model_name != 'Random Forest' else predict_random_forest(file),
             inputs=inputs_deep_learning,
             outputs=outputs_deep_learning,
-            title="Crop Yield Prediction using Deep Learning Models"
+            title="Crop Yield Prediction using Deep Learning Models and Random Forest"
         )
 
 demo.launch()