Update app.py

app.py

@@ -1,194 +1,185 @@
 import os
-import cv2
-import torch
+import sys
 import numpy as np
-from PIL import Image
-import gradio as gr
+import PIL.Image
+import torch
+import torchvision.transforms as T
 from huggingface_hub import hf_hub_download
+import gradio as gr
 import time
-import matplotlib.pyplot as plt
+
+# افزودن مسیر مورد نیاز برای ماژول‌های CelebAMask-HQ
+celebamask_path = "/home/user/app/CelebAMask-HQ"
+sys.path.insert(0, os.path.join(celebamask_path, "face_parsing"))
+
+# ایمپورت ماژول‌های مورد نیاز
+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")
+print(f"Using device: {device}")
 
 # تنظیم مسیرهای کش
 os.environ["HF_HOME"] = "/home/user/app/hf_cache"
 
-# کلاس‌های Face Parsing (19 کلاس)
+# تعریف transform
+transform = T.Compose([
+    T.Resize((512, 512), interpolation=PIL.Image.NEAREST),
+    T.ToTensor(),
+    T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+])
+
+# کلاس‌های Face Parsing
 CELEBA_CLASSES = [
     'background', 'skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r',
     'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat'
 ]
 
-class FaceParsing:
-    def __init__(self, model_path):
-        self.model_path = model_path
+class FaceParsingModel:
+    def __init__(self):
         self.model = None
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = device
         self.load_model()
     
     def load_model(self):
         """لود مدل Face Parsing"""
         try:
-            # اینجا باید معماری مدل را بر اساس ریپوی اصلی تنظیم کنید
-            # برای سادگی از یک مدل ساده استفاده می‌کنیم
-            self.model = torch.jit.load(self.model_path, map_location=self.device)
+            print("📥 Downloading model...")
+            model_path = hf_hub_download(
+                repo_id="public-data/CelebAMask-HQ-Face-Parsing",
+                filename="models/model.pth",
+                cache_dir="/home/user/app/hf_cache"
+            )
+            print(f"✅ Model downloaded to: {model_path}")
+            
+            # لود state dict
+            state_dict = torch.load(model_path, map_location="cpu")
+            
+            # ایجاد مدل
+            self.model = unet()
+            self.model.load_state_dict(state_dict)
             self.model.eval()
-            print("[Success] Model loaded successfully")
+            self.model.to(self.device)
+            
+            print("✅ Model loaded successfully")
+            
         except Exception as e:
-            print(f"[Error] Failed to load model: {e}")
-            # اگر مدل قابل لود نیست، یک مدل ساده ایجاد می‌کنیم
-            self.model = SimpleFaceParser()
+            print(f"❌ Failed to load model: {e}")
+            self.model = None
     
-    def preprocess_image(self, image):
-        """پیش‌پردازش تصویر ورودی"""
-        # تبدیل به RGB اگر لازم است
+    def predict(self, image):
+        """پردازش تصویر و تولید ماسک"""
+        if self.model is None:
+            raise ValueError("Model not loaded properly")
+        
+        # تبدیل به PIL Image اگر لازم است
         if isinstance(image, str):
-            image = Image.open(image).convert('RGB')
+            image = PIL.Image.open(image).convert('RGB')
         elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-        
-        # تغییر سایز به 512x512
-        image = image.resize((512, 512))
+            image = PIL.Image.fromarray(image)
         
-        # تبدیل به tensor و نرمال‌سازی
-        image_tensor = torch.from_numpy(np.array(image)).float() / 255.0
-        image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)
+        # ذخیره تصویر اصلی
+        original_image = image.copy()
         
