Update yolo_predictor.py

yolo_predictor.py

@@ -5,7 +5,6 @@ import rasterio
 from ultralytics import YOLO
 from ndvi_predictor import normalize_rgb, predict_ndvi
 import tempfile
-from rasterio.transform import from_bounds
 from PIL import Image
 import tifffile
 
@@ -46,38 +45,60 @@ def predict_yolo(yolo_model, image_path, conf=0.001):
     """
     # Verify the image has 4 channels before prediction
     try:
-        # Check image format and channels
-        with Image.open(image_path) as img:
-            if hasattr(img, 'n_frames'):
-                # Multi-frame TIFF
-                channels = img.n_frames
-            else:
-                # Regular image
-                channels = len(img.getbands()) if hasattr(img, 'getbands') else 3
+        # Use tifffile for 32-bit TIFF support
+        img_array = tifffile.imread(image_path)
         
-        # If not 4 channels, try with tifffile
-        if channels != 4:
-            img_array = tifffile.imread(image_path)
-            if len(img_array.shape) == 3:
-                if img_array.shape[0] == 4:
-                    channels = 4
-                elif img_array.shape[2] == 4:
-                    channels = 4
-                else:
-                    channels = img_array.shape[0] if img_array.shape[0] <= 4 else img_array.shape[2]
-            else:
-                channels = 1
+        # Handle different array shapes
+        if len(img_array.shape) == 3:
+            if img_array.shape[0] == 4:
+                # Shape is (4, H, W) - transpose to (H, W, 4)
+                img_array = np.transpose(img_array, (1, 2, 0))
+            elif img_array.shape[2] != 4:
+                raise ValueError(f"YOLO model expects 4-channel images, but got {img_array.shape[2]} channels")
+        else:
+            raise ValueError(f"Unexpected image shape: {img_array.shape}")
         
-        if channels != 4:
-            raise ValueError(f"YOLO model expects 4-channel images, but got {channels} channels")
+        # Convert 32-bit float to uint8 if necessary
+        if img_array.dtype != np.uint8:
+            # Normalize to [0, 255] for RGB channels (first 3)
+            rgb_array = img_array[:, :, :3]
+            if rgb_array.max() > 1.0:
+                rgb_array = np.clip(rgb_array / rgb_array.max() * 255, 0, 255).astype(np.uint8)
+            else:
+                rgb_array = np.clip(rgb_array * 255, 0, 255).astype(np.uint8)
             
-    except Exception as e:
-        raise ValueError(f"Error reading image channels: {str(e)}")
-    
-    # Run YOLO prediction
-    results = yolo_model([image_path], conf=conf)
+            # Normalize NDVI (4th channel) from [-1, 1] to [0, 255]
+            ndvi_array = img_array[:, :, 3]
+            ndvi_normalized = ((ndvi_array + 1) * 127.5).astype(np.uint8)
+            
+            # Recombine into 4-channel uint8 array
+            img_array = np.zeros((img_array.shape[0], img_array.shape[1], 4), dtype=np.uint8)
+            img_array[:, :, :3] = rgb_array
+            img_array[:, :, 3] = ndvi_normalized
+            
+            # Save normalized image to temporary file
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
+                temp_path = tmp_file.name
+                tifffile.imwrite(
+                    temp_path,
+                    img_array,
+                    photometric='rgb',
+                    compress='lzw',
+                    metadata={'axes': 'YXC', 'resolution': (1, 1)}  # DPI=1
+                )
+                image_path = temp_path
+        
+        # Run YOLO prediction
+        results = yolo_model([image_path], conf=conf)
+        
+        # Clean up temporary file if created
+        if 'temp_path' in locals() and os.path.exists(temp_path):
+            os.unlink(temp_path)
+        
+        return results[0]  # Return first result
     
-    return results[0]  # Return first result
+    except Exception as e:
+        raise ValueError(f"Error processing image: {str(e)}")
 
 def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     """
@@ -90,7 +111,7 @@ def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     """
     height, width = rgb_array.shape[:2]
     
-    # Ensure RGB is in uint8 format for better compatibility
+    # Ensure RGB is in uint8 format
     if rgb_array.dtype != np.uint8:
         if rgb_array.max() <= 1.0:
             rgb_normalized = (rgb_array * 255).astype(np.uint8)
@@ -109,13 +130,14 @@ def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     four_channel[:, :, 2] = rgb_normalized[:, :, 2]  # Blue
     four_channel[:, :, 3] = ndvi_normalized          # NDVI
     
-    # Save using tifffile with proper format for YOLO compatibility
+    # Save using tifffile with explicit 32-bit compatibility and DPI=1
     tifffile.imwrite(
-        output_path, 
-        four_channel, 
+        output_path,
+        four_channel,
         photometric='rgb',
         compress='lzw',
-        metadata={'axes': 'YXC'}
+        metadata={'axes': 'YXC', 'resolution': (1, 1)},  # DPI=1
+        bitspersample=8  # Explicitly set to 8-bit per channel
     )
 
 def load_4channel_tiff(image_path):
@@ -225,7 +247,7 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
         except Exception as e2:
             try:
                 # Method 3: Fall back to PIL for standard formats
-                img = Image.open(image_path)
+                img = Image.PIL(image_path)
                 if img.mode != 'RGB':
                     img = img.convert('RGB')
                 rgb_array = np.array(img)
@@ -238,12 +260,10 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
     
     # Ensure RGB is in correct format and range
     if rgb_array.dtype == np.uint8:
-        # Keep as uint8 but also create float version for NDVI prediction
         rgb_float = rgb_array.astype(np.float32) / 255.0
     else:
-        # Already float, ensure range is [0, 1]
         if rgb_array.max() > 1.0:
-            rgb_float = rgb_array / 255.0
+            rgb_float = rgb_array / rgb_array.max()
         else:
             rgb_float = rgb_array
         rgb_array = (rgb_float * 255).astype(np.uint8)
@@ -261,13 +281,11 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
         
         # Verify the created file can be read
         try:
-            test_img = Image.open(temp_4ch_path)
-            if hasattr(test_img, 'n_frames'):
-                channels = test_img.n_frames
+            test_array = tifffile.imread(temp_4ch_path)
+            if len(test_array.shape) == 3 and (test_array.shape[0] == 4 or test_array.shape[2] == 4):
+                channels = 4
             else:
-                channels = len(test_img.getbands())
-            test_img.close()
-            
+                channels = test_array.shape[2] if len(test_array.shape) == 3 else 1
             if channels != 4:
                 raise ValueError(f"Created TIFF has {channels} channels instead of 4")