File size: 3,396 Bytes
f4d52c1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import commands.exec_path
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import os
import random
from pathlib import Path
def autoCropImages(image,scale_factor):
# Load a model
model = YOLO("./cv_files/AniFaceDet.pt")
test_images = []
test_images.append(image)
# Create a directory for saving cropped images
relative_dir = './Images/cropped'
cropped_dir = os.path.abspath(relative_dir)
if not os.path.exists(cropped_dir):
os.makedirs(cropped_dir)
imagesToReturn = []
# Load the test image
for img in test_images:
if img.split(".")[-1] not in ["jpg", "jpeg", "png"]:
continue
image_path = Path(image) / img
image = Image.open(image_path)
# Get the size of the image
image_width, image_height = image.size
# Calculate the scaling factor based on the image size for font size
scaling_factor = max(image_width, image_height) / 200
# Calculate the final font size by scaling the base font size
base_font_size = 10
font_size = int(base_font_size * scaling_factor)
font = ImageFont.load_default()
# Predict the bounding boxes #Defaults conf=0.25, iou=0.7
pred = model.predict(image, conf=0.65, iou=0.7)
# Extract the bounding box coordinates, class labels, and confidence scores
boxes = pred[0].boxes.xyxy.tolist()
classes = pred[0].boxes.cls.tolist()
scores = pred[0].boxes.conf.tolist()
# Choose a scale factor for the cropped image
scale_factor = scale_factor
# Loop over all detected faces and draw bounding boxes, crop, and save
for i in range(len(boxes)):
box = boxes[i]
score = scores[i]
x1, y1, x2, y2 = box
# Calculate the width and height of the bounding box and apply the scale factor
box_width = x2 - x1
box_height = y2 - y1
scaled_width = int(box_width * scale_factor)
scaled_height = int(box_height * scale_factor)
# Calculate the top-left corner coordinates of the cropped region
cropped_x1 = max(0, int(x1 - (scaled_width - box_width) / 2))
cropped_y1 = max(0, int(y1 - (scaled_height - box_height) / 2))
# Calculate the bottom-right corner coordinates of the cropped region
cropped_x2 = min(int(x2 + (scaled_width - box_width) / 2), image_width)
cropped_y2 = min(int(y2 + (scaled_height - box_height) / 2), image_height)
# Crop the image based on the detected face
cropped_image = image.crop((cropped_x1, cropped_y1, cropped_x2, cropped_y2))
# Save the cropped image with the original filename and an index
cropped_image_name = '{}_cropped_{}_scale_{}.jpg'.format(os.path.splitext(os.path.split(img)[-1])[0], i, scale_factor)
cropped_image_path = os.path.join(cropped_dir, cropped_image_name)
cropped_image.save(fp=cropped_image_path)
# Appending Cropped images in an array to display in gradio for end-user
imagesToReturn.append(cropped_image_path)
print('Cropped image saved:', cropped_image_path)
# Draw bounding boxes on the original image
draw = ImageDraw.Draw(image)
return imagesToReturn
|