Update app.py

app.py

@@ -177,6 +177,9 @@ def upload_page():
             
             st.success(f"Image saved to {filepath}")
             
+            # Save image to Hugging Face repo
+            save_to_hf_repo(filepath, f'uploaded_images/{filename}')
+            
             # Check if the uploaded file is a GeoTIFF
             if file_extension in ['.tiff', '.tif']:
                 st.info('Processing GeoTIFF image...')
@@ -193,29 +196,35 @@ def upload_page():
             # Store the full path of the converted image
             st.session_state.filename = converted_filename
 
-            # Process the image
-            st.write("Processing image...")
-            with st.spinner('Analyzing...'):
-                full_mask = process_large_image(model, converted_filepath)
-            
+            # 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:
+                st.info('Large image detected. Using patch-based processing.')
+                with st.spinner('Analyzing large image...'):
+                    full_mask = process_large_image(model, converted_filepath)
+            else:
+                st.info('Small image detected. Processing whole image at once.')
+                with st.spinner('Analyzing image...'):
+                    img_transformed = transforms(img)
+                    prediction = predict(img_transformed)
+                    full_mask = (prediction > 0.5).astype(np.uint8) * 255
+
             # Save the full mask
             mask_filename = f"mask_{timestamp}.png"
             mask_filepath = os.path.join(MASK_DIR, mask_filename)
             cv2.imwrite(mask_filepath, full_mask)
             st.session_state.mask_filename = mask_filename
 
-            st.success("Image processed successfully")
+            # Save mask to Hugging Face repo
+            save_to_hf_repo(mask_filepath, f'generated_masks/{mask_filename}')
 
             # Log image details
             log_image_details(timestamp, converted_filename, mask_filename)
 
             st.session_state.file_uploaded = True
-            
-            # Clean up temporary patch files
-            st.info('Cleaning up temporary files...')
-            for file in os.listdir(PRED_PATCHES_DIR):
-                os.remove(os.path.join(PRED_PATCHES_DIR, file))
-            st.success('Temporary files cleaned up')
+            st.success("Image processed successfully")
             
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")