Update app.py

app.py

@@ -57,8 +57,38 @@ def predict_traditional(model_name, year, state, crop, yield_):
     else:
         return "Model not found"
 
-# Load the pre-trained RandomForestRegressor model
-rf_model = joblib.load('crop_yield_model.joblib')
+# 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:
@@ -146,13 +176,12 @@ def predict_deep_learning(model_name, file):
             plt.imshow(img_data, cmap='gray', alpha=0.5)
             plt.imshow(overlay, cmap='jet', alpha=0.5)
             plt.title('Crop Yield Prediction Overlay')
+            plt.colorbar()
             plt.savefig('/tmp/dl_prediction_overlay.png')
 
             return '/tmp/dl_prediction_overlay.png'
         else:
             return "No file uploaded"
-    elif model_name == 'Random Forest':
-        return predict_random_forest(file)
     else:
         return "Model not found"
 
@@ -182,7 +211,7 @@ 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 and Random Forest"