Spaces:
Running
Running
Update pred_color.py
Browse files- pred_color.py +67 -3
pred_color.py
CHANGED
@@ -5,6 +5,7 @@
|
|
5 |
!pip install -r SPIGA/requirements.txt
|
6 |
!pip install datasets
|
7 |
!pip install retinaface-py>=0.0.2
|
|
|
8 |
!huggingface-cli login
|
9 |
'''
|
10 |
|
@@ -70,8 +71,8 @@ def conditioning_from_landmarks(landmarks, size=512):
|
|
70 |
|
71 |
fig.canvas.draw()
|
72 |
buffer, (width, height) = fig.canvas.print_to_buffer()
|
73 |
-
|
74 |
-
|
75 |
|
76 |
buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
|
77 |
buffer = buffer[:, :, 0:3]
|
@@ -113,7 +114,70 @@ def single_pred_features(image):
|
|
113 |
landmarks = face_features["landmarks"][0]
|
114 |
face_features["spiga"] = landmarks
|
115 |
face_features['spiga_seg'] = conditioning_from_landmarks(landmarks)
|
116 |
-
return face_features
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
117 |
|
118 |
if __name__ == "__main__":
|
119 |
from datasets import load_dataset, Dataset
|
|
|
5 |
!pip install -r SPIGA/requirements.txt
|
6 |
!pip install datasets
|
7 |
!pip install retinaface-py>=0.0.2
|
8 |
+
!pip install bounding-box
|
9 |
!huggingface-cli login
|
10 |
'''
|
11 |
|
|
|
71 |
|
72 |
fig.canvas.draw()
|
73 |
buffer, (width, height) = fig.canvas.print_to_buffer()
|
74 |
+
assert width == height
|
75 |
+
assert width == size
|
76 |
|
77 |
buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
|
78 |
buffer = buffer[:, :, 0:3]
|
|
|
114 |
landmarks = face_features["landmarks"][0]
|
115 |
face_features["spiga"] = landmarks
|
116 |
face_features['spiga_seg'] = conditioning_from_landmarks(landmarks)
|
117 |
+
return features ,face_features
|
118 |
+
|
119 |
+
def produce_center_crop_image(features ,face_features, draw_rect = False):
|
120 |
+
left, top, right, bottom, _ = features["bbox"][0]
|
121 |
+
color = "red"
|
122 |
+
label = ""
|
123 |
+
from bounding_box import bounding_box as bb
|
124 |
+
img = np.asarray(face_features["spiga_seg"])
|
125 |
+
step = 20
|
126 |
+
if draw_rect:
|
127 |
+
bb.add(img, left - step, top - step, right + step, bottom + step, label, color)
|
128 |
+
crop_img = Image.fromarray(img[ int(top - step):int(bottom + step) ,int(left - step):int(right + step), :])
|
129 |
+
crop_img = crop_img.resize((256, 256))
|
130 |
+
req = Image.fromarray(
|
131 |
+
np.concatenate(
|
132 |
+
[np.full([512, 128, 3], fill_value=0),
|
133 |
+
np.concatenate([np.full([128, 256, 3], fill_value=0) ,np.asarray(crop_img),
|
134 |
+
np.full([128, 256, 3], fill_value=0)], axis = 0),
|
135 |
+
np.full([512, 128, 3], fill_value=0)
|
136 |
+
], axis = 1
|
137 |
+
).astype(np.uint8))
|
138 |
+
return req
|
139 |
+
|
140 |
+
'''
|
141 |
+
from pred_color import *
|
142 |
+
img = "babyxiang_ai.png"
|
143 |
+
img = "Protector_Cromwell_style.png"
|
144 |
+
features ,face_features = single_pred_features(img)
|
145 |
+
|
146 |
+
fix_img = produce_center_crop_image(features ,face_features, draw_rect = False)
|
147 |
+
fix_img
|
148 |
+
|
149 |
+
fix_r_img = produce_center_crop_image(features ,face_features, draw_rect = True)
|
150 |
+
fix_r_img
|
151 |
+
|
152 |
+
|
153 |
+
from pred_color import *
|
154 |
+
img = "babyxiang_ai.png"
|
155 |
+
img = "Protector_Cromwell_style.png"
|
156 |
+
features ,face_features = single_pred_features(img)
|
157 |
+
|
158 |
+
left, top, right, bottom, _ = features["bbox"][0]
|
159 |
+
color = "red"
|
160 |
+
label = ""
|
161 |
+
|
162 |
+
from bounding_box import bounding_box as bb
|
163 |
+
img = np.asarray(face_features["spiga_seg"])
|
164 |
+
step = 20
|
165 |
+
bb.add(img, left - step, top - step, right + step, bottom + step, label, color)
|
166 |
+
Image.fromarray(img)
|
167 |
+
|
168 |
+
crop_img = Image.fromarray(img[ int(top - step):int(bottom + step) ,int(left - step):int(right + step), :])
|
169 |
+
crop_img = crop_img.resize((256, 256))
|
170 |
+
crop_img
|
171 |
+
|
172 |
+
Image.fromarray(
|
173 |
+
np.concatenate(
|
174 |
+
[np.full([512, 128, 3], fill_value=0),
|
175 |
+
np.concatenate([np.full([128, 256, 3], fill_value=0) ,np.asarray(crop_img),
|
176 |
+
np.full([128, 256, 3], fill_value=0)], axis = 0),
|
177 |
+
np.full([512, 128, 3], fill_value=0)
|
178 |
+
], axis = 1
|
179 |
+
).astype(np.uint8))
|
180 |
+
'''
|
181 |
|
182 |
if __name__ == "__main__":
|
183 |
from datasets import load_dataset, Dataset
|