-        return image_tensor.to(self.device), image
-    
-    def postprocess_mask(self, mask):
-        """پس‌پردازش ماسک خروجی"""
-        mask = mask.squeeze().cpu().numpy()
-        return mask.astype(np.uint8)
-    
-    def predict(self, image):
-        """پیش‌بینی روی تصویر ورودی"""
-        try:
-            image_tensor, original_image = self.preprocess_image(image)
-            
-            with torch.no_grad():
-                if hasattr(self.model, 'predict'):
-                    output = self.model.predict(image_tensor)
-                else:
-                    output = self.model(image_tensor)
-            
-            # گرفتن ماسک پیش‌بینی شده
-            if isinstance(output, tuple):
-                mask = output[0]
-            else:
-                mask = output
-            
-            parsed_mask = self.postprocess_mask(mask)
-            return self.visualize_result(np.array(original_image), parsed_mask)
-            
-        except Exception as e:
-            print(f"[Error] Prediction failed: {e}")
-            # بازگشت تصویر اصلی در صورت خطا
-            if isinstance(image, str):
-                original_img = Image.open(image)
-            else:
-                original_img = image
-            return original_img, original_img
+        # پیش‌پردازش
+        data = transform(image)
+        data = data.unsqueeze(0).to(self.device)
+
+        # پیش‌بینی
+        with torch.no_grad():
+            out = self.model(data)
+            out = generate_label(out, 512)
+            out = out[0].cpu().numpy().transpose(1, 2, 0)
+            out = np.clip(np.round(out * 255), 0, 255).astype(np.uint8)
 
-    def visualize_result(self, original_image, mask):
-        """ویژوالایز کردن نتایج"""
-        # ایجاد تصویر رنگی از ماسک
-        colored_mask = self.colorize_mask(mask)
-        
         # ترکیب تصویر اصلی با ماسک
-        overlay = cv2.addWeighted(original_image, 0.7, colored_mask, 0.3, 0)
-        
-        return overlay, colored_mask
-    
-    def colorize_mask(self, mask):
-        """رنگ‌آمیزی ماسک بر اساس کلاس‌ها"""
-        # ایجاد پالت رنگ برای کلاس‌ها
-        cmap = plt.get_cmap('tab20', len(CELEBA_CLASSES))
-        colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
-        
-        for i in range(len(CELEBA_CLASSES)):
-            colored_mask[mask == i] = np.array(cmap(i)[:3]) * 255
-        
-        return colored_mask.astype(np.uint8)
+        resized_image = np.asarray(original_image.resize((512, 512))).astype(float)
+        blended = resized_image * 0.5 + out.astype(float) * 0.5
+        blended = np.clip(np.round(blended), 0, 255).astype(np.uint8)
 
-class SimpleFaceParser:
-    """یک پارسر ساده برای مواقعی که مدل اصلی کار نمی‌کند"""
-    def __init__(self):
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    
-    def __call__(self, x):
-        # یک خروجی ساده شبیه ماسک تولید می‌کند
-        batch_size, channels, height, width = x.shape
-        return torch.randint(0, len(CELEBA_CLASSES), (batch_size, 1, height, width)).float().to(self.device)
+        return out, blended
 
 def initialize_app():
-    """Initialize application and download models"""
+    """Initialize application"""
     print("===== Application Startup at {} =====".format(time.strftime("%Y-%m-%d %H:%M:%S")))
     
-    celeb_path = "/home/user/app/huggingface_models/CelebAMask-HQ"
-    face_parsing_path = os.path.join(celeb_path, "face_parsing")
-
     print("[Info] PYTHONPATH:", os.environ.get("PYTHONPATH"))
-    print("[Info] CelebAMask-HQ path exists:", os.path.exists(celeb_path))
-    print("[Info] face_parsing folder exists:", os.path.exists(face_parsing_path))
-
+    print("[Info] CelebAMask-HQ path exists:", os.path.exists(celebamask_path))
+    print("[Info] face_parsing folder exists:", os.path.exists(os.path.join(celebamask_path, "face_parsing")))
+    print("[Info] Module import success:", IMPORT_SUCCESS)
+    
     try:
-        model_path = hf_hub_download(
-            repo_id="public-data/CelebAMask-HQ-Face-Parsing",
-            filename="models/model.pth",
-            cache_dir="/home/user/app/hf_cache"
-        )
-        print("[Success] Model downloaded to:", model_path)
-        
-        # لود مدل
-        face_parser = FaceParsing(model_path)
-        return True, model_path, face_parser
-        
+        face_parser = FaceParsingModel()
+        return True, "Model loaded successfully", face_parser
     except Exception as e:
-        print("[Error] Failed to download model:", str(e))
-        # استفاده از پارسر ساده
-        face_parser = FaceParsing(None)
-        return False, str(e), face_parser
+        print(f"[Error] Initialization failed: {e}")
+        return False, f"Initialization failed: {e}", None
 
 # Initialize the application
-success, model_info, face_parser = initialize_app()
+success, status_msg, face_parser = initialize_app()
 
 def process_image(input_image):
