Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageDraw
 import torchvision.transforms as transforms
 import pandas as pd
 import os
@@ -11,11 +11,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from torch.optim.lr_scheduler import ReduceLROnPlateau
-import numpy as np
-import pandas as pd
-from PIL import Image
-import os
-import torchvision.transforms as transforms
 from sklearn.model_selection import train_test_split
 import glob
 
@@ -24,11 +19,8 @@ class ClimateNet(nn.Module):
         super(ClimateNet, self).__init__()
         self.input_size = input_size
         self.output_size = output_size
-
-        # Feature map sizes after two max pooling layers
         self.feature_size = (input_size[0] // 4, input_size[1] // 4)
 
-        # Improved RGB Encoder with residual connections
         self.rgb_encoder = nn.Sequential(
             nn.Conv2d(3, 64, kernel_size=3, padding=1),
             nn.BatchNorm2d(64),
@@ -49,7 +41,6 @@ class ClimateNet(nn.Module):
             nn.Dropout2d(0.2)
         )
 
-        # Improved NDVI Encoder
         self.ndvi_encoder = nn.Sequential(
             nn.Conv2d(1, 64, kernel_size=3, padding=1),
             nn.BatchNorm2d(64),
@@ -67,7 +58,6 @@ class ClimateNet(nn.Module):
             nn.Dropout2d(0.2)
         )
 
-        # Improved Terrain Encoder
         self.terrain_encoder = nn.Sequential(
             nn.Conv2d(1, 64, kernel_size=3, padding=1),
             nn.BatchNorm2d(64),
@@ -85,7 +75,6 @@ class ClimateNet(nn.Module):
             nn.Dropout2d(0.2)
         )
 
-        # Improved Weather Encoder with deeper architecture
         self.weather_encoder = nn.Sequential(
             nn.Linear(4, 64),
             nn.ReLU(),
@@ -96,7 +85,6 @@ class ClimateNet(nn.Module):
             nn.Linear(128, 128)
         )
 
-        # Improved Feature Fusion
         self.fusion = nn.Sequential(
             nn.Conv2d(512, 512, kernel_size=1),
             nn.BatchNorm2d(512),
@@ -107,7 +95,6 @@ class ClimateNet(nn.Module):
             nn.ReLU()
         )
 
-        # Improved Decoders with skip connections
         self.wind_decoder = nn.Sequential(
             nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
             nn.BatchNorm2d(256),
@@ -143,18 +130,15 @@ class ClimateNet(nn.Module):
     def forward(self, x):
         batch_size = x['rgb'].size(0)
 
-        # Resize all inputs to input_size
         rgb_input = F.interpolate(x['rgb'], size=self.input_size, mode='bilinear', align_corners=False)
         ndvi_input = F.interpolate(x['ndvi'], size=self.input_size, mode='bilinear', align_corners=False)
         terrain_input = F.interpolate(x['terrain'], size=self.input_size, mode='bilinear', align_corners=False)
 
-        # Extract features
-        rgb_features = self.rgb_encoder(rgb_input)  # [B, 128, H/4, W/4]
-        ndvi_features = self.ndvi_encoder(ndvi_input)  # [B, 128, H/4, W/4]
-        terrain_features = self.terrain_encoder(terrain_input)  # [B, 128, H/4, W/4]
+        rgb_features = self.rgb_encoder(rgb_input)
+        ndvi_features = self.ndvi_encoder(ndvi_input)
+        terrain_features = self.terrain_encoder(terrain_input)
 
-        # Process weather features and expand to match feature map size
-        weather_features = self.weather_encoder(x['weather_features'])  # [B, 128]
+        weather_features = self.weather_encoder(x['weather_features'])
         weather_features = weather_features.view(batch_size, 128, 1, 1)
         weather_features = F.interpolate(
             weather_features,
@@ -162,7 +146,6 @@ class ClimateNet(nn.Module):
             mode='nearest'
         )
 
-        # Combine features
         combined_features = torch.cat([
             rgb_features,
             ndvi_features,
@@ -170,10 +153,8 @@ class ClimateNet(nn.Module):
             weather_features
         ], dim=1)
 
-        # Apply fusion
         fused_features = self.fusion(combined_features)
 
-        # Generate predictions and resize to output_size
         wind_heatmap = self.wind_decoder(fused_features)
         solar_heatmap = self.solar_decoder(fused_features)
 
@@ -191,7 +172,6 @@ class ClimatePredictor:
             
         print(f"Using device: {self.device}")
         
-        # Load model
         self.model = ClimateNet(input_size=(256, 256), output_size=(64, 64)).to(self.device)
         checkpoint = torch.load(model_path, map_location=self.device)
         
@@ -214,15 +194,44 @@ class ClimatePredictor:
         ])
 
     def convert_to_single_channel(self, image_array):
