Update app.py

app.py

@@ -3,10 +3,12 @@ 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
 import time
+import cv2
 
 # افزودن مسیر مورد نیاز برای ماژول‌های CelebAMask-HQ
 celebamask_path = "/home/user/app/CelebAMask-HQ"
@@ -18,28 +20,95 @@ 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))
 
-# ایمپورت ماژول‌های مورد نیاز
-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}")
-    # تعریف توابع جایگزین در صورت نیاز
-    def unet(**kwargs):
-        from unet import UNet
-        return UNet(**kwargs)
+# تعریف معماری مدل مطابق با 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)
     
-    def generate_label(inputs, imsize=512):
-        pred_batch = []
-        for input in inputs:
-            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)
-        return pred_batch
+    return label_batch
 
 # تنظیمات دستگاه
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -50,9 +119,9 @@ os.environ["HF_HOME"] = "/home/user/app/hf_cache"
 
 # تعریف transform
 transform = T.Compose([
-    T.Resize((512, 512), interpolation=PIL.Image.NEAREST),
+    T.Resize((512, 512)),
     T.ToTensor(),
-    T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
 # کلاس‌های Face Parsing
@@ -78,20 +147,20 @@ class FaceParsingModel:
             )
             print(f"✅ Model downloaded to: {model_path}")
             
+            # ایجاد مدل با معماری صحیح
+            self.model = unet(n_channels=3, n_classes=19)
+            
             # لود state dict
             state_dict = torch.load(model_path, map_location="cpu")
             
-            # ایجاد مدل
-            self.model = unet(n_channels=3, n_classes=19)
-            
-            # تطبیق state dict در صورت نیاز
+            # اگر state dict از DataParallel باشد، module. را حذف می‌کنیم
             new_state_dict = {}
             for k, v in state_dict.items():
                 if k.startswith('module.'):
-                    k = k[7:]  # حذف 'module.' اگر وجود دارد
+                    k = k[7:]
                 new_state_dict[k] = v
             
-            self.model.load_state_dict(new_state_dict, strict=False)
+            self.model.load_state_dict(new_state_dict)
             self.model.eval()
             self.model.to(self.device)
             
@@ -124,27 +193,41 @@ class FaceParsingModel:
         # پیش‌بینی
         with torch.no_grad():
             out = self.model(data)
-            out = generate_label(out, 512)
-            out = out[0].cpu().numpy()
+            label_out = generate_label(out, 512)
+            mask = label_out[0].cpu().numpy()
             
-            # تبدیل به تصویر رنگی
-            colored_mask = self.colorize_mask(out)
+            # رنگ‌آمیزی ماسک
+            colored_mask = self.colorize_mask(mask)
             
         # ترکیب تصویر اصلی با ماسک
         resized_image = np.asarray(original_image.resize((512, 512)))
-        blended = resized_image * 0.6 + colored_mask * 0.4
-        blended = np.clip(blended, 0, 255).astype(np.uint8)
+        blended = cv2.addWeighted(resized_image, 0.7, colored_mask, 0.3, 0)
 
         return colored_mask, blended
 
     def colorize_mask(self, mask):
         """رنگ‌آمیزی ماسک بر اساس کلاس‌ها"""
-        # پالت رنگ برای 19 کلاس
+        # پالت رنگ برای 19 کلاس (متفاوت برای تشخیص بهتر)
         palette = [
-            [0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0],
-            [255, 0, 255], [0, 255, 255], [128, 0, 0], [0, 128, 0], [0, 0, 128],
-            [128, 128, 0], [128, 0, 128], [0, 128, 128], [128, 128, 128], [255, 128, 0],
-            [255, 0, 128], [128, 255, 0], [0, 255, 128], [255, 128, 128]
+            [0, 0, 0],        # background - سیاه
+            [255, 200, 200],  # skin - پوست
+            [0, 255, 0],      # l_brow - سبز
+            [0, 200, 0],      # r_brow - سبز تیره
+            [255, 0, 0],      # l_eye - قرمز
+            [200, 0, 0],      # r_eye - قرمز تیره
+            [255, 255, 0],    # eye_g - زرد
+            [0, 0, 255],      # l_ear - آبی
+            [0, 0, 200],      # r_ear - آبی تیره
+            [128, 0, 128],    # ear_r - بنفش
+            [255, 165, 0],    # nose - نارنجی
+            [255, 0, 255],    # mouth - صورتی
+            [200, 0, 200],    # u_lip - صورتی تیره
+            [165, 42, 42],    # l_lip - قهوه‌ای
+            [0, 255, 255],    # neck - فیروزه‌ای
+            [0, 200, 200],    # neck_l - فیروزه‌ای تیره
+            [128, 128, 128],  # cloth - خاکستری
+            [255, 255, 255],  # hair - سفید
+            [255, 215, 0]     # hat - طلایی
         ]
         
         colored = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
@@ -160,7 +243,6 @@ 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()
@@ -210,4 +292,73 @@ def process_image(input_image):
         traceback.print_exc()
         return None, None, error_msg
 
-# ادامه کد مشابه قبل برای Gradio interface...
+# ادامه کد Gradio (مشابه قبل)
+
+# ایجاد اینترفیس Gradio
+with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🎭 CelebAMask-HQ Face Parsing Demo
+    **آپلود یک تصویر صورت و دریافت خروجی Face Parsing**
+    
+    این مدل صورت را به 19 بخش مختلف تقسیم می‌کند (پوست، چشم، ابرو، بینی، دهان، مو و ...)
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="📷 تصویر ورودی",
+                type="filepath",
+                sources=["upload"],
+                height=300
+            )
+            process_btn = gr.Button("🚀 پردازش تصویر", variant="primary", size="lg")
+            
+            with gr.Accordion("ℹ️ وضعیت برنامه", open=False):
+                status_display = gr.Markdown(f"""
+                **وضعیت:**
+                - 🎯 مدل: {'✅ لود شده' if success else '❌ خطا در لود'}
+                - 💻 دستگاه: `{device}`
+                - 📦 ماژول‌ها: {'✅ ایمپورت شده' if IMPORT_SUCCESS else '❌ خطا در ایمپورت'}
+                - 🗂️ کلاس‌ها: {len(CELEBA_CLASSES)}
+                """)
+        
+        with gr.Column():
+            output_blended = gr.Image(
+                label="🎨 نتیجه ترکیبی (تصویر + ماسک)",
+                height=300
+            )
+            output_mask = gr.Image(
+                label="🎭 ماسک سگمنتیشن",
+                height=300
+            )
+    
+    with gr.Row():
+        info_output = gr.Textbox(
+            label="📊 اطلاعات پردازش",
+            lines=3,
+            max_lines=6
+        )
+    
+    with gr.Row():
+        gr.HTML(create_legend())
+    
+    # اتصال رویدادها
+    process_btn.click(
+        fn=process_image,
+        inputs=[input_image],
+        outputs=[output_blended, output_mask, info_output]
+    )
+    
+    input_image.upload(
+        fn=process_image,
+        inputs=[input_image],
+        outputs=[output_blended, output_mask, info_output]
+    )
+
+if __name__ == "__main__":
+    print("🚀 Starting Face Parsing Application...")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )