ahadhassan commited on
Commit
7ea7f71
·
1 Parent(s): db8255f

integrating new ndvi model

Browse files
Files changed (4) hide show
  1. app.py +18 -1
  2. ndvi_best_model.keras +0 -3
  3. ndvi_predictor.py +92 -51
  4. resize_image.py +62 -0
app.py CHANGED
@@ -13,6 +13,7 @@ from rasterio.transform import from_bounds
13
  import tempfile
14
  import os
15
  import logging
 
16
 
17
  # Configure logging
18
  logging.basicConfig(level=logging.INFO)
@@ -39,6 +40,7 @@ except Exception as e:
39
  async def root():
40
  return {"message": "Welcome to the NDVI and YOLO prediction API!"}
41
 
 
42
  @app.post("/predict_ndvi/")
43
  async def predict_ndvi_api(file: UploadFile = File(...)):
44
  """Predict NDVI from RGB image"""
@@ -46,11 +48,25 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
46
  return JSONResponse(status_code=500, content={"error": "NDVI model not loaded"})
47
 
48
  try:
 
 
 
49
  contents = await file.read()
50
  img = Image.open(BytesIO(contents)).convert("RGB")
51
- norm_img = normalize_rgb(np.array(img))
 
 
 
 
 
 
 
 
 
 
52
  pred_ndvi = predict_ndvi(ndvi_model, norm_img)
53
 
 
54
  # Visualization image as PNG
55
  vis_img_bytes = create_visualization(norm_img, pred_ndvi)
56
  vis_img_bytes.seek(0)
@@ -77,6 +93,7 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
77
  logger.error(f"Error in predict_ndvi_api: {e}")
78
  return JSONResponse(status_code=500, content={"error": str(e)})
79
 
 
80
  @app.post("/predict_yolo/")
81
  async def predict_yolo_api(file: UploadFile = File(...)):
82
  """Predict YOLO results from 4-channel TIFF image"""
 
13
  import tempfile
14
  import os
15
  import logging
16
+ from resize_image import resize_image_optimized, resize_image_simple
17
 
18
  # Configure logging
19
  logging.basicConfig(level=logging.INFO)
 
40
  async def root():
41
  return {"message": "Welcome to the NDVI and YOLO prediction API!"}
42
 
43
+ # Example usage in your predict_ndvi endpoint:
44
  @app.post("/predict_ndvi/")
45
  async def predict_ndvi_api(file: UploadFile = File(...)):
46
  """Predict NDVI from RGB image"""
 
48
  return JSONResponse(status_code=500, content={"error": "NDVI model not loaded"})
49
 
50
  try:
51
+ # Define target size (height, width)
52
+ target_size = (640, 640)
53
+
54
  contents = await file.read()
55
  img = Image.open(BytesIO(contents)).convert("RGB")
56
+
57
+ # Convert to numpy array
58
+ rgb_array = np.array(img)
59
+
60
+ # Resize image to target size
61
+ rgb_resized = resize_image_optimized(rgb_array, target_size)
62
+
63
+ # Normalize the resized image
64
+ norm_img = normalize_rgb(rgb_resized)
65
+
66
+ # Predict NDVI
67
  pred_ndvi = predict_ndvi(ndvi_model, norm_img)
68
 
69
+ # Rest of the endpoint remains the same...
70
  # Visualization image as PNG
71
  vis_img_bytes = create_visualization(norm_img, pred_ndvi)
72
  vis_img_bytes.seek(0)
 
93
  logger.error(f"Error in predict_ndvi_api: {e}")
94
  return JSONResponse(status_code=500, content={"error": str(e)})
95
 
96
+
97
  @app.post("/predict_yolo/")
98
  async def predict_yolo_api(file: UploadFile = File(...)):
99
  """Predict YOLO results from 4-channel TIFF image"""
ndvi_best_model.keras DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:9e85ff83bd897d6c2fd8b4a6f4aada10a1064e8b8dfa53e4e7435078c8abba67
3
- size 122132449
 
 
 
 
ndvi_predictor.py CHANGED
@@ -4,13 +4,14 @@ import os
4
  os.environ["SM_FRAMEWORK"] = "tf.keras"
