Upload character_splitter/app.py with huggingface_hub

character_splitter/app.py

@@ -0,0 +1,62 @@
+import os
+
+import gradio as gr
+from imgutils.detect import detect_person, detect_halfbody, detect_heads, detection_visualize
+
+
+def _split_image(image, head_scale: float):
+    retval = []
+    all_detects = []
+    for i, (px, _, score) in enumerate(detect_person(image), start=1):
+        person_image = image.crop(px)
+        person_label = f'Person #{i} ({score * 100.0:.1f}%)'
+        retval.append((person_image, person_label))
+        all_detects.append((px, 'person', score))
+        px0, py0, _, _ = px
+
+        half_detects = detect_halfbody(person_image)
+        if half_detects:
+            halfbody_image = person_image.crop(half_detects[0][0])
+            halfbody_label = f'Person #{i} - Half Body'
+            retval.append((halfbody_image, halfbody_label))
+            bx0, by0, bx1, by1 = half_detects[0][0]
+            all_detects.append(((bx0 + px0, by0 + py0, bx1 + px0, by1 + py0), 'halfbody', half_detects[0][2]))
+
+        head_detects = detect_heads(person_image)
+        if head_detects:
+            (hx0, hy0, hx1, hy1), _, head_score = head_detects[0]
+            cx, cy = (hx0 + hx1) / 2, (hy0 + hy1) / 2
+            width, height = hx1 - hx0, hy1 - hy0
+            width = height = max(width, height) * head_scale
+            x0, y0 = int(max(cx - width / 2, 0)), int(max(cy - height / 2, 0))
+            x1, y1 = int(min(cx + width / 2, person_image.width)), int(min(cy + height / 2, person_image.height))
+            head_image = person_image.crop((x0, y0, x1, y1))
+            head_label = f'Person #{i} - Head'
+            retval.append((head_image, head_label))
+            all_detects.append(((x0 + px0, y0 + py0, x1 + px0, y1 + py0), 'head', head_score))
+
+    return detection_visualize(image, all_detects), retval
+
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                gr_input = gr.Image(type='pil', label='Original Image')
+                gr_head_scale = gr.Slider(0.8, 2.5, 1.5, label='Head Scale')
+                gr_button = gr.Button(value='Crop', variant='primary')
+
+            with gr.Column():
+                with gr.Tabs():
+                    with gr.Tab('Detected'):
+                        gr_detected = gr.Image(type='pil', label='Detection')
+                    with gr.Tab('Cropped'):
+                        gr_gallery = gr.Gallery(label='Cropped Images')
+
+            gr_button.click(
+                _split_image,
+                inputs=[gr_input, gr_head_scale],
+                outputs=[gr_detected, gr_gallery],
+            )
+
+    demo.queue(os.cpu_count()).launch()