K-A-Uthman commited on
Commit
215c17e
·
1 Parent(s): 1664376

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
3
+
4
+ # Load the feature extractor and model directly
5
+ feature_extractor = AutoFeatureExtractor.from_pretrained("Devarshi/Brain_Tumor_Classification")
6
+ model = AutoModelForImageClassification.from_pretrained("Devarshi/Brain_Tumor_Classification")
7
+
8
+ # Define the prediction function using the loaded model
9
+ def classify_image(image):
10
+ # Preprocess the image and obtain features
11
+ inputs = feature_extractor(images=image, return_tensors="pt")
12
+ # Make prediction using the model
13
+ outputs = model(**inputs)
14
+ logits = outputs.logits
15
+ # Get the predicted class and confidence of the prediction
16
+ predicted_class = logits.argmax(dim=1).item()
17
+ confidence = logits.softmax(dim=1).max().item()
18
+
19
+ # Map the predicted class to the correct names
20
+ class_names = ["glioma_tumor", "meningioma_tumor", "no_tumor", "pituitary_tumor"]
21
+ predicted_class_text = class_names[predicted_class]
22
+
23
+ return {"prediction": predicted_class_text, "confidence": confidence}
24
+
25
+ # Define the Gradio interface
26
+ iface = gr.Interface(
27
+ fn=classify_image,
28
+ inputs=gr.inputs.Image(),
29
+ outputs="json",
30
+ title="Brain Tumor Image Classification",
31
+ description="This app classifies images of brain tumors into different classes.",
32
+ )
33
+
34
+ # Launch the Gradio interface
35
+ if __name__ == "__main__":
36
+ iface.launch()