|
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): |
|
|
|
model = YOLO("./cv_files/AniFaceDet.pt") |
|
|
|
test_images = [] |
|
test_images.append(image) |
|
|
|
|
|
relative_dir = './Images/cropped' |
|
cropped_dir = os.path.abspath(relative_dir) |
|
if not os.path.exists(cropped_dir): |
|
os.makedirs(cropped_dir) |
|
|
|
imagesToReturn = [] |
|
|
|
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) |
|
|
|
|
|
image_width, image_height = image.size |
|
|
|
|
|
scaling_factor = max(image_width, image_height) / 200 |
|
|
|
base_font_size = 10 |
|
font_size = int(base_font_size * scaling_factor) |
|
font = ImageFont.load_default() |
|
|
|
|
|
pred = model.predict(image, conf=0.65, iou=0.7) |
|
|
|
|
|
boxes = pred[0].boxes.xyxy.tolist() |
|
classes = pred[0].boxes.cls.tolist() |
|
scores = pred[0].boxes.conf.tolist() |
|
|
|
|
|
scale_factor = scale_factor |
|
|
|
|
|
for i in range(len(boxes)): |
|
box = boxes[i] |
|
score = scores[i] |
|
|
|
x1, y1, x2, y2 = box |
|
|
|
|
|
box_width = x2 - x1 |
|
box_height = y2 - y1 |
|
scaled_width = int(box_width * scale_factor) |
|
scaled_height = int(box_height * scale_factor) |
|
|
|
|
|
cropped_x1 = max(0, int(x1 - (scaled_width - box_width) / 2)) |
|
cropped_y1 = max(0, int(y1 - (scaled_height - box_height) / 2)) |
|
|
|
|
|
cropped_x2 = min(int(x2 + (scaled_width - box_width) / 2), image_width) |
|
cropped_y2 = min(int(y2 + (scaled_height - box_height) / 2), image_height) |
|
|
|
|
|
cropped_image = image.crop((cropped_x1, cropped_y1, cropped_x2, cropped_y2)) |
|
|
|
|
|
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) |
|
|
|
|
|
imagesToReturn.append(cropped_image_path) |
|
|
|
print('Cropped image saved:', cropped_image_path) |
|
|
|
|
|
draw = ImageDraw.Draw(image) |
|
|
|
return imagesToReturn |
|
|