Update app.py

app.py

@@ -10,7 +10,13 @@ import time
 
 # افزودن مسیر مورد نیاز برای ماژول‌های CelebAMask-HQ
 celebamask_path = "/home/user/app/CelebAMask-HQ"
-sys.path.insert(0, os.path.join(celebamask_path, "face_parsing"))
+face_parsing_path = os.path.join(celebamask_path, "face_parsing")
+sys.path.insert(0, celebamask_path)
+sys.path.insert(0, face_parsing_path)
+
+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:
@@ -21,6 +27,19 @@ try:
 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)
+    
+    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
 
 # تنظیمات دستگاه
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -33,7 +52,7 @@ os.environ["HF_HOME"] = "/home/user/app/hf_cache"
 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)),
+    T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
 ])
 
 # کلاس‌های Face Parsing
@@ -63,8 +82,16 @@ class FaceParsingModel:
             state_dict = torch.load(model_path, map_location="cpu")
             
             # ایجاد مدل
-            self.model = unet()
-            self.model.load_state_dict(state_dict)
+            self.model = unet(n_channels=3, n_classes=19)
+            
+            # تطبیق state dict در صورت نیاز
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                if k.startswith('module.'):
+                    k = k[7:]  # حذف 'module.' اگر وجود دارد
+                new_state_dict[k] = v
+            
+            self.model.load_state_dict(new_state_dict, strict=False)
             self.model.eval()
             self.model.to(self.device)
             
@@ -72,6 +99,8 @@ class FaceParsingModel:
             
         except Exception as e:
             print(f"❌ Failed to load model: {e}")
+            import traceback
+            traceback.print_exc()
             self.model = None
     
     def predict(self, image):
@@ -96,15 +125,33 @@ class FaceParsingModel:
         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)
-
+            out = out[0].cpu().numpy()
+            
+            # تبدیل به تصویر رنگی
+            colored_mask = self.colorize_mask(out)
+            
         # ترکیب تصویر اصلی با ماسک
-        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)
-
-        return out, blended
+        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)
+
+        return colored_mask, blended
+
+    def colorize_mask(self, mask):
+        """رنگ‌آمیزی ماسک بر اساس کلاس‌ها"""
+        # پالت رنگ برای 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]
+        ]
+        
+        colored = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
+        for i in range(len(palette)):
+            colored[mask == i] = palette[i]
+        
+        return colored
 
 def initialize_app():
     """Initialize application"""
@@ -112,12 +159,14 @@ 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(os.path.join(celebamask_path, "face_parsing")))
+    print("[Info] face_parsing folder exists:", os.path.exists(face_parsing_path))
     print("[Info] Module import success:", IMPORT_SUCCESS)
     
     try:
         face_parser = FaceParsingModel()
-        return True, "Model loaded successfully", face_parser
+        success = face_parser.model is not None
+        status_msg = "Model loaded successfully" if success else "Model failed to load"
+        return success, status_msg, face_parser
     except Exception as e:
         print(f"[Error] Initialization failed: {e}")
         return False, f"Initialization failed: {e}", None
@@ -143,7 +192,6 @@ def process_image(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"""
         ✅ پردازش انجام شد!
@@ -158,97 +206,8 @@ def process_image(input_image):
     except Exception as e:
         error_msg = f"❌ خطا در پردازش تصویر: {str(e)}"
         print(error_msg)
+        import traceback
+        traceback.print_exc()
         return None, None, error_msg
 
-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:
-    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
-    )
+# ادامه کد مشابه قبل برای Gradio interface...