minor refactor (#12)

* minor-refactor * minor refactor
fcakyon · Jan 13, 2023 · 909c6c9 · 909c6c9
1 parent ca76274
commit 909c6c9
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 13 deletions.
diff --git a/ultralyticsplus/__init__.py b/ultralyticsplus/__init__.py
@@ -1,4 +1,4 @@
 from .hf_utils import push_to_hfhub, download_from_hub
 from .ultralytics_utils import YOLO, render_predictions
 
-__version__ = "0.0.6"
+__version__ = "0.0.7"
diff --git a/ultralyticsplus/hf_utils.py b/ultralyticsplus/hf_utils.py
@@ -4,7 +4,7 @@
 import pandas as pd
 from PIL import Image
 
-from ultralyticsplus.file_utils import add_text_to_image
+from ultralyticsplus.other_utils import add_text_to_image
 
 LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
 logging.basicConfig(level=LOGLEVEL)

diff --git a/ultralyticsplus/file_utils.py → ultralyticsplus/other_utils.py b/ultralyticsplus/file_utils.py → ultralyticsplus/other_utils.py
@@ -1,13 +1,14 @@
-from PIL import ImageDraw, ImageFont, ImageEnhance
-
-
 def add_text_to_image(
     text,
     pil_image,
     brightness: float = 0.75,
     text_font: int = 50,
     crop_margin: int = None,
 ):
+    from PIL import ImageDraw, ImageFont, ImageEnhance
+
+    RESIZE_WIDTH_TO = 1280
+    MAX_TEXT_WIDTH = 860
 
     # Create an ImageEnhance object
     enhancer = ImageEnhance.Brightness(pil_image)
@@ -16,10 +17,10 @@ def add_text_to_image(
     pil_image = enhancer.enhance(brightness)
 
     # Resize the image
-    pil_image = pil_image.resize((1280, 1280))
+    pil_image = pil_image.resize((RESIZE_WIDTH_TO, RESIZE_WIDTH_TO))
 
+    # crop image
     if crop_margin is not None:
-        # crop image
         width, height = pil_image.size
         mid_height = int(height / 2)
         min_height = max(0, mid_height - crop_margin)
@@ -38,13 +39,14 @@ def add_text_to_image(
     # Get the width and height of the text
     text_width = bbox[2] - bbox[0]
     text_height = bbox[3] - bbox[1]
-    
+
     # update font size
-    text_font = int(860 / text_width * text_font)
-    font = ImageFont.truetype("OpenSans-Bold.ttf", text_font)
-    bbox = draw.textbbox((0, 0), text, font=font)
-    text_width = bbox[2] - bbox[0]
-    text_height = bbox[3] - bbox[1]
+    if text_width > MAX_TEXT_WIDTH:
+        text_font = int(MAX_TEXT_WIDTH / text_width * text_font)
+        font = ImageFont.truetype("OpenSans-Bold.ttf", text_font)
+        bbox = draw.textbbox((0, 0), text, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
 
     # Define the coordinates of the smaller rounded rectangle
     box_margin = 20