Update app.py

app.py

@@ -248,8 +248,15 @@ def predict_text(image: Image.Image):
 
 
 # New helper function: generate label images grid
-def generate_labels(font_file=None, num_labels: int = 20):
+CHARS = string.ascii_letters + string.digits + string.punctuation
+
+FONT_SIZE = 32
+PADDING = 8
+LABEL_DIR = "./labels"
+
+def generate_labels(font_file=None, num_labels: int = 25):
     global font_path
+    os.makedirs(LABEL_DIR, exist_ok=True)
 
     try:
         if font_file:
@@ -257,44 +264,46 @@ def generate_labels(font_file=None, num_labels: int = 20):
             with open(font_file.name, "rb") as uploaded:
                 with open(font_path, "wb") as f:
                     f.write(uploaded.read())
-
-        if font_path is None or not os.path.exists(font_path):
+        if not font_file or not os.path.exists(font_path):
             font = ImageFont.load_default()
         else:
-            font = ImageFont.truetype(font_path, 25)
+            font = ImageFont.truetype(font_path, FONT_SIZE)
 
-        labels = ["".join(np.random.choice(list(CHARS), np.random.randint(4, 7))) for _ in range(num_labels)]
+        images = []
+        labels = [
+            "".join(np.random.choice(list(CHARS), np.random.randint(4, 7)))
+            for _ in range(num_labels)
+        ]
 
-        cols = min(5, num_labels)
-        rows = math.ceil(num_labels / cols)
-        cell_w, cell_h = IMAGE_WIDTH, IMAGE_HEIGHT
+        for label in labels:
+            # Measure actual text bounding box
+            dummy_img = Image.new("L", (1, 1))
+            draw = ImageDraw.Draw(dummy_img)
+            bbox = draw.textbbox((0, 0), label, font=font)
 
-        grid_img = Image.new("L", (cols * cell_w, rows * cell_h), color=255)
-        draw = ImageDraw.Draw(grid_img)
+            text_width = bbox[2] - bbox[0]
+            text_height = bbox[3] - bbox[1]
 
-        spacing = 0  # <-- spacing between characters
+            # Compute final image size with padding
+            image_width = text_width + 2 * PADDING
+            image_height = text_height + 2 * PADDING
 
-        for idx, label in enumerate(labels):
-            x = (idx % cols) * cell_w
-            y = (idx // cols) * cell_h + 10
+            img = Image.new("L", (image_width, image_height), color=255)
+            draw = ImageDraw.Draw(img)
+            draw.text((PADDING, PADDING), label, font=font, fill=0)
 
-            # Draw each character with spacing
-            char_x = x + 10  # small left margin
-            char_y = y + (cell_h - font.size) // 2
+            filename = f"{LABEL_DIR}/{label}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.png"
+            img.save(filename)
+            images.append(img)
 
-            for char in label:
-                draw.text((char_x, char_y), char, font=font, fill=0)
-                char_w = font.getbbox(char)[2] - font.getbbox(char)[0]
-                char_x += char_w + spacing  # move right with spacing
-
-        return grid_img
+        return images
 
     except Exception as e:
         print("Error in generate_labels:", e)
         error_img = Image.new("RGB", (512, 128), color=(255, 255, 255))
         draw = ImageDraw.Draw(error_img)
         draw.text((10, 50), f"Error: {str(e)}", fill=(255, 0, 0))
-        return error_img
+        return [error_img]
 
 
 
@@ -334,7 +343,12 @@ with gr.Blocks() as demo:
         gen_button = gr.Button("Generate Label Grid")
         label_image = gr.Image(label="Generated Labels", type="pil")
 
-        gen_button.click(fn=generate_labels, inputs=[font_file_labels, num_labels], outputs=label_image)
+        gen_button.click(
+            fn=generate_labels,
+            inputs=[font_file_labels, num_labels],
+            outputs=gr.Gallery(label="Label Images", show_label=True)
+        )
+