ahadhassan commited on
Commit
ade40fc
·
verified ·
1 Parent(s): a13fdf1

dev_ahad (#14)

Browse files

- Updated pipepline predict (3143e36b9f145f99918325dcf5255a17ae4a3b27)

Files changed (2) hide show
  1. app.py +1 -1
  2. yolo_predictor.py +18 -11
app.py CHANGED
@@ -195,7 +195,7 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
195
  rgb_array = np.array(img)
196
  logger.info(f"Converted to RGB array with shape: {rgb_array.shape}")
197
 
198
- # Run the full pipeline
199
  results = predict_pipeline(ndvi_model, yolo_model, rgb_array)
200
  logger.info("Pipeline processing completed successfully")
201
 
 
195
  rgb_array = np.array(img)
196
  logger.info(f"Converted to RGB array with shape: {rgb_array.shape}")
197
 
198
+ # Run the full pipeline (now includes resizing internally)
199
  results = predict_pipeline(ndvi_model, yolo_model, rgb_array)
200
  logger.info("Pipeline processing completed successfully")
201
 
yolo_predictor.py CHANGED
@@ -7,6 +7,7 @@ import tifffile
7
  from rasterio.transform import from_bounds
8
  from ultralytics import YOLO
9
  from ndvi_predictor import normalize_rgb, predict_ndvi
 
10
 
11
  # Configure logging
12
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -119,26 +120,32 @@ def predict_pipeline(ndvi_model, yolo_model, rgb_array, conf=0.001):
119
  logger.info("Starting full prediction pipeline")
120
  logger.info(f"Input RGB array shape: {rgb_array.shape}, dtype: {rgb_array.dtype}")
121
 
122
- # Step 1: Normalize RGB image
123
- logger.info("Step 1: Normalizing RGB image")
124
- normalized_rgb = normalize_rgb(rgb_array)
 
 
 
 
 
 
125
  logger.info(f"Normalized RGB shape: {normalized_rgb.shape}, range: [{normalized_rgb.min():.3f}, {normalized_rgb.max():.3f}]")
126
 
127
- # Step 2: Predict NDVI
128
- logger.info("Step 2: Predicting NDVI from RGB")
129
  ndvi_prediction = predict_ndvi(ndvi_model, normalized_rgb)
130
  logger.info(f"NDVI prediction shape: {ndvi_prediction.shape}, range: [{ndvi_prediction.min():.3f}, {ndvi_prediction.max():.3f}]")
131
 
132
- # Step 3: Create 4-channel TIFF file
133
- logger.info("Step 3: Creating 4-channel TIFF file (RGB+NDVI)")
134
 
135
  # Create temporary file for the 4-channel TIFF
136
  with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
137
  tiff_path = tmp_file.name
138
 
139
  try:
140
- # Create the 4-channel TIFF
141
- create_4channel_tiff(rgb_array, ndvi_prediction, tiff_path)
142
 
143
  # Verify the created file
144
  if not os.path.exists(tiff_path):
@@ -147,8 +154,8 @@ def predict_pipeline(ndvi_model, yolo_model, rgb_array, conf=0.001):
147
  file_size = os.path.getsize(tiff_path)
148
  logger.info(f"Created 4-channel TIFF file size: {file_size} bytes")
149
 
150
- # Step 4: Run YOLO prediction on the 4-channel TIFF
151
- logger.info("Step 4: Running YOLO prediction on 4-channel TIFF")
152
  results = predict_yolo(yolo_model, tiff_path, conf=conf)
153
 
154
  logger.info("Full pipeline completed successfully")
 
7
  from rasterio.transform import from_bounds
8
  from ultralytics import YOLO
9
  from ndvi_predictor import normalize_rgb, predict_ndvi
10
+ from resize_image import resize_image_optimized
11
 
12
  # Configure logging
13
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
120
  logger.info("Starting full prediction pipeline")
121
  logger.info(f"Input RGB array shape: {rgb_array.shape}, dtype: {rgb_array.dtype}")
122
 
123
+ # Step 1: Resize RGB image to target size
124
+ logger.info("Step 1: Resizing RGB image to target size")
125
+ target_size = (640, 640) # (height, width)
126
+ rgb_resized = resize_image_optimized(rgb_array, target_size)
127
+ logger.info(f"Resized RGB shape: {rgb_resized.shape}")
128
+
129
+ # Step 2: Normalize RGB image
130
+ logger.info("Step 2: Normalizing RGB image")
131
+ normalized_rgb = normalize_rgb(rgb_resized)
132
  logger.info(f"Normalized RGB shape: {normalized_rgb.shape}, range: [{normalized_rgb.min():.3f}, {normalized_rgb.max():.3f}]")
133
 
134
+ # Step 3: Predict NDVI
135
+ logger.info("Step 3: Predicting NDVI from RGB")
136
  ndvi_prediction = predict_ndvi(ndvi_model, normalized_rgb)
137
  logger.info(f"NDVI prediction shape: {ndvi_prediction.shape}, range: [{ndvi_prediction.min():.3f}, {ndvi_prediction.max():.3f}]")
138
 
139
+ # Step 4: Create 4-channel TIFF file
140
+ logger.info("Step 4: Creating 4-channel TIFF file (RGB+NDVI)")
141
 
142
  # Create temporary file for the 4-channel TIFF
143
  with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
144
  tiff_path = tmp_file.name
145
 
146
  try:
147
+ # Create the 4-channel TIFF using resized RGB and predicted NDVI
148
+ create_4channel_tiff(rgb_resized, ndvi_prediction, tiff_path)
149
 
150
  # Verify the created file
151
  if not os.path.exists(tiff_path):
 
154
  file_size = os.path.getsize(tiff_path)
155
  logger.info(f"Created 4-channel TIFF file size: {file_size} bytes")
156
 
157
+ # Step 5: Run YOLO prediction on the 4-channel TIFF
158
+ logger.info("Step 5: Running YOLO prediction on 4-channel TIFF")
159
  results = predict_yolo(yolo_model, tiff_path, conf=conf)
160
 
161
  logger.info("Full pipeline completed successfully")