Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
from transformers import AutoFeatureExtractor, AutoModelForImageClassification | |
from PIL import Image | |
import numpy as np | |
import cv2 | |
class ThyroidTumorClassifierApp: | |
def __init__(self): | |
# Load the feature extractor and model | |
self.extractor = AutoFeatureExtractor.from_pretrained("SerdarHelli/ThyroidTumorClassificationModel") | |
self.model = AutoModelForImageClassification.from_pretrained("SerdarHelli/ThyroidTumorClassificationModel") | |
def classify_image(self, image): | |
# Preprocess the image using the extractor | |
inputs = self.extractor(images=image, return_tensors="pt") | |
# Pass the image through the model | |
outputs = self.model(**inputs) | |
# Get the class probabilities | |
logits = outputs.logits | |
# Calculate the final probabilities using softmax | |
probabilities = torch.softmax(logits, dim=1) | |
# Get the class with the highest probability | |
predicted_class = torch.argmax(probabilities, dim=1).item() | |
# Customize class labels based on your model | |
class_labels = ["Sem Tumor", "Tumor"] | |
# Predicted class label | |
predicted_label = class_labels[predicted_class] | |
# Add information to the output image using OpenCV | |
output_image_with_info = self.add_info_to_image(image, predicted_label, probabilities) | |
# Return the modified output image as an array | |
return output_image_with_info | |
def add_info_to_image(self, image, predicted_label, probabilities): | |
# Convert the image to RGB format | |
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
# Add the predicted class label and probabilities to the image using OpenCV | |
font = cv2.FONT_HERSHEY_SIMPLEX | |
font_scale = 0.6 | |
font_thickness = 1 | |
text_color = (255, 255, 255) | |
text_position = (10, 30) | |
# Add predicted class label | |
cv2.putText(image_rgb, f"Classe Prevista: {predicted_label}", text_position, font, font_scale, text_color, font_thickness) | |
# Add class probabilities | |
for i, prob in enumerate(probabilities[0]): | |
y_offset = 60 + i * 30 | |
class_name = f"Classe {i}:" | |
probability = f"{prob:.2f}" | |
cv2.putText(image_rgb, f"{class_name} {probability}", (10, text_position[1] + y_offset), font, font_scale, text_color, font_thickness) | |
# Convert back to BGR format for display | |
output_image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) | |
return output_image | |
def run_interface(self): | |
# Create a Gradio interface | |
input_interface = gr.Interface( | |
fn=self.classify_image, | |
inputs=gr.inputs.Image(), | |
outputs=gr.outputs.Image(), | |
title="Tumor da Tireoide Classificação", | |
description="Faça o upload de uma imagem de um tumor da tireoide para classificação. A saída inclui o rótulo da classe prevista e as probabilidades com informações adicionais.", | |
) | |
# Launch the Gradio interface | |
input_interface.launch() | |
if __name__ == "__main__": | |
app = ThyroidTumorClassifierApp() | |
app.run_interface() | |