5
  import segmentation_models as sm
6
  import tensorflow as tf
7
- from tensorflow.keras.models import model_from_json
8
- from efficientnet.tfkeras import EfficientNetB2
9
  import numpy as np
10
  import rasterio
11
  import matplotlib.pyplot as plt
12
  from PIL import Image
13
  import io
 
 
 
14
 
15
  # Custom loss functions and activation functions
16
  def balanced_mse_loss(y_true, y_pred):
@@ -27,31 +28,36 @@ def custom_mae(y_true, y_pred):
27
 
28
  def load_model(models_dir):
29
  """Load NDVI prediction model with custom objects"""
 
30
  # Define custom objects dictionary
31
  custom_objects = {
32
  'balanced_mse_loss': balanced_mse_loss,
33
  'custom_mae': custom_mae
34
  }
35
 
36
- # Load model architecture
37
- with open(os.path.join(models_dir, "model_architecture.json"), "r") as json_file:
38
- model_json = json_file.read()
39
-
40
- model = model_from_json(model_json, custom_objects=custom_objects)
41
-
42
- # Load weights
43
- model.load_weights(os.path.join(models_dir, "best_model_weights.weights.h5"))
44
-
45
- # Compile model with custom functions
46
- optimizer = tf.keras.optimizers.AdamW(learning_rate=0.0005, weight_decay=1e-4)
47
-
48
- model.compile(
49
- optimizer=optimizer,
50
- loss=balanced_mse_loss,
51
- metrics=[custom_mae, 'mse']
52
- )
53
-
54
- return model
 
 
 
 
55
 
56
  def normalize_rgb(rgb):
57
  """Normalize RGB image to [0, 1] range using percentile normalization"""
@@ -71,7 +77,7 @@ def normalize_rgb(rgb):
71
 
72
  def predict_ndvi(model, rgb_np):
73
  """
74
- Predict NDVI from RGB image using tiled approach for large images
75
 
76
  Args:
77
  model: Loaded NDVI prediction model
@@ -81,46 +87,81 @@ def predict_ndvi(model, rgb_np):
81
  ndvi_pred: Predicted NDVI as numpy array (H, W) in range [-1, 1]
82
  """
83
  height, width = rgb_np.shape[:2]
 
 
84
  tile_size = 512
85
- stride = int(tile_size * 0.7)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86
 
87
  # Initialize output arrays
88
  ndvi_pred = np.zeros((height, width), dtype=np.float32)
89
  weight_map = np.zeros((height, width), dtype=np.float32)
90
 