-        """RGB 이미지를 단일 채널로 변환"""
         if len(image_array.shape) == 3:
             return np.dot(image_array[...,:3], [0.2989, 0.5870, 0.1140])
         return image_array
 
+    def overlay_emojis(self, rgb_image, wind_map, solar_map, threshold=0.7):
+        """히트맵을 기반으로 이모지 오버레이"""
+        img = Image.fromarray(rgb_image)
+        draw = ImageDraw.Draw(img)
+        
+        h, w = rgb_image.shape[:2]
+        wind_map_resized = Image.fromarray((wind_map * 255).astype(np.uint8)).resize((w, h))
+        solar_map_resized = Image.fromarray((solar_map * 255).astype(np.uint8)).resize((w, h))
+        
+        wind_map_np = np.array(wind_map_resized) / 255.0
+        solar_map_np = np.array(solar_map_resized) / 255.0
+        
+        emoji_size = min(w, h) // 20
+        grid_step = emoji_size * 2
+        
+        for y in range(0, h - emoji_size, grid_step):
+            for x in range(0, w - emoji_size, grid_step):
+                wind_val = wind_map_np[y:y+emoji_size, x:x+emoji_size].mean()
+                solar_val = solar_map_np[y:y+emoji_size, x:x+emoji_size].mean()
+                
+                text = ""
+                if wind_val > threshold:
+                    text += "💨"
+                if solar_val > threshold:
+                    text += "☀️"
+                
+                if text:
+                    draw.text((x, y), text, fill="white")
+        
+        return np.array(img)
+
     def predict_from_inputs(self, rgb_image, ndvi_image, terrain_image, 
                           elevation_data, wind_speed, wind_direction, 
                           temperature, humidity):
-        """Gradio 인터페이스용 예측 함수"""
         try:
             # RGB 이미지 전처리
             rgb_tensor = self.rgb_transform(Image.fromarray(rgb_image)).unsqueeze(0)
@@ -262,17 +271,37 @@ class ClimatePredictor:
             solar_map = solar_pred.cpu().numpy()[0, 0]
             
             # 결과 시각화
-            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
+            fig = plt.figure(figsize=(20, 10))
             
-            sns.heatmap(wind_map, ax=ax1, cmap='YlOrRd', cbar_kws={'label': 'Wind Power Potential'})
-            ax1.set_title('Wind Power Potential Map')
+            # 1. 원본 이미지와 이모지 오버레이
+            ax1 = plt.subplot(2, 2, 1)
+            overlay_img = self.overlay_emojis(rgb_image, wind_map, solar_map)
+            ax1.imshow(overlay_img)
+            ax1.set_title('Predictions Overlay on RGB Image')
+            ax1.axis('off')
             
-            sns.heatmap(solar_map, ax=ax2, cmap='YlOrRd', cbar_kws={'label': 'Solar Power Potential'})
-            ax2.set_title('Solar Power Potential Map')
+            # 2. 풍력 발전 잠재량 히트맵
+            ax2 = plt.subplot(2, 2, 2)
+            wind_heatmap = sns.heatmap(wind_map, ax=ax2, cmap='YlOrRd', 
+                                     cbar_kws={'label': 'Wind Power Potential'})
+            ax2.set_title('Wind Power Potential Map')
             
-            plt.tight_layout()
+            # 3. 태양광 발전 잠재량 히트맵
+            ax3 = plt.subplot(2, 2, 3)
+            solar_heatmap = sns.heatmap(solar_map, ax=ax3, cmap='YlOrRd',
+                                      cbar_kws={'label': 'Solar Power Potential'})
+            ax3.set_title('Solar Power Potential Map')
+            
+            # 4. 컴바인드 히트맵
+            ax4 = plt.subplot(2, 2, 4)
+            combined_map = np.stack([solar_map, wind_map, np.zeros_like(wind_map)], axis=-1)
+            ax4.imshow(combined_map)
+            ax4.set_title('Combined Potential Map (Red: Solar, Green: Wind)')
+            ax4.axis('off')
             
+            plt.tight_layout()
             return fig
+            
         except Exception as e:
             print(f"Error in prediction: {str(e)}")
             raise e
@@ -284,21 +313,18 @@ def load_examples_from_directory(base_dir):
     
     for sample_dir in sample_dirs:
         try:
