Update app.py

app.py

@@ -42,6 +42,7 @@ def predict(image):
 
 def log_image_details(image_id, image_filename, mask_filename):
     file_exists = os.path.exists(CSV_LOG_PATH)
+    
     current_time = datetime.now()
     date = current_time.strftime('%Y-%m-%d')
     time = current_time.strftime('%H:%M:%S')
@@ -51,89 +52,131 @@ def log_image_details(image_id, image_filename, mask_filename):
         if not file_exists:
             writer.writerow(['S.No', 'Date', 'Time', 'Image ID', 'Image Filename', 'Mask Filename'])
         
-        sno = sum(1 for row in open(CSV_LOG_PATH)) if file_exists else 1
+        # Get the next S.No
+        if file_exists:
+            with open(CSV_LOG_PATH, mode='r') as f:
+                reader = csv.reader(f)
+                sno = sum(1 for row in reader)
+        else:
+            sno = 1
+        
         writer.writerow([sno, date, time, image_id, image_filename, mask_filename])
 
-def reset_state():
-    st.session_state.file_uploaded = False
-    st.session_state.filename = None
-    st.session_state.mask_filename = None
-    st.session_state.tr_img = None
-    if 'page' in st.session_state:
-        del st.session_state.page
-
-def process_image(image, timestamp):
-    filename = f"image_{timestamp}{os.path.splitext(image.name)[1]}"
-    filepath = os.path.join(UPLOAD_DIR, filename)
+def overlay_mask(image, mask, alpha=0.5, rgb=[255, 0, 0]):
+    # Ensure image is 3-channel
+    if len(image.shape) == 2:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
     
-    with open(filepath, "wb") as f:
-        f.write(image.getvalue())
+    # Ensure mask is binary and same shape as image
+    mask = mask.astype(bool)
+    if mask.shape[:2] != image.shape[:2]:
+        raise ValueError("Mask and image must have the same dimensions")
     
-    if filename.lower().endswith(('.tiff', '.tif')):
-        st.info('Processing GeoTIFF image...')
-        convert_gtiff_to_8bit(filepath)
-        st.success('GeoTIFF converted to 8-bit image')
+    # Create color overlay
+    color_mask = np.zeros_like(image)
+    color_mask[mask] = rgb
     
-    return filename, filepath
-
-def predict_image(img_array, filename, timestamp):
-    if img_array.shape[0] > 650 or img_array.shape[1] > 650:
-        split(os.path.join(UPLOAD_DIR, filename), patch_size=256)
-        
-        with st.spinner('Analyzing...'):
-            for patch_filename in os.listdir(PATCHES_DIR):
-                if patch_filename.endswith(".png"):
-                    patch_path = os.path.join(PATCHES_DIR, patch_filename)
-                    patch_img = Image.open(patch_path)
-                    patch_tr_img = transforms(patch_img)
-                    prediction = predict(patch_tr_img)
-                    mask = (prediction > 0.5).astype(np.uint8) * 255
-                    mask_filename = f"mask_{patch_filename}"
-                    mask_filepath = os.path.join(PRED_PATCHES_DIR, mask_filename)
-                    Image.fromarray(mask).save(mask_filepath)
-
-            merged_mask_filename = f"mask_{timestamp}.png"
-            merged_mask_filepath = os.path.join(MASK_DIR, merged_mask_filename)
-            merge(PRED_PATCHES_DIR, merged_mask_filepath, img_array.shape)
-            
-            st.info('Cleaning up temporary files...')
-            for dir in [PATCHES_DIR, PRED_PATCHES_DIR]:
-                shutil.rmtree(dir)
-                os.makedirs(dir)
-            st.success('Temporary files cleaned up')
-    else:
-        tr_img = transforms(Image.open(os.path.join(UPLOAD_DIR, filename)))
-        prediction = predict(tr_img)
-        mask = (prediction > 0.5).astype(np.uint8) * 255
-        merged_mask_filename = f"mask_{timestamp}.png"
-        merged_mask_filepath = os.path.join(MASK_DIR, merged_mask_filename)
-        Image.fromarray(mask).save(merged_mask_filepath)
+    # Blend the image and color mask
+    output = cv2.addWeighted(image, 1, color_mask, alpha, 0)
     
-    return merged_mask_filepath
+    return output
+
+import shutil  # Add this import at the top of your file
 
 def upload_page():
     if 'file_uploaded' not in st.session_state:
         st.session_state.file_uploaded = False
     
+    if 'filename' not in st.session_state:
+        st.session_state.filename = None
+    
+    if 'mask_filename' not in st.session_state:
+        st.session_state.mask_filename = None
+
     image = st.file_uploader('Choose a satellite image', type=['jpg', 'png', 'jpeg', 'tiff', 'tif'])
 
-    if image is not None:
-        reset_state()
+    if image is not None and not st.session_state.file_uploaded:
+        bytes_data = image.getvalue()
+        
         timestamp = int(time.time())
-        filename, filepath = process_image(image, timestamp)
+        original_filename = image.name
+        file_extension = os.path.splitext(original_filename)[1].lower()
+        
+        if file_extension in ['.tiff', '.tif']:
+            filename = f"image_{timestamp}.tif"
+            converted_filename = f"image_{timestamp}_converted.png"
+        else:
+            filename = f"image_{timestamp}.png"
+            converted_filename = filename
+        
+        filepath = os.path.join(UPLOAD_DIR, filename)
+        converted_filepath = os.path.join(UPLOAD_DIR, converted_filename)
+        
+        with open(filepath, "wb") as f:
+            f.write(bytes_data)
+        
+        # Check if the uploaded file is a GeoTIFF
+        if file_extension in ['.tiff', '.tif']:
+            st.info('Processing GeoTIFF image...')
+            rgb_image = read_pansharpened_rgb(filepath)
+            cv2.imwrite(converted_filepath, cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR))
+            st.success(f'GeoTIFF converted to 8-bit image and saved as {converted_filename}')
+            img = Image.open(converted_filepath)
+        else:
+            img = Image.open(filepath)
         
