Update app.py

app.py

@@ -212,57 +212,78 @@ class ClimatePredictor:
             transforms.Normalize(mean=[0.5], std=[0.5])
         ])
 
+    def convert_to_single_channel(self, image_array):
+        """RGB 이미지를 단일 채널로 변환"""
+        if len(image_array.shape) == 3:
+            # RGB to grayscale conversion
+            return np.dot(image_array[...,:3], [0.2989, 0.5870, 0.1140])
+        return image_array
+
     def predict_from_inputs(self, rgb_image, ndvi_image, terrain_image, 
                           elevation_data, wind_speed, wind_direction, 
                           temperature, humidity):
         """Gradio 인터페이스용 예측 함수"""
-        # 이미지 전처리
-        rgb_tensor = self.rgb_transform(Image.fromarray(rgb_image)).unsqueeze(0)
-        ndvi_tensor = self.single_channel_transform(Image.fromarray(ndvi_image)).unsqueeze(0)
-        terrain_tensor = self.single_channel_transform(Image.fromarray(terrain_image)).unsqueeze(0)
-        
-        # 고도 데이터 처리
-        elevation_tensor = torch.from_numpy(elevation_data).float().unsqueeze(0).unsqueeze(0)
-        elevation_tensor = (elevation_tensor - elevation_tensor.min()) / (elevation_tensor.max() - elevation_tensor.min())
-        
-        # 기상 데이터 처리
-        weather_features = np.array([wind_speed, wind_direction, temperature, humidity])
-        weather_features = (weather_features - weather_features.min()) / (weather_features.max() - weather_features.min())
-        weather_features = torch.tensor(weather_features, dtype=torch.float32).unsqueeze(0)
-        
-        # 디바이스로 이동
-        sample = {
-            'rgb': rgb_tensor.to(self.device),
-            'ndvi': ndvi_tensor.to(self.device),
-            'terrain': terrain_tensor.to(self.device),
-            'elevation': elevation_tensor.to(self.device),
-            'weather_features': weather_features.to(self.device)
-        }
-        
-        # 예측
-        with torch.no_grad():
-            wind_pred, solar_pred = self.model(sample)
-        
-        # 결과를 numpy 배열로 변환
-        wind_map = wind_pred.cpu().numpy()[0, 0]
-        solar_map = solar_pred.cpu().numpy()[0, 0]
-        
-        # 결과 시각화
-        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
-        
-        # 풍력 발전 잠재량 시각화
-        sns.heatmap(wind_map, ax=ax1, cmap='YlOrRd', cbar_kws={'label': 'Wind Power Potential'})
-        ax1.set_title('Wind Power Potential Map')
-        
-        # 태양광 발전 잠재량 시각화
-        sns.heatmap(solar_map, ax=ax2, cmap='YlOrRd', cbar_kws={'label': 'Solar Power Potential'})
-        ax2.set_title('Solar Power Potential Map')
-        
-        plt.tight_layout()
-        
-        return fig
+        try:
+            # RGB 이미지 전처리
+            rgb_tensor = self.rgb_transform(Image.fromarray(rgb_image)).unsqueeze(0)
+            
+            # NDVI 이미지를 단일 채널로 변환 후 전처리
+            ndvi_gray = self.convert_to_single_channel(ndvi_image)
+            ndvi_tensor = self.single_channel_transform(Image.fromarray(ndvi_gray.astype(np.uint8))).unsqueeze(0)
+            
+            # Terrain 이미지를 단일 채널로 변환 후 전처리
+            terrain_gray = self.convert_to_single_channel(terrain_image)
+            terrain_tensor = self.single_channel_transform(Image.fromarray(terrain_gray.astype(np.uint8))).unsqueeze(0)
+            
+            # Print shapes for debugging
+            print(f"RGB tensor shape: {rgb_tensor.shape}")
+            print(f"NDVI tensor shape: {ndvi_tensor.shape}")
+            print(f"Terrain tensor shape: {terrain_tensor.shape}")
+            
+            # 고도 데이터 처리
+            elevation_tensor = torch.from_numpy(elevation_data).float().unsqueeze(0).unsqueeze(0)
+            elevation_tensor = (elevation_tensor - elevation_tensor.min()) / (elevation_tensor.max() - elevation_tensor.min())
+            
+            # 기상 데이터 처리
+            weather_features = np.array([wind_speed, wind_direction, temperature, humidity])
+            weather_features = (weather_features - weather_features.min()) / (weather_features.max() - weather_features.min())
+            weather_features = torch.tensor(weather_features, dtype=torch.float32).unsqueeze(0)
+            
+            # 디바이스로 이동
+            sample = {
+                'rgb': rgb_tensor.to(self.device),
+                'ndvi': ndvi_tensor.to(self.device),
+                'terrain': terrain_tensor.to(self.device),
+                'elevation': elevation_tensor.to(self.device),
+                'weather_features': weather_features.to(self.device)
+            }
+            
+            # 예측
+            with torch.no_grad():
+                wind_pred, solar_pred = self.model(sample)
+            
+            # 결과를 numpy 배열로 변환
+            wind_map = wind_pred.cpu().numpy()[0, 0]
+            solar_map = solar_pred.cpu().numpy()[0, 0]
+            
+            # 결과 시각화
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
+            
+            # 풍력 발전 잠재량 시각화
+            sns.heatmap(wind_map, ax=ax1, cmap='YlOrRd', cbar_kws={'label': 'Wind Power Potential'})
+            ax1.set_title('Wind Power Potential Map')
+            
+            # 태양광 발전 잠재량 시각화
+            sns.heatmap(solar_map, ax=ax2, cmap='YlOrRd', cbar_kws={'label': 'Solar Power Potential'})
+            ax2.set_title('Solar Power Potential Map')
+            
+            plt.tight_layout()
+            
+            return fig
+        except Exception as e:
+            print(f"Error in prediction: {str(e)}")
+            raise e
 
-# Gradio 인터페이스 생성
 def create_gradio_interface():
     predictor = ClimatePredictor('best_model.pth')
     
@@ -282,8 +303,8 @@ def create_gradio_interface():
         fn=predict_and_visualize,
         inputs=[
             gr.Image(label="RGB Satellite Image", type="numpy"),
-            gr.Image(label="NDVI Image", type="numpy"),
-            gr.Image(label="Terrain Map", type="numpy"),
+            gr.Image(label="NDVI Image (will be converted to grayscale)", type="numpy"),
+            gr.Image(label="Terrain Map (will be converted to grayscale)", type="numpy"),
             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),
@@ -292,7 +313,8 @@ def create_gradio_interface():
         ],
         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.",
+        description="""Upload satellite imagery and environmental data to predict wind and solar power potential.
+        Note: NDVI and Terrain images will be automatically converted to grayscale.""",
         examples=[
             [
                 "examples/rgb_example.png",
@@ -305,6 +327,7 @@ def create_gradio_interface():
     )
     return interface
 
+# Hugging Face Spaces에서 앱 실행
 if __name__ == "__main__":
     interface = create_gradio_interface()
-    interface.launch()
+    interface.launch()