-            # 파일 경로 구성
             rgb_path = os.path.join(sample_dir, "satellite", "sentinel2_rgb_2023-07-15_to_2023-09-01.png")
             ndvi_path = os.path.join(sample_dir, "satellite", "sentinel2_ndvi_2023-07-15_to_2023-09-01.png")
             terrain_path = os.path.join(sample_dir, "terrain", "terrain_map.png")
             elevation_path = os.path.join(sample_dir, "terrain", "elevation_data.npy")
             weather_path = os.path.join(sample_dir, "weather", "weather_data.csv")
             
-            # 기상 데이터 읽기
             weather_data = pd.read_csv(weather_path)
             wind_speed = weather_data['wind_speed'].mean()
             wind_direction = weather_data['wind_direction'].mean()
             temperature = weather_data['temperature'].mean()
             humidity = weather_data['humidity'].mean()
             
-            # 예제 리스트에 추가
             examples.append([
                 rgb_path,
                 ndvi_path,
@@ -318,46 +344,52 @@ def load_examples_from_directory(base_dir):
 def create_gradio_interface():
     predictor = ClimatePredictor('best_model.pth')
     
-    def predict_and_visualize(rgb_path, ndvi_path, terrain_path, elevation_path,
-                            wind_speed, wind_direction, temperature, humidity):
-        # 이미지 로드
-        rgb_image = np.array(Image.open(rgb_path))
-        ndvi_image = np.array(Image.open(ndvi_path))
-        terrain_image = np.array(Image.open(terrain_path))
-        elevation_data = np.load(elevation_path)
+    with gr.Blocks() as interface:
+        gr.Markdown("# Renewable Energy Potential Predictor")
+        gr.Markdown("Upload satellite imagery and environmental data to predict wind and solar power potential.")
         
-        # 예측 및 시각화
-        result = predictor.predict_from_inputs(
-            rgb_image, ndvi_image, terrain_image, elevation_data,
-            wind_speed, wind_direction, temperature, humidity
+        with gr.Row():
+            # 입력 섹션 (1/3 크기)
+            with gr.Column(scale=1):
+                rgb_input = gr.Image(label="RGB Satellite Image", type="numpy")
+                ndvi_input = gr.Image(label="NDVI Image", type="numpy")
+                terrain_input = gr.Image(label="Terrain Map", type="numpy")
+                elevation_input = gr.File(label="Elevation Data (NPY file)")
+                
+                with gr.Row():
+                    wind_speed = gr.Number(label="Wind Speed (m/s)", value=5.0)
+                    wind_direction = gr.Number(label="Wind Direction (°)", value=180.0)
+                
+                with gr.Row():
+                    temperature = gr.Number(label="Temperature (°C)", value=25.0)
+                    humidity = gr.Number(label="Humidity (%)", value=60.0)
+                
+                predict_btn = gr.Button("Generate Predictions", variant="primary")
+            
+            # 출력 섹션 (2/3 크기)
+            with gr.Column(scale=2):
+                output_plot = gr.Plot(label="Prediction Results")
+        
+        # 예측 버튼 클릭 이벤트 연결
+        predict_btn.click(
+            fn=predictor.predict_from_inputs,
+            inputs=[rgb_input, ndvi_input, terrain_input, elevation_input,
+                   wind_speed, wind_direction, temperature, humidity],
+            outputs=output_plot
+        )
+        
+        # 예제 섹션
+        examples = load_examples_from_directory("filtered_climate_data")
+        gr.Examples(
+            examples=examples,
+            inputs=[rgb_input, ndvi_input, terrain_input, elevation_input,
+                   wind_speed, wind_direction, temperature, humidity],
+            outputs=output_plot,
+            cache_examples=True
         )
-        return result
-    
-    # 예제 데이터 로드
-    examples = load_examples_from_directory("filtered_climate_data")
-    print(f"Loaded {len(examples)} examples")
     
-    interface = gr.Interface(
-        fn=predict_and_visualize,
-        inputs=[
-            gr.Image(label="RGB Satellite Image", type="filepath"),
-            gr.Image(label="NDVI Image", type="filepath"),
-            gr.Image(label="Terrain Map", type="filepath"),
-            gr.File(label="Elevation Data (NPY file)"),
-            gr.Number(label="Wind Speed (m/s)", value=5.0),
-            gr.Number(label="Wind Direction (degrees)", value=180.0),
-            gr.Number(label="Temperature (°C)", value=25.0),
-            gr.Number(label="Humidity (%)", value=60.0)
-        ],
-        outputs=gr.Plot(label="Prediction Results"),
-        title="Renewable Energy Potential Predictor",
-        description="""Upload satellite imagery and environmental data to predict wind and solar power potential.
-        You can also try various examples from our dataset using the Examples section below.""",
-        examples=examples,
-        cache_examples=True
-    )
     return interface
 
 if __name__ == "__main__":
     interface = create_gradio_interface()
-    interface.launch()
+    interface.launch()