Update app.py

app.py

@@ -33,39 +33,27 @@ def generate_caption(image):
     meme_text = response.choices[0].message.content
     print(meme_text)
 
-    pil_image = Image.fromarray((image.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
+    image = Image.fromarray((image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)).convert('RGB')
 
-    # Create a copy of the input PIL image
-    meme_image = pil_image.copy()
+    meme_image = image.copy()
+    draw = ImageDraw.Draw(meme_image)
 
-    # Convert the copied image to RGBA mode to support transparency
-    meme_image = meme_image.convert('RGBA')
+    font_size = 1
+    text_width = 0
+    max_width = int(224 * 0.9)
+    while text_width < max_width:
+        font_size += 1
+        font = ImageFont.truetype("impact.ttf", font_size)
+        text_width, _ = draw.textsize(meme_text, font=font)
 
-    # Add the text to the image
-    draw = ImageDraw.Draw(meme_image)    
-    
-    # Calculate font size based on text length
-    text_length = len(meme_text)
-    if text_length <= 15:
-        font_size = 24
-    elif text_length <= 30:
-        font_size = 18
-    else:
-        font_size = 14
-
-    font = ImageFont.load_default()
     text_width, text_height = draw.textsize(meme_text, font=font)
-    
-    # Calculate the position to place the text at the top and center horizontally
-    x = (meme_image.width - text_width) // 2
-    y = 10  # Adjust this value to add more or less padding from the top
+    text_x = (224 - text_width) // 2
+    text_y = (224 - text_height) // 2
 
-    draw.text((x, y), meme_text, fill=(255, 255, 255), font=font)
+    draw.text((text_x, 0), meme_text, font=font, fill=(255, 255, 255))
 
-    # Convert the meme_image back to RGB mode
     meme_image = meme_image.convert('RGB')
 
-
     if torch.cuda.is_available():
         torch.cuda.empty_cache()