Create cropImage.py

app.py

@@ -6,10 +6,11 @@ import gradio as gr
 import PIL.Image
 import zipfile
 from genTag import genTag
+from cropImage import cropImage
 from checkIgnore import is_ignore
 from createTagDom import create_tag_dom
 
-def predict(image: PIL.Image.Image, score_threshold: float):
+def predict(image: PIL.Image.Image, score_threshold: float):    
     result_threshold = genTag(image, score_threshold)
     result_html = ''
     for label, prob in result_threshold.items():
@@ -17,7 +18,8 @@ def predict(image: PIL.Image.Image, score_threshold: float):
     result_html = '<div>' + result_html + '</div>'
     result_filter = {key: value for key, value in result_threshold.items() if not is_ignore(key, 1)}
     result_text = '<div id="m5dd_result">' + ', '.join(result_filter.keys()) + '</div>'
-    return result_html, result_text
+    crop_image = cropImage(image)
+    return result_html, result_text, crop_image
 
 def predict_batch(zip_file, score_threshold: float, progress=gr.Progress()):
     result = ''
@@ -49,6 +51,11 @@ with gr.Blocks(css="style.css", js="script.js") as demo:
                                             value=0.3)
                 run_button = gr.Button('Run')
                 result_text = gr.HTML(value="")
+                with gr.Accordion("Crop Image"):
+                    crop_image = gr.Image(elem_classes='m5dd_image',
+                                          format='jpg',
+                                          show_label=False,
+                                          container=False)
             with gr.Column(scale=2):
                 result_html = gr.HTML(value="")
     with gr.Tab(label='Batch'):
@@ -72,7 +79,7 @@ with gr.Blocks(css="style.css", js="script.js") as demo:
     run_button.click(
         fn=predict,
         inputs=[image, score_threshold],
-        outputs=[result_html, result_text],
+        outputs=[result_html, result_text, crop_image],
         api_name='predict',
     )
     run_button2.click(

cropImage.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import PIL.Image
+
+def cropImage(image: PIL.Image.Image):
+    original_width, original_height = image.size
+    scale = max(original_width, original_height) / min(original_width, original_height)
+
+    target_width = 512
+    target_height = 768
+
+    if scale < 1.1:
+        target_width = 640
+        target_height = 640
+    elif original_width > original_height:
+        target_width = 768
+        target_height = 512
+
+    if original_width / original_height > target_width / target_height:
+        new_width = int(original_height * (target_width / target_height))
+        crop_box = ((original_width - new_width) // 2, 0, (original_width + new_width) // 2, original_height)
+    else:
+        new_height = int(original_width * (target_height / target_width))
+        crop_box = (0, (original_height - new_height) // 2, original_width, (original_height + new_height) // 2)
+    cropped_image = image.crop(crop_box)
+
+    cropped_image = cropped_image.resize((target_width, target_height))
+
+    return cropped_image