Update app.py

app.py

@@ -72,18 +72,22 @@ def refine_mask(mask, blur_kernel=5, threshold_value=127, morph_kernel_size=3, m
 
 # save to dir func
 def save_to_hf_repo(local_path, repo_path):
+    if not os.path.isfile(local_path):
+        st.error(f"File not found at {local_path}")
+        return
+    
     try:
-        hf_api.upload_file(
-            path_or_fileobj=local_path,
-            path_in_repo=repo_path,
-            repo_id=REPO_ID,
-            repo_type=REPO_TYPE,
-            token=HF_TOKEN
-        )
+        with open(local_path, 'rb') as f:
+            hf_api.upload_file(
+                path_or_fileobj=f,
+                path_in_repo=repo_path,
+                repo_id=REPO_ID,
+                repo_type=REPO_TYPE,
+                token=HF_TOKEN
+            )
         st.success(f"File uploaded successfully to {repo_path}")
     except Exception as e:
-        st.error(f"Uploaded")
-        
+        st.error(f"Error during upload: {str(e)}")
         st.exception(e)
 
 
@@ -228,18 +232,8 @@ def upload_page():
             filepath = os.path.join(UPLOAD_DIR, filename)
             converted_filepath = os.path.join(UPLOAD_DIR, converted_filename)
             
-            
-
-            
-            #st.success(f"Image saved to {filepath}")
-            
-            # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
-          
-            try:
-                image_repo_path = f"images/{converted_filename}"
-                save_to_hf_repo(converted_filepath, image_repo_path)
-            except Exception as e:
-                st.error(f"Error saving image to Hugging Face repo: {str(e)}")
+            with open(filepath, "wb") as f:
+                f.write(bytes_data)
             
             # Check if the uploaded file is a GeoTIFF
             if file_extension in ['.tiff', '.tif']:
@@ -250,6 +244,14 @@ def upload_page():
                 img = Image.open(converted_filepath)
             else:
                 img = Image.open(filepath)
+                img.save(converted_filepath)
+            
+            if os.path.exists(converted_filepath):
+                st.success(f"Image saved successfully: {converted_filepath}")
+                file_size = os.path.getsize(converted_filepath)
+                st.write(f"File size: {file_size} bytes")
+            else:
+                st.error(f"Failed to save image: {converted_filepath}")
             
             st.image(img, caption='Uploaded Image', use_column_width=True)
             st.success(f'Image processed and saved as {converted_filename}')
@@ -272,22 +274,43 @@ def upload_page():
                     prediction = predict(img_transformed)
                     full_mask = (prediction > 0.5).astype(np.uint8) * 255
 
-                  
-            
-            # Save the full mask---------------------------------------------------------------------------------------------------
+            # Save the full mask
             full_mask = refine_mask(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
 
-   
+            # Upload to Hugging Face repo
             try:
-                mask_repo_path = f"masks/{mask_filename}"
-                save_to_hf_repo(mask_filepath, mask_repo_path)
+                with open(converted_filepath, 'rb') as f:
+                    image_repo_path = f"images/{converted_filename}"
+                    hf_api.upload_file(
+                        path_or_fileobj=f,
+                        path_in_repo=image_repo_path,
+                        repo_id=REPO_ID,
+                        repo_type=REPO_TYPE,
+                        token=HF_TOKEN
+                    )
+                st.success(f"Image uploaded successfully to {image_repo_path}")
+            except Exception as e:
+                st.error(f"Error saving image to Hugging Face repo: {str(e)}")
+                st.exception(e)
+
+            try:
+                with open(mask_filepath, 'rb') as f:
+                    mask_repo_path = f"masks/{mask_filename}"
+                    hf_api.upload_file(
+                        path_or_fileobj=f,
+                        path_in_repo=mask_repo_path,
+                        repo_id=REPO_ID,
+                        repo_type=REPO_TYPE,
+                        token=HF_TOKEN
+                    )
+                st.success(f"Mask uploaded successfully to {mask_repo_path}")
             except Exception as e:
                 st.error(f"Error saving mask to Hugging Face repo: {str(e)}")
-            
+                st.exception(e)
 
             # Log image details
             log_image_details(timestamp, converted_filename, mask_filename)