Update app.py

app.py

@@ -85,10 +85,12 @@ def detect(img,thr=0.2):
       lab = labels[i][scr > thrh]
       box = boxes[i][scr > thrh]
   
-  
+      label_dict = {7:'person',6:'vehicle',5:'swimming-pool',4:'ship',3:'storage-tank',2:'bridge',1:'airplane'}         #AITOD labels
+      label_color_dict = {7:'red',6:'blue',5:'green',4:'yellow',3:'cyan',2:'magenta',1:'white'}
       for idx,b in enumerate(box):
-          draw.rectangle(list(b), outline='red', )
-          draw.text((b[0], b[1]), text=str(lab[idx]), fill='blue', )
+          label_i = lab[idx].item()
+          draw.rectangle(list(b), outline=label_color_dict[label_i], )
+          draw.text((b[0], b[1]), text=label_dict[label_i], fill='blue', )
   
   #save_path = Path('./output') / img_path.name
   return img