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,34 @@ def predict_traditional(model_name, year, state, crop, yield_):
     else:
         return "Model not found"
 
-# Load pre-trained deep learning models
+# Load the pre-trained RandomForestRegressor model
+rf_model = joblib.load('crop_yield_model.joblib')
+
+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
@@ -105,27 +134,22 @@ def predict_deep_learning(model_name, file):
             predictions = model.predict(preprocessed_patches)
             predictions = predictions.reshape((n_patches_y, n_patches_x))
 
-            # Set a threshold to highlight areas with higher predicted yields
-            threshold = np.percentile(predictions, 90)  # Adjust the percentile as needed
+            threshold = np.percentile(predictions, 90)
 
-            # Create an overlay image to visualize predictions
             overlay = np.zeros_like(img_data, dtype=np.float32)
             for i in range(n_patches_y):
                 for j in range(n_patches_x):
                     if predictions[i, j] > threshold:
                         overlay[i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size] = predictions[i, j]
 
-            # Plot the overlay on the original image
             plt.figure(figsize=(10, 10))
             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()
-
-            # Save the plot to a file
-            plt.savefig('/tmp/prediction_overlay.png')
+            # plt.colorbar()
+            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 +165,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 +181,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()