Update app.py

app.py

@@ -3,7 +3,6 @@ import sys
 import numpy as np
 import PIL.Image
 import torch
-import torch.nn as nn
 import torchvision.transforms as T
 from huggingface_hub import hf_hub_download
 import gradio as gr
@@ -20,95 +19,15 @@ print("Python path:", sys.path)
 print("CelebAMask path exists:", os.path.exists(celebamask_path))
 print("Face parsing path exists:", os.path.exists(face_parsing_path))
 
-# تعریف معماری مدل مطابق با state dict دانلود شده
-class SimpleFaceParser(nn.Module):
-    def __init__(self, n_channels=3, n_classes=19):
-        super(SimpleFaceParser, self).__init__()
-        
-        def conv_block(in_channels, out_channels):
-            return nn.Sequential(
-                nn.Conv2d(in_channels, out_channels, 3, padding=1),
-                nn.BatchNorm2d(out_channels),
-                nn.ReLU(inplace=True),
-                nn.Conv2d(out_channels, out_channels, 3, padding=1),
-                nn.BatchNorm2d(out_channels),
-                nn.ReLU(inplace=True)
-            )
-        
-        # Encoder
-        self.enc1 = conv_block(n_channels, 16)
-        self.enc2 = conv_block(16, 32)
-        self.enc3 = conv_block(32, 64) 
-        self.enc4 = conv_block(64, 128)
-        self.enc5 = conv_block(128, 256)
-        
-        # Decoder
-        self.dec4 = conv_block(256 + 128, 128)
-        self.dec3 = conv_block(128 + 64, 64)
-        self.dec2 = conv_block(64 + 32, 32)
-        self.dec1 = conv_block(32 + 16, 16)
-        
-        # Pooling and upsample
-        self.pool = nn.MaxPool2d(2)
-        self.upsample4 = nn.ConvTranspose2d(256, 128, 2, 2)
-        self.upsample3 = nn.ConvTranspose2d(128, 64, 2, 2)
-        self.upsample2 = nn.ConvTranspose2d(64, 32, 2, 2)
-        self.upsample1 = nn.ConvTranspose2d(32, 16, 2, 2)
-        
-        # Final layer
-        self.final = nn.Conv2d(16, n_classes, 1)
-        
-    def forward(self, x):
-        # Encoder
-        e1 = self.enc1(x)
-        e2 = self.enc2(self.pool(e1))
-        e3 = self.enc3(self.pool(e2))
-        e4 = self.enc4(self.pool(e3))
-        e5 = self.enc5(self.pool(e4))
-        
-        # Decoder with skip connections
-        d4 = self.upsample4(e5)
-        d4 = torch.cat([d4, e4], dim=1)
-        d4 = self.dec4(d4)
-        
-        d3 = self.upsample3(d4)
-        d3 = torch.cat([d3, e3], dim=1)
-        d3 = self.dec3(d3)
-        
-        d2 = self.upsample2(d3)
-        d2 = torch.cat([d2, e2], dim=1)
-        d2 = self.dec2(d2)
-        
-        d1 = self.upsample1(d2)
-        d1 = torch.cat([d1, e1], dim=1)
-        d1 = self.dec1(d1)
-        
-        return self.final(d1)
-
-def unet(**kwargs):
-    return SimpleFaceParser(**kwargs)
-
-# تابع generate_label
-def generate_label(inputs, imsize=512):
-    """Generate label maps from model outputs"""
-    pred_batch = []
-    for input in inputs:
-        input = input.unsqueeze(0)
-        pred = np.squeeze(input.data.max(1)[1].cpu().numpy(), axis=0)
-        pred_batch.append(pred)
-
-    pred_batch = np.array(pred_batch)
-    pred_batch = torch.from_numpy(pred_batch)
-            
-    label_batch = []
-    for p in pred_batch:
-        p = p.view(1, imsize, imsize)
-        label_batch.append(p.data.cpu())
-    
-    label_batch = torch.cat(label_batch, 0)
-    label_batch = label_batch.type(torch.LongTensor)
-    
-    return label_batch
+# ایمپورت ماژول‌های مورد نیاز
+try:
+    from unet import unet
+    from utils import generate_label
+    IMPORT_SUCCESS = True
+    print("✅ Successfully imported CelebAMask-HQ modules")
+except ImportError as e:
+    IMPORT_SUCCESS = False
+    print(f"❌ Failed to import CelebAMask-HQ modules: {e}")
 
 # تنظیمات دستگاه
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -148,7 +67,13 @@ class FaceParsingModel:
             print(f"✅ Model downloaded to: {model_path}")
             
             # ایجاد مدل با معماری صحیح
-            self.model = unet(n_channels=3, n_classes=19)
+            self.model = unet(
+                feature_scale=4,
+                n_classes=19,
+                is_deconv=True,
+                in_channels=3,
+                is_batchnorm=True
+            )
             
             # لود state dict
             state_dict = torch.load(model_path, map_location="cpu")
@@ -243,6 +168,7 @@ def initialize_app():
     print("[Info] PYTHONPATH:", os.environ.get("PYTHONPATH"))
     print("[Info] CelebAMask-HQ path exists:", os.path.exists(celebamask_path))
     print("[Info] face_parsing folder exists:", os.path.exists(face_parsing_path))
+    print("[Info] Module import success:", IMPORT_SUCCESS)
     
     try:
         face_parser = FaceParsingModel()
@@ -292,9 +218,29 @@ def process_image(input_image):
         traceback.print_exc()
         return None, None, error_msg
 
-# ادامه کد Gradio (مشابه قبل)
-
-# ادامه کد Gradio (مشابه قبل)
+def create_legend():
+    """ایجاد لیجند برای کلاس‌ها"""
+    import matplotlib.pyplot as plt
+    
+    legend_html = """
+    <div style='max-height: 300px; overflow-y: auto; border: 1px solid #ccc; padding: 10px; border-radius: 5px;'>
+    <h4>🎨 Legend - کلاس‌های Face Parsing:</h4>
+    """
+    
+    colors = plt.get_cmap('tab20', len(CELEBA_CLASSES))
+    
+    for i, class_name in enumerate(CELEBA_CLASSES):
+        color = colors(i)
+        color_hex = '#%02x%02x%02x' % (int(color[0]*255), int(color[1]*255), int(color[2]*255))
+        text_color = 'white' if color[0] * 0.299 + color[1] * 0.587 + color[2] * 0.114 < 0.5 else 'black'
+        legend_html += f"""
+        <div style='margin: 2px; padding: 5px; background-color: {color_hex}; color: {text_color}; border-radius: 3px;'>
+        <strong>{i}:</strong> {class_name}
+        </div>
+        """
+    
+    legend_html += "</div>"
+    return legend_html
 
 # ایجاد اینترفیس Gradio
 with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as demo: