|
|
import os |
|
|
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_path = "/home/user/app/CelebAMask-HQ" |
|
|
face_parsing_path = os.path.join(celebamask_path, "face_parsing") |
|
|
sys.path.insert(0, celebamask_path) |
|
|
sys.path.insert(0, face_parsing_path) |
|
|
|
|
|
|
|
|
try: |
|
|
from unet import unet as celebamask_unet |
|
|
from utils import generate_label |
|
|
HAS_CELEBAMASK = True |
|
|
except ImportError: |
|
|
HAS_CELEBAMASK = False |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
os.environ["HF_HOME"] = "/home/user/app/hf_cache" |
|
|
|
|
|
|
|
|
class BiSeNet(nn.Module): |
|
|
def __init__(self, n_classes=19): |
|
|
super(BiSeNet, self).__init__() |
|
|
|
|
|
self.conv1 = nn.Sequential( |
|
|
nn.Conv2d(3, 64, 3, stride=2, padding=1), |
|
|
nn.BatchNorm2d(64), |
|
|
nn.ReLU(), |
|
|
nn.Conv2d(64, 64, 3, stride=2, padding=1), |
|
|
nn.BatchNorm2d(64), |
|
|
nn.ReLU() |
|
|
) |
|
|
|
|
|
self.final = nn.Conv2d(64, n_classes, 1) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.conv1(x) |
|
|
x = self.final(x) |
|
|
return x |
|
|
|
|
|
|
|
|
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 AdvancedFaceParsing: |
|
|
def __init__(self): |
|
|
self.model = None |
|
|
self.device = device |
|
|
self.model_type = "unknown" |
|
|
self.load_best_model() |
|
|
|
|
|
def load_best_model(self): |
|
|
"""سعی میکند بهترین مدل موجود را لود کند""" |
|
|
models_to_try = [ |
|
|
|
|
|
{ |
|
|
"name": "BiSeNet-Face-Parsing", |
|
|
"repo_id": "yangyuke001/bisenet-face-parsing", |
|
|
"filename": "model.pth", |
|
|
"constructor": self.create_bisenet |
|
|
}, |
|
|
{ |
|
|
"name": "CelebAMask-HQ-Improved", |
|
|
"repo_id": "public-data/CelebAMask-HQ-Face-Parsing", |
|
|
"filename": "models/model.pth", |
|
|
"constructor": self.create_celebamask_unet |
|
|
}, |
|
|
|
|
|
{ |
|
|
"name": "CelebAMask-HQ-Original", |
|
|
"repo_id": "public-data/CelebAMask-HQ-Face-Parsing", |
|
|
"filename": "model.pth", |
|
|
"constructor": self.create_celebamask_unet |
|
|
} |
|
|
] |
|
|
|
|
|
for model_info in models_to_try: |
|
|
try: |
|
|
print(f"🔄 Trying {model_info['name']}...") |
|
|
model_path = hf_hub_download( |
|
|
repo_id=model_info["repo_id"], |
|
|
filename=model_info["filename"], |
|
|
cache_dir="/home/user/app/hf_cache" |
|
|
) |
|
|
|
|
|
self.model = model_info["constructor"]() |
|
|
state_dict = torch.load(model_path, map_location="cpu") |
|
|
|
|
|
|
|
|
new_state_dict = {} |
|
|
for k, v in state_dict.items(): |
|
|
if k.startswith('module.'): |
|
|
k = k[7:] |
|
|
new_state_dict[k] = v |
|
|
|
|
|
self.model.load_state_dict(new_state_dict, strict=False) |
|
|
self.model.eval() |
|
|
self.model.to(self.device) |
|
|
self.model_type = model_info["name"] |
|
|
|
|
|
print(f"✅ Successfully loaded {model_info['name']}") |
|
|
return |
|
|
|
|
|
except Exception as e: |
|
|
print(f"❌ Failed to load {model_info['name']}: {e}") |
|
|
continue |
|
|
|
|
|
print("⚠️ Could not load any model, using simple fallback") |
|
|
self.model = self.create_simple_model() |
|
|
self.model_type = "Simple-Fallback" |
|
|
|
|
|
def create_bisenet(self): |
|
|
"""ایجاد مدل BiSeNet""" |
|
|
return BiSeNet(n_classes=19) |
|
|
|
|
|
def create_celebamask_unet(self): |
|
|
"""ایجاد مدل CelebAMask-HQ U-Net""" |
|
|
if HAS_CELEBAMASK: |
|
|
return celebamask_unet( |
|
|
feature_scale=4, |
|
|
n_classes=19, |
|
|
is_deconv=True, |
|
|
in_channels=3, |
|
|
is_batchnorm=True |
|
|
) |
|
|
else: |
|
|
return self.create_simple_model() |
|
|
|
|
|
def create_simple_model(self): |
|
|
"""مدل ساده fallback""" |
|
|
return nn.Sequential( |
|
|
nn.Conv2d(3, 64, 3, padding=1), |
|
|
nn.ReLU(), |
|
|
nn.Conv2d(64, 19, 1) |
|
|
) |
|
|
|
|
|
def predict(self, image): |
|
|
"""پردازش تصویر""" |
|
|
if self.model is None: |
|
|
raise ValueError("Model not loaded") |
|
|
|
|
|
|
|
|
if isinstance(image, str): |
|
|
image = PIL.Image.open(image).convert('RGB') |
|
|
elif isinstance(image, np.ndarray): |
|
|
image = PIL.Image.fromarray(image) |
|
|
|
|
|
original_image = image.copy() |
|
|
|
|
|
|
|
|
transform = T.Compose([ |
|
|
T.Resize((512, 512)), |
|
|
T.ToTensor(), |
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
|
]) |
|
|
|
|
|
data = transform(image).unsqueeze(0).to(self.device) |
|
|
|
|
|
|
|
|
with torch.no_grad(): |
|
|
out = self.model(data) |
|
|
|
|
|
|
|
|
if hasattr(self, 'generate_label') and HAS_CELEBAMASK: |
|
|
mask = generate_label(out, 512)[0].cpu().numpy() |
|
|
else: |
|
|
|
|
|
mask = torch.argmax(out, dim=1)[0].cpu().numpy() |
|
|
|
|
|
colored_mask = self.colorize_mask(mask) |
|
|
|
|
|
|
|
|
resized_image = np.asarray(original_image.resize((512, 512))) |
|
|
blended = cv2.addWeighted(resized_image, 0.7, colored_mask, 0.3, 0) |
|
|
|
|
|
return colored_mask, blended, self.model_type |
|
|
|
|
|
def colorize_mask(self, mask): |
|
|
"""رنگآمیزی ماسک""" |
|
|
palette = [ |
|
|
[0, 0, 0], [255, 200, 200], [0, 255, 0], [0, 200, 0], |
|
|
[255, 0, 0], [200, 0, 0], [255, 255, 0], [0, 0, 255], |
|
|
[0, 0, 200], [128, 0, 128], [255, 165, 0], [255, 0, 255], |
|
|
[200, 0, 200], [165, 42, 42], [0, 255, 255], [0, 200, 200], |
|
|
[128, 128, 128], [255, 255, 255], [255, 215, 0] |
|
|
] |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
face_parser = AdvancedFaceParsing() |
|
|
print(f"🎯 Loaded model: {face_parser.model_type}") |
|
|
|
|
|
def process_image(input_image): |
|
|
if input_image is None: |
|
|
return None, None, "لطفاً یک تصویر آپلود کنید" |
|
|
|
|
|
try: |
|
|
mask, blended, model_type = face_parser.predict(input_image) |
|
|
|
|
|
info_text = f""" |
|
|
✅ پردازش انجام شد با {model_type}! |
|
|
- مدل: {model_type} |
|
|
- کلاسهای تشخیص: {len(CELEBA_CLASSES)} |
|
|
- دستگاه: {device} |
|
|
""" |
|
|
|
|
|
return blended, mask, info_text |
|
|
|
|
|
except Exception as e: |
|
|
return None, None, f"❌ خطا: {str(e)}" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
) |
