Pavan2k4/Building_area
0
1import streamlit as st2import sys3import os4import shutil5import time6from datetime import datetime7import csv8import cv29import numpy as np10from PIL import Image11import torch12sys.path.append('Utils')13sys.path.append('model')14from model.CBAM.reunet_cbam import reunet_cbam15from model.transform import transforms16from model.unet import UNET17from Utils.area import pixel_to_sqft, process_and_overlay_image18from Utils.convert import read_pansharpened_rgb19 20 21@st.cache_resource22def load_model():23 model = reunet_cbam()24 model.load_state_dict(torch.load('latest.pth', map_location='cpu', weights_only = True)['model_state_dict'])25 model.eval()26 return model27# Load model28model = load_model()29 30 31 32def refine_mask(mask, blur_kernel=5, threshold_value=127, morph_kernel_size=3, min_object_size=100):33 """Refine and clean the mask with Gaussian blur, thresholding, morphological operations, and small object removal."""34 35 # Ensure mask is grayscale36 if len(mask.shape) > 2:37 mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)38 39 # Apply Gaussian blur to smooth edges40 mask = cv2.GaussianBlur(mask, (blur_kernel, blur_kernel), 0)41 42 # Apply binary threshold43 _, mask = cv2.threshold(mask, threshold_value, 255, cv2.THRESH_BINARY)44 45 # Apply morphological operations (opening and closing)46 kernel = np.ones((morph_kernel_size, morph_kernel_size), np.uint8)47 mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)48 mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)49 50 # Remove small objects based on area51 num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)52 for i in range(1, num_labels):53 if stats[i, cv2.CC_STAT_AREA] < min_object_size:54 mask[labels == i] = 055 56 return mask57 58 59 60 61# save to dir func62 63 64 65 66base = os.getcwd()67# Define subdirectories68UPLOAD_DIR = os.path.join(base,"Images")69MASK_DIR = os.path.join(base,"Masks")70 71CSV_LOG_PATH = "image_log.csv"72 73 74# Create directories with read and write permissions75for directory in [UPLOAD_DIR, MASK_DIR]:76 os.makedirs(directory, exist_ok=True)77 78 79 80 81def predict(image):82 with torch.no_grad():83 output = model(image.unsqueeze(0))84 return output.squeeze().cpu().numpy()85 86def split_image(image, patch_size=512):87 h, w, _ = image.shape88 patches = []89 for y in range(0, h, patch_size):90 for x in range(0, w, patch_size):91 patch = image[y:min(y+patch_size, h), x:min(x+patch_size, w)]92 patches.append((f"patch_{y}_{x}.png", patch))93 return patches94 95def merge(patch_folder, dest_image='out.png', image_shape=None):96 merged = np.zeros(image_shape[:-1] + (3,), dtype=np.uint8)97 for filename in os.listdir(patch_folder):98 if filename.endswith(".png"):99 patch_path = os.path.join(patch_folder, filename)100 patch = cv2.imread(patch_path)101 patch_height, patch_width, _ = patch.shape102 103 # Extract patch coordinates from filename104 parts = filename.split("_")105 x, y = None, None106 for part in parts:107 if part.endswith(".png"):108 x = int(part.split(".")[0])109 elif part.isdigit():110 y = int(part)111 if x is None or y is None:112 raise ValueError(f"Invalid filename: {filename}")113 114 # Check if patch fits within image boundaries115 if x + patch_width > image_shape[1] or y + patch_height > image_shape[0]:116 # Adjust patch position to fit within image boundaries117 if x + patch_width > image_shape[1]:118 x = image_shape[1] - patch_width119 if y + patch_height > image_shape[0]:120 y = image_shape[0] - patch_height121 122 # Merge patch into the main image123 merged[y:y+patch_height, x:x+patch_width, :] = patch124 125 cv2.imwrite(dest_image, merged)126 return merged127 128def process_large_image(model, image_path, patch_size=512):129 # Read the image130 img = cv2.imread(image_path)131 if img is None:132 raise ValueError(f"Failed to read image from {image_path}")133 134 h, w, _ = img.shape135 st.write(f"Processing image of size {w}x{h}")136 137 # Split the image into patches138 patches = split_image(img, patch_size)139 140 # Process each patch141 for filename, patch in patches:142 patch_pil = Image.fromarray(cv2.cvtColor(patch, cv2.COLOR_BGR2RGB))143 patch_transformed = transforms(patch_pil)144 prediction = predict(patch_transformed)145 mask = (prediction > 0.5).astype(np.uint8) * 255146 147 # Save the mask patch148 mask_filepath = os.path.join(PRED_PATCHES_DIR, filename)149 cv2.imwrite(mask_filepath, mask)150 151 # Merge the predicted patches152 merged_mask = merge(PRED_PATCHES_DIR, dest_image='merged_mask.png', image_shape=img.shape)153 154 return merged_mask155 156def log_image_details(image_id, image_filename, mask_filename):157 file_exists = os.path.exists(CSV_LOG_PATH)158 159 current_time = datetime.now()160 date = current_time.strftime('%Y-%m-%d')161 time = current_time.strftime('%H:%M:%S')162 163 with open(CSV_LOG_PATH, mode='a', newline='') as file:164 writer = csv.writer(file)165 if not file_exists:166 writer.writerow(['S.No', 'Date', 'Time', 'Image ID', 'Image Filename', 'Mask Filename'])167 168 # Get the next S.No169 if file_exists:170 with open(CSV_LOG_PATH, mode='r') as f:171 reader = csv.reader(f)172 sno = sum(1 for row in reader)173 else:174 sno = 1175 176 writer.writerow([sno, date, time, image_id, image_filename, mask_filename])177 178 179def upload_page():180 if 'file_uploaded' not in st.session_state:181 st.session_state.file_uploaded = False182 if 'filename' not in st.session_state:183 st.session_state.filename = None184 if 'mask_filename' not in st.session_state:185 st.session_state.mask_filename = None186 187 image = st.file_uploader('Choose a satellite image', type=['jpg', 'png', 'jpeg', 'tiff', 'tif'])188 189 if image is not None and not st.session_state.file_uploaded:190 try:191 bytes_data = image.getvalue()192 timestamp = int(time.time())193 original_filename = image.name194 file_extension = os.path.splitext(original_filename)[1].lower()195 196 if file_extension in ['.tiff', '.tif']:197 filename = f"image_{timestamp}.tif"198 converted_filename = f"image_{timestamp}_converted.png"199 else:200 filename = f"image_{timestamp}.png"201 converted_filename = filename202 203 filepath = os.path.join(UPLOAD_DIR, filename)204 converted_filepath = os.path.join(UPLOAD_DIR, converted_filename)205 206 with open(filepath, "wb") as f:207 f.write(bytes_data)208 209 if file_extension in ['.tiff', '.tif']:210 st.info('Processing GeoTIFF image...')211 rgb_image = read_pansharpened_rgb(filepath)212 cv2.imwrite(converted_filepath, cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR))213 st.success(f'GeoTIFF converted to 8-bit image and saved as {converted_filename}')214 img = Image.open(converted_filepath)215 else:216 img = Image.open(filepath)217 img.save(converted_filepath)218 219 if os.path.exists(converted_filepath):220 st.success(f"Image saved successfully: {converted_filepath}")221 file_size = os.path.getsize(converted_filepath)222 st.write(f"File size: {file_size} bytes")223 else:224 st.error(f"Failed to save image: {converted_filepath}")225 226 st.image(img, caption='Uploaded Image', use_column_width=True)227 st.success(f'Image processed and saved as {converted_filename}')228 229 st.session_state.filename = converted_filename230 231 img_array = np.array(img)232 233 if img_array.shape[0] > 650 or img_array.shape[1] > 650:234 st.info('Large image detected. Using patch-based processing.')235 with st.spinner('Analyzing large image...'):236 full_mask = process_large_image(model, converted_filepath)237 else:238 st.info('Small image detected. Processing whole image at once.')239 with st.spinner('Analyzing image...'):240 img_transformed = transforms(img)241 prediction = predict(img_transformed)242 full_mask = (prediction > 0.5).astype(np.uint8) * 255243 244 full_mask = refine_mask(full_mask)#-----------------------------------------------------------------------245 mask_filename = f"mask_{timestamp}.png"246 mask_filepath = os.path.join(MASK_DIR, mask_filename)247 cv2.imwrite(mask_filepath, full_mask)248 st.session_state.mask_filename = mask_filename249 250 log_image_details(timestamp, converted_filename, mask_filename)251 252 st.session_state.file_uploaded = True253 st.success("Image processed successfully")254 255 except Exception as e:256 st.error(f"An error occurred: {str(e)}")257 st.error("Please check the logs for more details.")258 print(f"Error in upload_page: {str(e)}")259 260 if st.session_state.file_uploaded and st.button('View result'):261 if st.session_state.filename is None:262 st.error("Please upload an image before viewing the result.")263 else:264 st.success('Image analyzed')265 st.session_state.page = 'result'266 st.rerun()267def result_page():268 st.title('Analysis Result')269 270 if 'filename' not in st.session_state or 'mask_filename' not in st.session_state:271 st.error("No image or mask file found. Please upload and process an image first.")272 if st.button('Back to Upload'):273 st.session_state.page = 'upload'274 st.session_state.file_uploaded = False275 st.session_state.filename = None276 st.session_state.mask_filename = None277 st.rerun()278 return279 280 col1, col2 = st.columns(2)281 282 # Display original image283 original_img_path = os.path.join(UPLOAD_DIR, st.session_state.filename)284 if os.path.exists(original_img_path):285 original_img = Image.open(original_img_path)286 col1.image(original_img, caption='Original Image', use_column_width=True)287 else:288 col1.error(f"Original image file not found: {original_img_path}")289 290 # Display predicted mask291 mask_path = os.path.join(MASK_DIR, st.session_state.mask_filename)292 if os.path.exists(mask_path):293 mask = Image.open(mask_path)294 col2.image(mask, caption='Predicted Mask', use_column_width=True)295 else:296 col2.error(f"Predicted mask file not found: {mask_path}")297 298 st.subheader("Overlay with Area of Buildings (sqft)")299 300 # Display overlayed image301 if os.path.exists(original_img_path) and os.path.exists(mask_path):302 original_np = cv2.imread(original_img_path)303 mask_np = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)304 305 # Ensure mask is binary306 _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)307 308 # Resize mask to match original image size if necessary309 if original_np.shape[:2] != mask_np.shape[:2]:310 mask_np = cv2.resize(mask_np, (original_np.shape[1], original_np.shape[0]))311 312 # Process and overlay image313 overlay_img = process_and_overlay_image(original_np, mask_np, 'output.png')314 315 st.image(overlay_img, caption='Overlay Image', use_column_width=True)316 else:317 st.error("Image or mask file not found for overlay.")318 319 if st.button('Back to Upload'):320 st.session_state.page = 'upload'321 st.session_state.file_uploaded = False322 st.session_state.filename = None323 st.session_state.mask_filename = None324 st.rerun()325 326def main():327 st.title('Building area estimation')328 329 if 'page' not in st.session_state:330 st.session_state.page = 'upload'331 332 if st.session_state.page == 'upload':333 upload_page()334 elif st.session_state.page == 'result':335 result_page()336 337if __name__ == '__main__':338 main()