Upload 4 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/Eda.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -21,7 +21,7 @@ class FaceCrop:
     def __init__(self):
         self.device = select_device()
         self.half = self.device.type != 'cpu'
-        self.results = {}
+        self.results = []
     
     def load_dataset(self, source):
         self.source = source
@@ -75,20 +75,16 @@ class FaceCrop:
                     det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
 
                     # Write results
-                    ind = 0
                     for *xyxy, conf, cls in det:
-                        if conf > 0.6:  # Write to file
-                            
+                        if conf > 0.6:  # Write to file                         
                             x, y, w, h = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
-                            self.results[ind] = crop(self.source, (x, y), mode=self.mode, size=self.target_size, box=(w, h), face_ratio=self.face_ratio, shreshold=self.threshold)
-
-                            ind += 1
+                            self.results.append(crop(self.source, (x, y), mode=self.mode, size=self.target_size, box=(w, h), face_ratio=self.face_ratio, shreshold=self.threshold))
                             
-def run(img, mode, width, height):
+def run(img, mode, width, height, face_ratio, threshold):
     face_crop_pipeline.load_dataset(img)
-    face_crop_pipeline.set_crop_config(mode=mode, target_size=(width,height))
+    face_crop_pipeline.set_crop_config(mode=mode, target_size=(width,height), face_ratio=face_ratio, threshold=threshold)
     face_crop_pipeline.process()
-    return face_crop_pipeline.results[0]
+    return face_crop_pipeline.results
 
 if __name__ == '__main__':
     model_path = huggingface_hub.hf_hub_download("Carzit/yolo5x_anime", "yolov5x_anime.pt")
@@ -97,17 +93,29 @@ if __name__ == '__main__':
     
     app = gr.Blocks()
     with app:
-        gr.Markdown("# Anime Face Crop")
+        gr.Markdown("# Face Crop Anime")
         with gr.Row():
-            input_img = gr.Image(label="input image", image_mode="RGB", type='filepath')
-            output_img = gr.Image(label="result", image_mode="RGB")
+            input_img = gr.Image(label="Input Image", image_mode="RGB", type='filepath')
+            output_img = gr.Gallery(label="Cropped Image")
+            
         with gr.Row():
-            crop_mode = gr.Dropdown([0, 1, 2, 3], label="Crop Mode", info="0:Auto; 1:No Scale; 2:Full Screen; 3:Fixed Face Ratio")
-            tgt_width = gr.Slider(10, 2048, value=512, label="Width")
-            tgt_height = gr.Slider(10, 2048, value=512, label="Height")
+            crop_mode = gr.Dropdown(['Auto', 'No Scale', 'Full Screen', 'Fixed Face Propotion'], label="Crop Mode", value='Auto', type='index')
+            tgt_width = gr.Slider(32, 2048, value=512, label="Width")
+            tgt_height = gr.Slider(32, 2048, value=512, label="Height")
 
+        with gr.Row():
+            face_ratio = gr.Slider(1, 5, step=0.1, value=2, label="Face Ratio", info="Necessary if choosing \'Auto\' or 'Fixed Face Propotion' Mode")
+            threshold = gr.Slider(1, 5, step=0.1, value=1.5, label="Threshold", info="Necessary if choosing \'Auto\' Mode")
             run_btn = gr.Button(variant="primary")
-        run_btn.click(run, [input_img, crop_mode, tgt_width, tgt_height], [output_img])
+        
+        with gr.Row():
+            examples_data = [["examples/Eda.png"],["examples/Fairies.png"]]
+            examples = gr.Dataset(components=[input_img], samples=examples_data)
+        
+        examples.click(lambda x: x[0], [examples], [input_img])
+        run_btn.click(run, [input_img, crop_mode, tgt_width, tgt_height, face_ratio, threshold], [output_img])
+        
+
     app.launch()
 
 

crop.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from PIL import Image
 import numpy as np
 
-def crop(image_path, point, mode=0, size=(512, 512), box=None, face_ratio=3, shreshold=1.5):
+def crop(image_path, point, mode=0, size=(512, 512), box=None, face_ratio=2, shreshold=1.5):
     img = Image.open(image_path)
     img_width, img_height = img.size
     tgt_width, tgt_height = size