add thumbnail generation (#8)

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.4"
+__version__ = "0.0.5"

ultralyticsplus/file_utils.py

@@ -0,0 +1,71 @@
+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,
+):
+
+    # Create an ImageEnhance object
+    enhancer = ImageEnhance.Brightness(pil_image)
+
+    # Darken the image by 25%
+    pil_image = enhancer.enhance(brightness)
+
+    # Resize the image
+    pil_image = pil_image.resize((1280, 1280))
+
+    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)
+        max_height = min(height, mid_height + crop_margin)
+        pil_image = pil_image.crop((0, min_height, width, max_height))
+
+    # Create an ImageDraw object
+    draw = ImageDraw.Draw(pil_image)
+
+    # Define the font and text to be written
+    font = ImageFont.truetype("OpenSans-Bold.ttf", text_font)
+
+    # Get the bounding box of the text
+    bbox = draw.textbbox((0, 0), text, font=font)
+
+    # Get the width and height of the text
+    text_width = bbox[2] - bbox[0]
+    text_height = bbox[3] - bbox[1]
+
+    # Define the coordinates of the smaller rounded rectangle
+    box_margin = 20
+    x1, y1 = (pil_image.width - text_width) / 2 - box_margin, (
+        pil_image.height - text_height
+    ) / 2 - box_margin / 3
+    x2, y2 = (pil_image.width + text_width) / 2 + box_margin, (
+        pil_image.height + text_height * 2
+    ) / 2 + box_margin / 3
+
+    # Define the radius of the rounded corners
+    radius = 15
+
+    # Draw the rounded rectangle
+    draw.rounded_rectangle(
+        [(x1, y1), (x2, y2)],
+        radius=radius,
+        fill=(255, 255, 255),
+        outline=None,
+        width=5,
+    )
+
+    # Draw the text on the image
+    draw.text(
+        ((pil_image.width - text_width) / 2, (pil_image.height - text_height) / 2),
+        text,
+        font=font,
+        fill=(0, 0, 0),
+    )
+
+    return pil_image

ultralyticsplus/hf_utils.py

@@ -2,7 +2,7 @@
 import os
 from pathlib import Path
 
-import pandas as pd
+from ultralyticsplus.file_utils import add_text_to_image
 
 LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
 logging.basicConfig(level=LOGLEVEL)
@@ -98,6 +98,45 @@ def generate_model_usage_markdown(
 """
 
 
+def generate_thumbnail(image_path, repo_id, task="object-detection"):
+    """
+    Generate thumbnail for the model card
+
+    USERNAME/yolov8n-garbage > YOLOv8 Garbage Detection
+    """
+    thumbnail_text = repo_id.split("/")[-1]
+    texts = thumbnail_text.split("-")
+    for text in texts:
+        if "yolo" not in text.lower():
+            text = text.title()
+        text.replace("yolo", "YOLO")
+
+    thumbnail_text = " ".join(texts)
+
+    if task == "object-detection":
+        thumbnail_text += " Detection"
+    elif task == "image-classification":
+        thumbnail_text += " Classification"
+    elif task == "instance-segmentation":
+        thumbnail_text += " Segmentation"
+    else:
+        raise ValueError(f"Task {task} is not supported.")
+
+    image = add_text_to_image(
+        text=thumbnail_text,
+        image_path=image_path,
+        brightness=0.60,
+        text_font=65,
+        crop_margin=None,
+    )
+
+    folder_path = Path(image_path).parent
+    thumbnail_path = folder_path / "thumbnail.jpg"
+    image.save(str(thumbnail_path), quality=100)
+
+    return thumbnail_path
+
+
 def push_model_card_to_hfhub(
     repo_id,
     exp_folder,
@@ -118,12 +157,13 @@ def push_model_card_to_hfhub(
         exist_ok=True,
     )
 
-    # upload sample visual to the repo
+    # upload thumbnail to the repo
     sample_visual_path = Path(exp_folder) / "val_batch0_labels.jpg"
+    thumbnail_path = generate_thumbnail(sample_visual_path, repo_id=repo_id, task=task)
     upload_file(
         repo_id=repo_id,
-        path_or_fileobj=str(sample_visual_path),
-        path_in_repo="sample_visuals.jpg",
+        path_or_fileobj=thumbnail_path,
+        path_in_repo="thumbnail.jpg",
         commit_message="upload sample visuals",
         token=hf_token,
         repo_type="model",