-    """پردازش تصویر ورودی و بازگشت نتایج"""
+    """پردازش تصویر ورودی"""
     if input_image is None:
         return None, None, "لطفاً یک تصویر آپلود کنید"
     
+    if not success or face_parser is None:
+        return None, None, "❌ مدل لود نشده است. لطفاً دوباره تلاش کنید."
+    
     try:
         # پردازش تصویر
-        overlay_result, mask_result = face_parser.predict(input_image)
+        mask, blended = face_parser.predict(input_image)
         
-        # اطلاعات درباره تصویر
+        # اطلاعات پردازش
         if isinstance(input_image, str):
-            img_size = Image.open(input_image).size
+            original_img = PIL.Image.open(input_image)
+            img_size = original_img.size
         else:
             img_size = input_image.size if hasattr(input_image, 'size') else input_image.shape[:2][::-1]
+            original_img = PIL.Image.fromarray(input_image) if isinstance(input_image, np.ndarray) else input_image
         
         info_text = f"""
         ✅ پردازش انجام شد!
         - اندازه تصویر ورودی: {img_size}
-        - مدل: {'CelebAMask-HQ' if success else 'Simple Parser'}
+        - اندازه خروجی: 512x512
         - کلاس‌های تشخیص: {len(CELEBA_CLASSES)}
+        - دستگاه پردازش: {device}
         """
         
-        return overlay_result, mask_result, info_text
+        return blended, mask, info_text
         
     except Exception as e:
-        error_msg = f"خطا در پردازش تصویر: {str(e)}"
+        error_msg = f"❌ خطا در پردازش تصویر: {str(e)}"
         print(error_msg)
         return None, None, error_msg
 
 def create_legend():
     """ایجاد لیجند برای کلاس‌ها"""
-    legend_html = "<div style='max-height: 300px; overflow-y: auto;'><h4>Legend - کلاس‌های Face Parsing:</h4>"
+    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))
-        legend_html += f"<div style='margin: 2px; padding: 5px; background-color: {color_hex}; color: white;'>{i}: {class_name}</div>"
+        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
@@ -198,6 +189,8 @@ with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as de
     gr.Markdown("""
     # 🎭 CelebAMask-HQ Face Parsing Demo
     **آپلود یک تصویر صورت و دریافت خروجی Face Parsing**
+    
+    این مدل صورت را به 19 بخش مختلف تقسیم می‌کند (پوست، چشم، ابرو، بینی، دهان، مو و ...)
     """)
     
     with gr.Row():
@@ -205,26 +198,27 @@ with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as de
             input_image = gr.Image(
                 label="📷 تصویر ورودی",
                 type="filepath",
-                sources=["upload", "webcam"],
+                sources=["upload"],
                 height=300
             )
-            process_btn = gr.Button("🚀 پردازش تصویر", variant="primary")
+            process_btn = gr.Button("🚀 پردازش تصویر", variant="primary", size="lg")
             
-            with gr.Accordion("ℹ️ اطلاعات برنامه", open=False):
-                status_text = gr.Markdown(f"""
-                **وضعیت برنامه:**
-                - مدل: {'✅ موفق' if success else '⚠️ ساده'}
-                - مسیر مدل: `{model_info if success else 'مدل پیش‌فرض'}`
-                - کلاس‌های تشخیص: {len(CELEBA_CLASSES)}
+            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_overlay = gr.Image(
-                label="🎨 نتیجه ترکیبی (Overlay)",
+            output_blended = gr.Image(
+                label="🎨 نتیجه ترکیبی (تصویر + ماسک)",
                 height=300
             )
             output_mask = gr.Image(
-                label="🎭 ماسک segmentation",
+                label="🎭 ماسک سگمنتیشن",
                 height=300
             )
     
@@ -242,18 +236,17 @@ with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as de
     process_btn.click(
         fn=process_image,
         inputs=[input_image],
-        outputs=[output_overlay, output_mask, info_output]
+        outputs=[output_blended, output_mask, info_output]
     )
     
-    # پردازش خودکار هنگام آپلود تصویر
     input_image.upload(
         fn=process_image,
         inputs=[input_image],
-        outputs=[output_overlay, output_mask, info_output]
+        outputs=[output_blended, output_mask, info_output]
     )
 
 if __name__ == "__main__":
-    print("Starting Face Parsing Application...")
+    print("🚀 Starting Face Parsing Application...")
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,