timri / app.py
Pranomvignesh
Added localization
510f06a
import gradio as gr
import yolov5
import os
from transformers import pipeline
imageClassifier = pipeline(task="image-classification",
model="Ara88/timri-model")
model = yolov5.load('./gentle-meadow.pt', device="cpu")
def predict(image):
predictions = imageClassifier(image)
maxScore = 0
predictedLabel = None
labelsForLocalization = ['meningioma', 'pituitary']
output = {}
for item in predictions:
output[item['label']] = item['score']
if (maxScore < item['score']):
maxScore = item['score']
predictedLabel = item['label']
if (predictedLabel in labelsForLocalization):
results = model([image], size=224)
imageWithLocalization = results.render()[0]
else:
imageWithLocalization = image
return output, imageWithLocalization
title = "Detecting Tumors in MRI Images"
description = """
Try the examples at bottom to get started.
"""
examples = [
[os.path.abspath('examples/sample_1.jpg')],
[os.path.abspath('examples/sample_2.jpg')],
[os.path.abspath('examples/sample_3.jpg')],
[os.path.abspath('examples/sample_4.jpg')],
[os.path.abspath('examples/sample_5.jpg')],
[os.path.abspath('examples/sample_6.jpg')],
[os.path.abspath('examples/sample_7.jpg')],
[os.path.abspath('examples/sample_8.jpg')],
]
inputs = gr.Image(type="pil", shape=(224, 224),
label="Upload your image for detection")
outputs = [
gr.Label(label="Tumor Classification"),
gr.Image(type="pil", label="Tumor Detections")
]
interface = gr.Interface(
fn=predict,
inputs=inputs,
outputs=outputs,
title=title,
examples=examples,
description=description,
cache_examples=True,
theme='huggingface'
)
interface.launch(debug=True, enable_queue=True)