Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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
|
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 =
|
73 |
-
MASK_DIR =
|
74 |
-
PATCHES_DIR =
|
75 |
-
PRED_PATCHES_DIR =
|
76 |
-
CSV_LOG_PATH =
|
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 |
-
|
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 |
-
|
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)
|