Spaces:

Kyo-Kai
/

Fsg-Pp

Sleeping

Fsg-Pp / ai /classifying_ai.py

Upload 18 files

f4d52c1 verified 5 months ago

No virus

1.42 kB

	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 = []

	# Create a directory for saving classified images
	folder_dir = './Images'
	if not os.path.exists(folder_dir):
	os.makedirs(folder_dir)

	# Classify images with "good" class in the images folder and save them in the image directory
	for img in test_images:
	img_loc = img
	img_class = model(img_loc, verbose=False)

	# If the first index is higher than the second index, the image is classified as "good"
	if img_class[0].probs.data[0] < img_class[0].probs.data[1]:

	# Save the image in the classified directory
	if classifer_type:
	image = Image.open(img_loc)
	image.save(folder_dir + img)

	# Appending Cropped images in an array to display in gradio for end-user
	imagesToReturn.append(folder_dir + img)
	return imagesToReturn

	# Downloading Thumbnail images so don't save them in the image directory
	else:
	return True

	else:
	return False