-        img = Image.open(filepath)
         st.image(img, caption='Uploaded Image', use_column_width=True)
-        st.success(f'Image saved as {filename}')
+        st.success(f'Image saved as {converted_filename}')
 
-        st.session_state.filename = filename
+        # Store the full path of the converted image
+        st.session_state.filename = converted_filename
+
+        # Convert image to numpy array
         img_array = np.array(img)
-        
-        mask_filepath = predict_image(img_array, filename, timestamp)
-        st.session_state.mask_filename = mask_filepath
 
-        log_image_details(timestamp, filename, os.path.basename(mask_filepath))
-        
+        # Check if image shape is more than 650x650
+        if img_array.shape[0] > 650 or img_array.shape[1] > 650:
+            # Split image into patches
+            split(converted_filepath, patch_size=512)
+
+            # Display buffer while analyzing
+            with st.spinner('Analyzing...'):
+                # Predict on each patch
+                for patch_filename in os.listdir(patches_folder):
+                    if patch_filename.endswith(".png"):
+                        patch_path = os.path.join(patches_folder, patch_filename)
+                        patch_img = Image.open(patch_path)
+                        patch_tr_img = transforms(patch_img)
+                        prediction = predict(patch_tr_img)
+                        mask = (prediction > 0.5).astype(np.uint8) * 255
+                        mask_filename = f"mask_{patch_filename}"
+                        mask_filepath = os.path.join(pred_patches, mask_filename)
+                        Image.fromarray(mask).save(mask_filepath)
+
+                # Merge predicted patches
+                merged_mask_filename = f"generated_masks/mask_{timestamp}.png"
+                merge(pred_patches, merged_mask_filename, img_array.shape)
+
+                # Save merged mask
+                st.session_state.mask_filename = merged_mask_filename
+
+                # Clean up temporary patch files
+                st.info('Cleaning up temporary files...')
+                shutil.rmtree(patches_folder)
+                shutil.rmtree(pred_patches)
+                os.makedirs(patches_folder)  # Recreate empty folders
+                os.makedirs(pred_patches)
+                st.success('Temporary files cleaned up')
+        else:
+            # Predict on whole image
+            st.session_state.tr_img = transforms(img)
+            prediction = predict(st.session_state.tr_img)
+            mask = (prediction > 0.5).astype(np.uint8) * 255
+            mask_filename = f"mask_{timestamp}.png"
+            mask_filepath = os.path.join(MASK_DIR, mask_filename)
+            Image.fromarray(mask).save(mask_filepath)
+            st.session_state.mask_filename = mask_filepath
+
         st.session_state.file_uploaded = True
 
     if st.session_state.file_uploaded and st.button('View result'):
@@ -143,28 +186,31 @@ def upload_page():
             st.success('Image analyzed')
             st.session_state.page = 'result'
             st.rerun()
-
 def result_page():
     st.title('Analysis Result')
     
     if 'filename' not in st.session_state or 'mask_filename' not in st.session_state:
         st.error("No image or mask file found. Please upload and process an image first.")
         if st.button('Back to Upload'):
-            reset_state()
+            st.session_state.page = 'upload'
+            st.session_state.file_uploaded = False
+            st.session_state.filename = None
+            st.session_state.mask_filename = None
             st.rerun()
         return
 
     col1, col2 = st.columns(2)
 
+    # Display original image
     original_img_path = os.path.join(UPLOAD_DIR, st.session_state.filename)
-    mask_path = st.session_state.mask_filename
-
     if os.path.exists(original_img_path):
         original_img = Image.open(original_img_path)
         col1.image(original_img, caption='Original Image', use_column_width=True)
     else:
         col1.error(f"Original image file not found: {original_img_path}")
 
+    # Display predicted mask
+    mask_path = st.session_state.mask_filename
     if os.path.exists(mask_path):
         mask = Image.open(mask_path)
         col2.image(mask, caption='Predicted Mask', use_column_width=True)
@@ -173,23 +219,35 @@ def result_page():
 
     st.subheader("Overlay with Area of Buildings (sqft)")
 
+    # Display overlayed image
     if os.path.exists(original_img_path) and os.path.exists(mask_path):
         original_np = cv2.imread(original_img_path)
         mask_np = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         
+        # Ensure mask is binary
         _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
 
+        # Resize mask to match original image size if necessary
         if original_np.shape[:2] != mask_np.shape[:2]:
             mask_np = cv2.resize(mask_np, (original_np.shape[1], original_np.shape[0]))
 
+        # Process and overlay image
         overlay_img = process_and_overlay_image(original_np, mask_np, 'output.png')
         
+        # Convert BGR to RGB for displaying with st.image
+       # overlay_rgb = cv2.cvtColor(overlay_img, cv2.COLOR_BGR2RGB)
+        
         st.image(overlay_img, caption='Overlay Image', use_column_width=True)
     else:
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        reset_state()
+        shutil.rmtree(patches_folder)
+        shutil.rmtree(pred_patches)
+        st.session_state.page = 'upload'
+        st.session_state.file_uploaded = False
+        st.session_state.filename = None
+        st.session_state.mask_filename = None
         st.rerun()
 
 def main():
@@ -204,4 +262,4 @@ def main():
         result_page()
 
 if __name__ == '__main__':
-    main()
+    main()