Update app.py

app.py

@@ -31,11 +31,15 @@ def fig2img(fig):
     buf = io.BytesIO()
     fig.savefig(buf)
     buf.seek(0)
-    img = Image.open(buf)
+    pil_img = Image.open(buf)
+    basewidth = 750
+    wpercent = (basewidth/float(pil_img.size[0]))
+    hsize = int((float(pil_img.size[1])*float(wpercent)))
+    img = pil_img.resize((basewidth,hsize), Image.ANTIALIAS) 
     return img
 
 
-def visualize_prediction(pil_img, output_dict, threshold=0.5, id2label=None):
+def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     keep = output_dict["scores"] > threshold
     boxes = output_dict["boxes"][keep].tolist()
     scores = output_dict["scores"][keep].tolist()
@@ -44,10 +48,6 @@ def visualize_prediction(pil_img, output_dict, threshold=0.5, id2label=None):
         labels = [id2label[x] for x in labels]
 
     plt.figure(figsize=(18, 13))
-    basewidth = 750
-    wpercent = (basewidth/float(pil_img.size[0]))
-    hsize = int((float(pil_img.size[1])*float(wpercent)))
-    img = pil_img.resize((basewidth,hsize), Image.ANTIALIAS)
     plt.imshow(img)
     ax = plt.gca()
     colors = COLORS * 100