svjack commited on
Commit
6859b03
1 Parent(s): 1459d4f

Update pred_color.py

Browse files
Files changed (1) hide show
  1. 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
- #assert width == height
74
- #assert width == size
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