|
import commands.exec_path |
|
from ultralytics import YOLO |
|
from PIL import Image, ImageDraw, ImageFont, ImageFile |
|
import os |
|
import random |
|
|
|
model_path = os.path.join(os.getcwd(), 'cv_files/AniClassifier.pt') |
|
model = YOLO(model_path) |
|
|
|
|
|
def img_classifier(image, classifer_type=0): |
|
|
|
test_images = [] |
|
test_images.append(image) |
|
imagesToReturn = [] |
|
|
|
|
|
folder_dir = './Images' |
|
if not os.path.exists(folder_dir): |
|
os.makedirs(folder_dir) |
|
|
|
|
|
for img in test_images: |
|
img_loc = img |
|
img_class = model(img_loc, verbose=False) |
|
|
|
|
|
if img_class[0].probs.data[0] < img_class[0].probs.data[1]: |
|
|
|
|
|
if classifer_type: |
|
image = Image.open(img_loc) |
|
image.save(folder_dir + img) |
|
|
|
|
|
imagesToReturn.append(folder_dir + img) |
|
return imagesToReturn |
|
|
|
|
|
else: |
|
return True |
|
|
|
else: |
|
return False |