Update pred_color.py

pred_color.py

@@ -5,6 +5,7 @@
 !pip install -r SPIGA/requirements.txt
 !pip install datasets
 !pip install retinaface-py>=0.0.2
+!pip install bounding-box
 !huggingface-cli login
 '''
 
@@ -70,8 +71,8 @@ def conditioning_from_landmarks(landmarks, size=512):
 
     fig.canvas.draw()
     buffer, (width, height) = fig.canvas.print_to_buffer()
-    #assert width == height
-    #assert width == size
+    assert width == height
+    assert width == size
 
     buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
     buffer = buffer[:, :, 0:3]
@@ -113,7 +114,70 @@ def single_pred_features(image):
         landmarks = face_features["landmarks"][0]
         face_features["spiga"] = landmarks
         face_features['spiga_seg'] = conditioning_from_landmarks(landmarks)
-        return face_features
+        return features ,face_features
+
+def produce_center_crop_image(features ,face_features, draw_rect = False):
+    left, top, right, bottom, _ = features["bbox"][0]
+    color = "red"
+    label = ""
+    from bounding_box import bounding_box as bb
+    img = np.asarray(face_features["spiga_seg"])
+    step = 20
+    if draw_rect:
+        bb.add(img, left - step, top - step, right + step, bottom + step, label, color)
+    crop_img = Image.fromarray(img[ int(top - step):int(bottom + step) ,int(left - step):int(right + step), :])
+    crop_img = crop_img.resize((256, 256))
+    req = Image.fromarray(
+    np.concatenate(
+        [np.full([512, 128, 3], fill_value=0),
+    np.concatenate([np.full([128, 256, 3], fill_value=0) ,np.asarray(crop_img),
+                    np.full([128, 256, 3], fill_value=0)], axis = 0),
+         np.full([512, 128, 3], fill_value=0)
+        ], axis = 1
+    ).astype(np.uint8))
+    return req
+
+'''
+from pred_color import *
+img = "babyxiang_ai.png"
+img = "Protector_Cromwell_style.png"
+features ,face_features = single_pred_features(img)
+
+fix_img = produce_center_crop_image(features ,face_features, draw_rect = False)
+fix_img
+
+fix_r_img = produce_center_crop_image(features ,face_features, draw_rect = True)
+fix_r_img
+
+
+from pred_color import *
+img = "babyxiang_ai.png"
+img = "Protector_Cromwell_style.png"
+features ,face_features = single_pred_features(img)
+
+left, top, right, bottom, _ = features["bbox"][0]
+color = "red"
+label = ""
+
+from bounding_box import bounding_box as bb
+img = np.asarray(face_features["spiga_seg"])
+step = 20
+bb.add(img, left - step, top - step, right + step, bottom + step, label, color)
+Image.fromarray(img)
+
+crop_img = Image.fromarray(img[ int(top - step):int(bottom + step) ,int(left - step):int(right + step), :])
+crop_img = crop_img.resize((256, 256))
+crop_img
+
+Image.fromarray(
+np.concatenate(
+    [np.full([512, 128, 3], fill_value=0),
+np.concatenate([np.full([128, 256, 3], fill_value=0) ,np.asarray(crop_img),
+                np.full([128, 256, 3], fill_value=0)], axis = 0),
+     np.full([512, 128, 3], fill_value=0)
+    ], axis = 1
+).astype(np.uint8))
+'''
 
 if __name__ == "__main__":
     from datasets import load_dataset, Dataset