Modify arrayLista and JSON

app.py

@@ -36,6 +36,7 @@ import pandas as pd
 import torch
 import logging
 import json
+import re
 from PIL import Image
 
 # Images
@@ -49,32 +50,51 @@ torch.hub.download_url_to_file('https://i.pinimg.com/564x/3b/2f/d4/3b2fd4b6881b6
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True, autoshape=True)  # local model  o google colab
 #model = torch.hub.load('path/to/yolov5', 'custom', path='/content/yolov56.pt', source='local')  # local repo
 
+def removeStr(string): 
+    return string.replace(" ", "")
+
 def listJSON(a,b,c,d,e,f):
-      #for n in 
-      if d =='Pelicano\nSp':
+      x = re.findall("obo mar", d)
+      y = re.findall("elica", d)
+      z = re.findall("elica", f)
+      if x:
+        d = 'Lobo marino'
+      if y:
         d = 'Pelicano'
-      if f!='Pelicano':
-        strlista = '"detail":[{"quantity":"'+str(c)+'","description":"'+str(d)+'"}]'
-      else:
-        strlista = '"detail":[{"quantity":"'+str(c)+'","description":"'+str(d)+'"},{"quantity":"'+str(e)+'","description":"'+str(f)+'"}]'
-      strlist = '{"image":"'+str(a)+'","size":"'+str(b)+'",'+strlista+'}'
-      json_string = json.loads(strlist)
-      return json_string
+      if z:
+        f = 'Pelicano'
+      if(d=='Lobo marino' or d=='Pelicano'):
+        if d =='Pelicano\nSp' or d =='Pelicano\nS':
+          d = 'Pelicano'
+        if f!='Pelicano':
+          strlista = '"detail":[{"quantity":"'+str(removeStr(c))+'","description":"'+str(d)+'"}]'
+        else:
+          strlista = '"detail":[{"quantity":"'+str(removeStr(c))+'","description":"'+str(d)+'"},{"quantity":"'+str(removeStr(e))+'","description":"'+str(f)+'"}]'
+        strlist = '{"image":"'+str(removeStr(a))+'","size":"'+str(removeStr(b))+'",'+strlista+'}'
+        json_string = json.loads(strlist)
+        return json_string
 
 def arrayLista(a,b,c,d):
-      strlist =[]
-      strlist2 =[]
-      strlist.append(a)
-      if b =='Pelicano\nSp':
-        strlist.append('Pelicano')
-      else:
+      x = re.findall("obo mar", b)
+      y = re.findall("elica", b)
+      z = re.findall("elica", d)
+      if x:
+        b = 'Lobo marino'
+      if y:
+        b = 'Pelicano'
+      if z:
+        d = 'Pelicano'
+      if(b=='Lobo marino' or b=='Pelicano'):
+        strlist =[]
+        strlist2 =[]
+        strlist.append(removeStr(a))
         strlist.append(b)
-      if d=='Pelicano':
-        strlist2.append(c)
-        strlist2.append(d)
-      strlista = [strlist,strlist2]
-      df = pd.DataFrame(strlista,columns=['Cantidad','Especie'])
-      return df
+        if d=='Pelicano':
+          strlist2.append(removeStr(c))
+          strlist2.append(d)
+        strlista = [strlist,strlist2]
+        df = pd.DataFrame(strlista,columns=['Cantidad','Especie'])
+        return df
 
 def yolo(size, iou, conf, im):
     try:
@@ -89,8 +109,8 @@ def yolo(size, iou, conf, im):
         results2 = model(im)  # inference
         results2.render()  # updates results.imgs with boxes and labels
         results3 = str(results2)
-        lista = listJSON(results3[0:9], results3[11:18] ,results3[19:20],results3[21:32], results3[35:36], results3[37:45])
-        lista2 = arrayLista(results3[19:20],results3[21:32], results3[35:36], results3[37:45])
+        lista = listJSON(results3[0:9], results3[11:18] ,results3[19:21],results3[22:32], results3[35:37], results3[37:45])
+        lista2 = arrayLista(results3[19:21],results3[22:32], results3[35:37], results3[37:45])        
         return Image.fromarray(results2.ims[0]), lista2, lista
     except Exception as e:
         logging.error(e, exc_info=True)