Added localization

.gitignore

@@ -0,0 +1 @@
+test.py

app.py

@@ -10,15 +10,23 @@ model = yolov5.load('./gentle-meadow.pt', device="cpu")
 
 
 def predict(image):
-    # results = model([image], size=224)
-    predictions = imageClassifier(image)
-    print(predictions)
 
+    predictions = imageClassifier(image)
+    maxScore = 0
+    predictedLabel = None
+    labelsForLocalization = ['meningioma', 'pituitary']
     output = {}
     for item in predictions:
         output[item['label']] = item['score']
-
-    return output
+        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"
@@ -40,8 +48,8 @@ inputs = gr.Image(type="pil", shape=(224, 224),
                   label="Upload your image for detection")
 
 outputs = [
-    # gr.Image(type="pil", label="Tumor Detections"),
-    gr.Label(label="Tumor Classification")
+    gr.Label(label="Tumor Classification"),
+    gr.Image(type="pil", label="Tumor Detections")
 ]
 
 interface = gr.Interface(

test.py

@@ -1,14 +0,0 @@
-from PIL import Image
-from gradio_client import Client
-
-client = Client("PranomVignesh/timri")
-
-with open('examples/sample_1.jpg', "rb") as f:
-    image = Image.open(f)
-    image_bytes = f.read()
-
-output = client.predict(
-    image_bytes
-)
-
-print(output)