91
- # Handle small images by padding
92
- if height < tile_size or width < tile_size:
93
- pad_height = max(0, tile_size - height)
94
- pad_width = max(0, tile_size - width)
95
- rgb_padded = np.pad(rgb_np, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
96
- height_padded, width_padded = rgb_padded.shape[0], rgb_padded.shape[1]
97
- else:
98
- rgb_padded = rgb_np
99
- height_padded, width_padded = height, width
100
-
101
- # Process image tiles
102
- for i in range(0, height_padded - tile_size + 1, stride):
103
- for j in range(0, width_padded - tile_size + 1, stride):
 
 
 
 
104
  # Extract tile
105
- tile = rgb_padded[i:i+tile_size, j:j+tile_size, :]
106
 
107
- # Create distance-based weights for blending
108
- y, x = np.mgrid[0:tile_size, 0:tile_size]
109
- weights = np.minimum(np.minimum(x, tile_size - x - 1), np.minimum(y, tile_size - y - 1))
110
- weights = np.clip(weights, 0, 50) / 50
 
111
 
112
- # Predict NDVI for tile
113
- tile_pred = model.predict(np.expand_dims(tile, axis=0), verbose=0)[0, :, :, 0]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
114
 
115
- # Determine valid region (handle edge cases)
116
- valid_height = min(tile_size, height - i)
117
- valid_width = min(tile_size, width - j)
118
 
119
- # Accumulate weighted predictions
120
- ndvi_pred[i:i+valid_height, j:j+valid_width] += (
121
- tile_pred[:valid_height, :valid_width] * weights[:valid_height, :valid_width]
122
- )
123
- weight_map[i:i+valid_height, j:j+valid_width] += weights[:valid_height, :valid_width]
 
124
 
125
  # Normalize by weights
126
  mask = weight_map > 0
 
4
  os.environ["SM_FRAMEWORK"] = "tf.keras"
5
  import segmentation_models as sm
6
  import tensorflow as tf
 
 
7
  import numpy as np
8
  import rasterio
9
  import matplotlib.pyplot as plt
10
  from PIL import Image
11
  import io
12
+ from tensorflow.keras.models import model_from_json
13
+ import traceback
14
+ import gc
15
 
16
  # Custom loss functions and activation functions
17
  def balanced_mse_loss(y_true, y_pred):
 
28
 
29
  def load_model(models_dir):
30
  """Load NDVI prediction model with custom objects"""
31
+
32
  # Define custom objects dictionary
33
  custom_objects = {
34
  'balanced_mse_loss': balanced_mse_loss,
35
  'custom_mae': custom_mae
36
  }
37
 
38
+ try:
39
+ # Load model architecture
40
+ with open(os.path.join(models_dir, "model_architecture.json"), "r") as json_file:
41
+ model_json = json_file.read()
42
+
43
+ model = model_from_json(model_json, custom_objects=custom_objects)
44
+
45
+ # Load weights
46
+ model.load_weights(os.path.join(models_dir, "best_model_weights.weights.h5"))
47
+
48
+ # Compile model with custom functions
49
+ optimizer = tf.keras.optimizers.AdamW(learning_rate=0.0005, weight_decay=1e-4)
50
+
51
+ model.compile(
52
+ optimizer=optimizer,
53
+ loss=balanced_mse_loss,
54
+ metrics=[custom_mae, 'mse']
55
+ )
56
+
57
+ return model
58
+ except Exception as e:
59
+ traceback.print_exc()
60
+ return None
61
 
62
  def normalize_rgb(rgb):
63
  """Normalize RGB image to [0, 1] range using percentile normalization"""
 
77
 
78
  def predict_ndvi(model, rgb_np):
79
  """
80
+ Faster NDVI prediction with larger tiles and more efficient processing
81
 
82
  Args:
83
  model: Loaded NDVI prediction model
 
87
  ndvi_pred: Predicted NDVI as numpy array (H, W) in range [-1, 1]
88
  """
89
  height, width = rgb_np.shape[:2]
90
+
91
+ # Larger tiles for faster processing
92
  tile_size = 512
93
+ stride = int(tile_size * 0.75) # 25% overlap
94
+
95
+ # For smaller images, process whole image at once
96
+ if height <= tile_size and width <= tile_size:
97
+ # Pad to tile size if needed
98
+ pad_height = max(0, tile_size - height)
99
+ pad_width = max(0, tile_size - width)
100
+ if pad_height > 0 or pad_width > 0:
101
+ rgb_padded = np.pad(rgb_np, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
102
+ else:
103
+ rgb_padded = rgb_np
104
+
105
+ # Single prediction
106
+ pred = model.predict(np.expand_dims(rgb_padded, axis=0), verbose=0, batch_size=1)[0, :, :, 0]
107
+ return pred[:height, :width]
108
 
109
  # Initialize output arrays
110
  ndvi_pred = np.zeros((height, width), dtype=np.float32)
111
  weight_map = np.zeros((height, width), dtype=np.float32)
112
 
113
+ # Pre-compute weights for efficiency
114
+ y, x = np.mgrid[0:tile_size, 0:tile_size]
115
+ base_weights = np.minimum(np.minimum(x, tile_size - x - 1), np.minimum(y, tile_size - y - 1))
116
+ base_weights = np.clip(base_weights, 0, 64) / 64
117
+
118
+ # Collect all tiles first
119
+ tiles = []
120
+ positions = []
121
+
122
+ for i in range(0, height, stride):
123
+ for j in range(0, width, stride):
124
+ # Calculate actual tile bounds
125
+ end_i = min(i + tile_size, height)
126
+ end_j = min(j + tile_size, width)
127
+ actual_height = end_i - i
128
+ actual_width = end_j - j
129
+
130
  # Extract tile
131
+ tile = rgb_np[i:end_i, j:end_j, :]
132
 
133
+ # Pad if necessary
134
+ if actual_height < tile_size or actual_width < tile_size:
135
+ pad_height = tile_size - actual_height
136
+ pad_width = tile_size - actual_width
137
+ tile = np.pad(tile, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
138
 
139
+ tiles.append(tile)
140
+ positions.append((i, j, actual_height, actual_width))
141
+
142
+ # Process all tiles in larger batches
143
+ batch_size = 8 # Process 8 tiles at once
144
+ for batch_start in range(0, len(tiles), batch_size):
145
+ batch_end = min(batch_start + batch_size, len(tiles))
146
+ batch_tiles = np.array(tiles[batch_start:batch_end])
147
+
148
+ # Predict batch
149
+ batch_preds = model.predict(batch_tiles, verbose=0, batch_size=batch_size)
150
+
151
+ # Apply predictions
152
+ for k in range(batch_end - batch_start):
153
+ pred = batch_preds[k, :, :, 0]
154
+ i, j, actual_height, actual_width = positions[batch_start + k]
155
 
156
+ # Use appropriate weights
157
+ weights = base_weights[:actual_height, :actual_width]
 
158
 
159
+ # Add to output
160
+ ndvi_pred[i:i+actual_height, j:j+actual_width] += pred[:actual_height, :actual_width] * weights
161
+ weight_map[i:i+actual_height, j:j+actual_width] += weights
162
+
163
+ # Clean up batch
164
+ del batch_tiles, batch_preds
165
 
166
  # Normalize by weights
167
  mask = weight_map > 0
resize_image.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image
2
+ import numpy as np
3
+
4
+ # Alternative: Simple resize function using PIL directly
5
+ def resize_image_simple(image_array, target_size):
6
+ """
7
+ Simple resize function using PIL
8
+
9
+ Args:
10
+ image_array: Input image array (H, W, C)
11
+ target_size: Tuple (height, width)
12
+
13
+ Returns:
14
+ Resized image array
15
+ """
16
+ # Ensure image is in correct format
17
+ if image_array.max() <= 1:
18
+ image_array = (image_array * 255).astype(np.uint8)
19
+
20
+ # Convert to PIL Image
21
+ image_pil = Image.fromarray(image_array)
22
+
23
+ # Resize (PIL uses width, height format)
24
+ resized_pil = image_pil.resize((target_size[1], target_size[0]), Image.LANCZOS)
25
+
26
+ # Convert back to numpy array and normalize back to [0, 1]
27
+ resized_array = np.array(resized_pil).astype(np.float32) / 255.0
28
+
29
+ return resized_array
30
+
31
+ def resize_image_optimized(image_array, target_size):
32
+ """
33
+ Resize image to target size with memory optimization
34
+
35
+ Args:
36
+ image_array: Input image array (H, W, C)
37
+ target_size: Tuple (height, width) representing target dimensions
38
+
39
+ Returns:
40
+ Resized image array
41
+ """
42
+ # Convert numpy array to PIL Image
43
+ if image_array.dtype != np.uint8:
44
+ # Convert to uint8 if not already
45
+ if image_array.max() <= 1:
46
+ image_array = (image_array * 255).astype(np.uint8)
47
+ else:
48
+ image_array = image_array.astype(np.uint8)
49
+
50
+ image_pil = Image.fromarray(image_array)
51
+
52
+ # Resize image (PIL uses (width, height) format)
53
+ resized_pil = image_pil.resize((target_size[1], target_size[0]), Image.LANCZOS)
54
+
55
+ # Convert back to numpy array
56
+ result = np.array(resized_pil)
57
+
58
+ # Clean up
59
+ image_pil.close()
60
+ resized_pil.close()
61
+
62
+ return result