Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,984 Bytes
6613c62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
from transformers import pipeline
import numpy as np
import cv2
from PIL import Image, ImageDraw
# Try to import insightface for face detection, fallback to OpenCV if not available
try:
from insightface.app import FaceAnalysis
USE_INSIGHTFACE = True
except ImportError:
USE_INSIGHTFACE = False
print("⚠️ InsightFace not available, using OpenCV for face detection")
# Initialize face detection
if USE_INSIGHTFACE:
try:
app = FaceAnalysis(providers=['CPUExecutionProvider']) # CPU only for HF Spaces
app.prepare(ctx_id=-1, det_size=(640, 640)) # -1 for CPU
except Exception as e:
print(f"⚠️ InsightFace initialization failed: {e}, falling back to OpenCV")
USE_INSIGHTFACE = False
# Initialize segmentation pipeline
segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes", device=-1) # -1 for CPU
def remove_face_opencv(img, mask):
"""Fallback face detection using OpenCV"""
img_arr = np.asarray(img.convert('RGB'))
gray = cv2.cvtColor(img_arr, cv2.COLOR_RGB2GRAY)
# Load OpenCV's pre-trained face detector
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
# Detect faces
faces = face_cascade.detectMultiScale(gray, 1.1, 4)
if len(faces) > 0:
# Get the first (largest) face
x, y, w, h = faces[0]
# Expand face region
x = max(0, int(x - w * 0.5))
y = max(0, int(y - h * 0.5))
w = int(w * 2.0)
h = int(h * 1.7)
# Draw black rectangle on mask
img1 = ImageDraw.Draw(mask)
img1.rectangle([(x, y), (x + w, y + h)], fill=0)
return mask
def remove_face(img, mask):
"""Remove face from mask using InsightFace or OpenCV fallback"""
if not USE_INSIGHTFACE:
return remove_face_opencv(img, mask)
try:
# Convert image to numpy array
img_arr = np.asarray(img)
# Run face detection
faces = app.get(img_arr)
if len(faces) == 0:
return mask
# Get the first face
bbox = faces[0]['bbox']
# Width and height of face
w = bbox[2] - bbox[0]
h = bbox[3] - bbox[1]
# Make face locations bigger
bbox[0] = bbox[0] - (w * 0.5) # x left
bbox[2] = bbox[2] + (w * 0.5) # x right
bbox[1] = bbox[1] - (h * 0.5) # y top
bbox[3] = bbox[3] + (h * 0.2) # y bottom
# Convert to [(x_left, y_top), (x_right, y_bottom)]
face_locations = [(bbox[0], bbox[1]), (bbox[2], bbox[3])]
# Draw black rect onto mask
img1 = ImageDraw.Draw(mask)
img1.rectangle(face_locations, fill=0)
return mask
except Exception as e:
print(f"⚠️ InsightFace face removal failed: {e}, using OpenCV")
return remove_face_opencv(img, mask)
def segment_body(original_img, face=True):
"""
Segment body from image
Args:
original_img: PIL Image
face: If True, includes face in mask. If False, excludes face.
Returns:
tuple: (segmented_image, mask_image)
"""
# Make a copy
img = original_img.copy()
# Segment image
segments = segmenter(img)
# Create list of masks
segment_include = ["Hat", "Hair", "Sunglasses", "Upper-clothes", "Skirt", "Pants",
"Dress", "Belt", "Left-shoe", "Right-shoe", "Face", "Left-leg",
"Right-leg", "Left-arm", "Right-arm", "Bag", "Scarf"]
mask_list = []
for s in segments:
if s['label'] in segment_include:
mask_list.append(s['mask'])
if len(mask_list) == 0:
# If no segments found, return full mask
print("⚠️ No body segments detected, using full mask")
final_mask = Image.new('L', img.size, 255)
else:
# Paste all masks on top of each other
final_mask = np.array(mask_list[0])
for mask in mask_list[1:]:
current_mask = np.array(mask)
final_mask = final_mask + current_mask
# Convert final mask from np array to PIL image
final_mask = Image.fromarray(final_mask)
# Remove face
if face == False:
final_mask = remove_face(img.convert('RGB'), final_mask)
# Apply mask to original image
img_copy = img.copy()
img_copy.putalpha(final_mask)
return img_copy, final_mask
def segment_torso(original_img):
"""
Segment only torso/upper body from image
Args:
original_img: PIL Image
Returns:
tuple: (segmented_image, mask_image)
"""
# Make a copy
img = original_img.copy()
# Segment image
segments = segmenter(img)
# Create list of masks (torso only)
segment_include = ["Upper-clothes", "Dress", "Belt", "Face", "Left-arm", "Right-arm"]
mask_list = []
for s in segments:
if s['label'] in segment_include:
mask_list.append(s['mask'])
if len(mask_list) == 0:
# If no segments found, return full mask
print("⚠️ No torso segments detected, using full mask")
final_mask = Image.new('L', img.size, 255)
else:
# Paste all masks on top of each other
final_mask = np.array(mask_list[0])
for mask in mask_list[1:]:
current_mask = np.array(mask)
final_mask = final_mask + current_mask
# Convert final mask from np array to PIL image
final_mask = Image.fromarray(final_mask)
# Remove face
final_mask = remove_face(img.convert('RGB'), final_mask)
# Apply mask to original image
img_copy = img.copy()
img_copy.putalpha(final_mask)
return img_copy, final_mask
|