Pavan2k4 commited on
Commit
c05fe39
·
verified ·
1 Parent(s): 6da3bf0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -20
app.py CHANGED
@@ -18,19 +18,13 @@ from model.transform import transforms
18
  from model.unet import UNET
19
  from Utils.area import pixel_to_sqft, process_and_overlay_image
20
  from Utils.convert import read_pansharpened_rgb
21
-
22
-
23
  from huggingface_hub import HfApi, login
24
- import os
25
-
26
- # Set up Hugging Face authentication
27
 
28
  HF_TOKEN = os.environ.get("HF_TOKEN")
29
  if not HF_TOKEN:
30
  raise ValueError("HF_TOKEN environment variable is not set")
31
 
32
- login(token=HF_TOKEN, add_to_git_credential=True)
33
-
34
  hf_api = HfApi()
35
 
36
  REPO_ID = "Pavan2k4/Building_area"
@@ -42,9 +36,36 @@ def load_model():
42
  model.load_state_dict(torch.load('latest.pth', map_location='cpu', weights_only = True)['model_state_dict'])
43
  model.eval()
44
  return model
 
 
45
 
46
 
47
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
 
49
 
50
 
@@ -69,20 +90,17 @@ def save_to_hf_repo(local_path, repo_path):
69
  BASE_DIR = os.getcwd()
70
 
71
  # Define subdirectories
72
- UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
73
- MASK_DIR = os.path.join(BASE_DIR, "generated_masks")
74
- PATCHES_DIR = os.path.join(BASE_DIR, "patches")
75
- PRED_PATCHES_DIR = os.path.join(BASE_DIR, "pred_patches")
76
- CSV_LOG_PATH = os.path.join(BASE_DIR, "image_log.csv")
77
 
78
  # Create directories
79
  for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
80
  os.makedirs(directory, exist_ok=True)
81
 
82
- # Load model
83
-
84
 
85
- model = load_model()
86
 
87
  def predict(image):
88
  with torch.no_grad():
@@ -218,7 +236,7 @@ def upload_page():
218
  #st.success(f"Image saved to {filepath}")
219
 
220
  # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
221
- # Save image to Hugging Face repo
222
  try:
223
  image_repo_path = f"images/{converted_filename}"
224
  save_to_hf_repo(converted_filepath, image_repo_path)
@@ -256,16 +274,16 @@ def upload_page():
256
  prediction = predict(img_transformed)
257
  full_mask = (prediction > 0.5).astype(np.uint8) * 255
258
 
259
-
260
 
261
- # Save the full mask
 
262
  mask_filename = f"mask_{timestamp}.png"
263
  mask_filepath = os.path.join(MASK_DIR, mask_filename)
264
  cv2.imwrite(mask_filepath, full_mask)
265
  st.session_state.mask_filename = mask_filename
266
 
267
- # Save mask to Hugging Face repo---------------------------------------------------------------------------------------------
268
- # Save mask to Hugging Face repo
269
  try:
270
  mask_repo_path = f"masks/{mask_filename}"
271
  save_to_hf_repo(mask_filepath, mask_repo_path)
 
18
  from model.unet import UNET
19
  from Utils.area import pixel_to_sqft, process_and_overlay_image
20
  from Utils.convert import read_pansharpened_rgb
 
 
21
  from huggingface_hub import HfApi, login
 
 
 
22
 
23
  HF_TOKEN = os.environ.get("HF_TOKEN")
24
  if not HF_TOKEN:
25
  raise ValueError("HF_TOKEN environment variable is not set")
26
 
27
+ login(token=HF_TOKEN)
 
28
  hf_api = HfApi()
29
 
30
  REPO_ID = "Pavan2k4/Building_area"
 
36
  model.load_state_dict(torch.load('latest.pth', map_location='cpu', weights_only = True)['model_state_dict'])
37
  model.eval()
38
  return model
39
+ # Load model
40
+ model = load_model()
41
 
42
 
43
 
44
+ def refine_mask(mask, blur_kernel=5, threshold_value=127, morph_kernel_size=3, min_object_size=100):
45
+ """Refine and clean the mask with Gaussian blur, thresholding, morphological operations, and small object removal."""
46
+
47
+ # Ensure mask is grayscale
48
+ if len(mask.shape) > 2:
49
+ mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
50
+
51
+ # Apply Gaussian blur to smooth edges
52
+ mask = cv2.GaussianBlur(mask, (blur_kernel, blur_kernel), 0)
53
+
54
+ # Apply binary threshold
55
+ _, mask = cv2.threshold(mask, threshold_value, 255, cv2.THRESH_BINARY)
56
+
57
+ # Apply morphological operations (opening and closing)
58
+ kernel = np.ones((morph_kernel_size, morph_kernel_size), np.uint8)
59
+ mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
60
+ mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
61
+
62
+ # Remove small objects based on area
63
+ num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
64
+ for i in range(1, num_labels):
65
+ if stats[i, cv2.CC_STAT_AREA] < min_object_size:
66
+ mask[labels == i] = 0
67
+
68
+ return mask
69
 
70
 
71
 
 
90
  BASE_DIR = os.getcwd()
91
 
92
  # Define subdirectories
93
+ UPLOAD_DIR = "uploaded_images"
94
+ MASK_DIR = "generated_masks"
95
+ PATCHES_DIR = "patches"
96
+ PRED_PATCHES_DIR = "pred_patches"
97
+ CSV_LOG_PATH = "image_log.csv"
98
 
99
  # Create directories
100
  for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
101
  os.makedirs(directory, exist_ok=True)
102
 
 
 
103
 
 
104
 
105
  def predict(image):
106
  with torch.no_grad():
 
236
  #st.success(f"Image saved to {filepath}")
237
 
238
  # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
239
+
240
  try:
241
  image_repo_path = f"images/{converted_filename}"
242
  save_to_hf_repo(converted_filepath, image_repo_path)
 
274
  prediction = predict(img_transformed)
275
  full_mask = (prediction > 0.5).astype(np.uint8) * 255
276
 
277
+
278
 
279
+ # Save the full mask---------------------------------------------------------------------------------------------------
280
+ full_mask = refine_mask(full_mask)
281
  mask_filename = f"mask_{timestamp}.png"
282
  mask_filepath = os.path.join(MASK_DIR, mask_filename)
283
  cv2.imwrite(mask_filepath, full_mask)
284
  st.session_state.mask_filename = mask_filename
285
 
286
+
 
287
  try:
288
  mask_repo_path = f"masks/{mask_filename}"
289
  save_to_hf_repo(mask_filepath, mask_repo_path)