Spaces:
Sleeping
Sleeping
File size: 4,442 Bytes
f64042e 53727cd f64042e 53727cd f64042e |
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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
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() |