Spaces:
Sleeping
Sleeping
import os | |
import random | |
import logging | |
import gradio as gr | |
from PIL import Image | |
from zipfile import ZipFile | |
from typing import Any, Dict,List | |
from transformers import pipeline | |
class Image_classification: | |
def __init__(self): | |
pass | |
def unzip_image_data(self) -> str: | |
""" | |
Unzips an image dataset into a specified directory. | |
Returns: | |
str: The path to the directory containing the extracted image files. | |
""" | |
try: | |
with ZipFile("image_dataset.zip","r") as extract: | |
directory_path=str("dataset") | |
os.mkdir(directory_path) | |
extract.extractall(f"{directory_path}") | |
return f"{directory_path}" | |
except Exception as e: | |
logging.error(f"An error occurred during extraction: {e}") | |
return "" | |
def example_images(self) -> List[str]: | |
""" | |
Unzips the image dataset and generates a list of paths to the individual image files and use image for showing example | |
Returns: | |
List[str]: A list of file paths to each image in the dataset. | |
""" | |
try: | |
image_dataset_folder = self.unzip_image_data() | |
image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'] | |
image_count = len([name for name in os.listdir(image_dataset_folder) if os.path.isfile(os.path.join(image_dataset_folder, name)) and os.path.splitext(name)[1].lower() in image_extensions]) | |
example=[] | |
for i in range(image_count): | |
for name in os.listdir(image_dataset_folder): | |
path=(os.path.join(os.path.dirname(image_dataset_folder),os.path.join(image_dataset_folder,name))) | |
example.append(path) | |
return example | |
except Exception as e: | |
logging.error(f"An error occurred in example images: {e}") | |
return "" | |
def classify(self, image: Image.Image, model: Any) -> Dict[str, float]: | |
""" | |
Classifies an image using a specified model. | |
Args: | |
image (Image.Image): The image to classify. | |
model (Any): The model used for classification. | |
Returns: | |
Dict[str, float]: A dictionary of classification labels and their corresponding scores. | |
""" | |
try: | |
classifier = pipeline("image-classification", model=model) | |
result= classifier(image) | |
return result | |
except Exception as e: | |
logging.error(f"An error occurred during image classification: {e}") | |
raise | |
def format_the_result(self, image: Image.Image, model: Any) -> Dict[str, float]: | |
""" | |
Formats the classification result by retaining the highest score for each label. | |
Args: | |
image (Image.Image): The image to classify. | |
model (Any): The model used for classification. | |
Returns: | |
Dict[str, float]: A dictionary with unique labels and the highest score for each label. | |
""" | |
try: | |
data=self.classify(image,model) | |
new_dict = {} | |
for item in data: | |
label = item['label'] | |
score = item['score'] | |
if label in new_dict: | |
if new_dict[label] < score: | |
new_dict[label] = score | |
else: | |
new_dict[label] = score | |
return new_dict | |
except Exception as e: | |
logging.error(f"An error occurred while formatting the results: {e}") | |
raise | |
def interface(self): | |
with gr.Blocks(css=""".gradio-container {background: #314755; | |
background: -webkit-linear-gradient(to right, #26a0da, #314755); | |
background: linear-gradient(to right, #26a0da, #314755);} | |
.block svelte-90oupt padded{background:314755;}""") as demo: | |
gr.HTML(""" | |
<center><h1 style="color:#fff">Image Classification</h1></center>""") | |
exam_img=self.example_images() | |
with gr.Row(): | |
model = gr.Dropdown(["facebook/regnet-x-040","google/vit-large-patch16-384","microsoft/resnet-50",""],label="Choose a model") | |
with gr.Row(): | |
image = gr.Image(type="filepath",sources="upload") | |
with gr.Column(): | |
output=gr.Label() | |
with gr.Row(): | |
button=gr.Button() | |
button.click(self.format_the_result,[image,model],output) | |
gr.Examples( | |
examples=exam_img, | |
inputs=[image], | |
outputs=output, | |
fn=self.format_the_result, | |
cache_examples=False, | |
) | |
demo.launch(debug=True) | |
if __name__=="__main__": | |
image_classification=Image_classification() | |
result=image_classification.interface() |