Update app.py

app.py

@@ -20,15 +20,18 @@ from Utils.area import pixel_to_sqft, process_and_overlay_image
 from split_merge import split, merge 
 from Utils.convert import read_pansharpened_rgb
 
-# Define directories for Hugging Face Spaces
-UPLOAD_DIR = "/tmp/uploaded_images/"
-MASK_DIR = "/tmp/generated_masks/"
-patches_folder = "/tmp/Patches/"
-pred_patches = "/tmp/Patch_pred/"
-CSV_LOG_PATH = "outputs/image_log.csv"
+# Define base directory for Hugging Face Spaces
+BASE_DIR = "/home/user"
+
+# Define subdirectories
+UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
+MASK_DIR = os.path.join(BASE_DIR, "generated_masks")
+PATCHES_DIR = os.path.join(BASE_DIR, "patches")
+PRED_PATCHES_DIR = os.path.join(BASE_DIR, "pred_patches")
+CSV_LOG_PATH = os.path.join(BASE_DIR, "image_log.csv")
 
 # Create directories
-for directory in [UPLOAD_DIR, MASK_DIR, patches_folder, pred_patches, "outputs"]:
+for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
     os.makedirs(directory, exist_ok=True)
 
 # Load model
@@ -81,89 +84,97 @@ def upload_page():
     image = st.file_uploader('Choose a satellite image', type=['jpg', 'png', 'jpeg', 'tiff', 'tif'])
 
     if image is not None and not st.session_state.file_uploaded:
-        bytes_data = image.getvalue()
-        
-        timestamp = int(time.time())
-        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)
-        
-        st.image(img, caption='Uploaded Image', use_column_width=True)
-        st.success(f'Image saved as {converted_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)
-
-        # 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"mask_{timestamp}.png"
-                merged_mask_path = os.path.join(MASK_DIR, merged_mask_filename)
-                merge(pred_patches, merged_mask_path, 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_filename
+        try:
+            bytes_data = image.getvalue()
+            
+            timestamp = int(time.time())
+            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)
+            
+            st.success(f"Image saved to {filepath}")
+            
+            # 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)
+            
+            st.image(img, caption='Uploaded Image', use_column_width=True)
+            st.success(f'Image processed and saved as {converted_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)
+
+            # 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_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)
+
+                    # Merge predicted patches
+                    merged_mask_filename = f"mask_{timestamp}.png"
+                    merged_mask_path = os.path.join(MASK_DIR, merged_mask_filename)
+                    merge(PRED_PATCHES_DIR, merged_mask_path, 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_DIR)
+                    shutil.rmtree(PRED_PATCHES_DIR)
+                    os.makedirs(PATCHES_DIR)  # Recreate empty folders
+                    os.makedirs(PRED_PATCHES_DIR)
+                    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_filename
 
-        st.session_state.file_uploaded = True
+            st.session_state.file_uploaded = True
+            
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}")
+            st.error("Please check the logs for more details.")
+            print(f"Error in upload_page: {str(e)}")  # This will appear in the Streamlit logs
 
     if st.session_state.file_uploaded and st.button('View result'):
         if st.session_state.filename is None:
@@ -226,8 +237,8 @@ def result_page():
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        shutil.rmtree(patches_folder)
-        shutil.rmtree(pred_patches)
+        shutil.rmtree(PATCHES_DIR)
+        shutil.rmtree(PRED_PATCHES_DIR)
         st.session_state.page = 'upload'
         st.session_state.file_uploaded = False
         st.session_state.filename = None