| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torchvision import transforms, models | |
| from PIL import Image | |
| from TranSalNet_Res import TranSalNet | |
| from tqdm import tqdm | |
| import torch.nn as nn | |
| from utils.data_process import preprocess_img, postprocess_img | |
| device = torch.device('cpu') | |
| model = TranSalNet() | |
| model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu'))) | |
| model.to(device) | |
| model.eval() | |
| def count_and_label_red_patches(heatmap, threshold=200): | |
| red_mask = heatmap[:, :, 2] > threshold | |
| _, labels, stats, _ = cv2.connectedComponentsWithStats(red_mask.astype(np.uint8), connectivity=8) | |
| num_red_patches = labels.max() | |
| for i in range(1, num_red_patches + 1): | |
| patch_mask = (labels == i) | |
| patch_centroid_x, patch_centroid_y = int(stats[i, cv2.CC_STAT_LEFT] + stats[i, cv2.CC_STAT_WIDTH] / 2), int(stats[i, cv2.CC_STAT_TOP] + stats[i, cv2.CC_STAT_HEIGHT] / 2) | |
| cv2.putText(heatmap, str(i), (patch_centroid_x, patch_centroid_y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) | |
| return heatmap, num_red_patches | |
| st.title('Saliency Detection App') | |
| st.write('Upload an image for saliency detection:') | |
| uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_image: | |
| image = Image.open(uploaded_image) | |
| st.image(image, caption='Uploaded Image', use_column_width=True) | |
| if st.button('Detect Saliency'): | |
| img = image.resize((384, 288)) | |
| img = np.array(img) / 255. | |
| img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0) | |
| img = torch.from_numpy(img) | |
| img = img.type(torch.FloatTensor).to(device) | |
| pred_saliency = model(img).squeeze().detach().numpy() | |
| heatmap = (pred_saliency * 255).astype(np.uint8) | |
| heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) | |
| heatmap = cv2.resize(heatmap, (image.width, image.height)) | |
| heatmap, num_red_patches = count_and_label_red_patches(heatmap) | |
| enhanced_image = np.array(image) | |
| b, g, r = cv2.split(enhanced_image) | |
| clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) | |
| b_enhanced = clahe.apply(b) | |
| enhanced_image = cv2.merge((b_enhanced, g, r)) | |
| alpha = 0.7 | |
| blended_img = cv2.addWeighted(enhanced_image, 1 - alpha, heatmap, alpha, 0) | |
| st.image(heatmap, caption='Enhanced Saliency Heatmap', use_column_width=True, channels='BGR') | |
| st.image(enhanced_image, caption='Enhanced Blue Image', use_column_width=True, channels='BGR') | |
| st.image(blended_img, caption=f'Blended Image with {num_red_patches} Red Patches', use_column_width=True, channels='BGR') | |
| cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200]) | |
| st.success('Saliency detection complete. Result saved as "example/result15.png".') | |