Contador

app.py

@@ -32,6 +32,9 @@ Este Notebook se acelera opcionalmente con un entorno de ejecución de GPU
 #!pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt gradio # install dependencies
 
 import os
+import re
+import json
+import pandas as pd
 import gradio as gr
 import torch
 from PIL import Image
@@ -47,11 +50,41 @@ torch.hub.download_url_to_file('https://i.pinimg.com/originals/c2/ce/e0/c2cee056
 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
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-os.environ["HF_DATASETS_OFFLINE"] = "1"
-os.environ["TRANSFORMERS_OFFLINE"] = "1"
-
-hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "datasets/AllAideas/demo-iazika-flags")
+def getQuantity(string): 
+    contador_raw = ''.join(string.split(" ")[3:])
+    
+    resultado_especie_1 = 'Aedes'
+    resultado_especie_2 = 'Mosquito'
+    resultado_especie_3 = 'Mosca'
+    resultado_cantidad_1 = ''.join(re.findall(r'\d+',''.join(re.findall(r'\d+'+resultado_especie_1, contador_raw))))
+    resultado_cantidad_2 = ''.join(re.findall(r'\d+',''.join(re.findall(r'\d+'+resultado_especie_2, contador_raw))))
+    resultado_cantidad_3 = ''.join(re.findall(r'\d+',''.join(re.findall(r'\d+'+resultado_especie_3, contador_raw))))
+    resultado_cantidad_1 = resultado_cantidad_1 if len(resultado_cantidad_1) > 0 else 0
+    resultado_cantidad_2 = resultado_cantidad_2 if len(resultado_cantidad_2) > 0 else 0
+    resultado_cantidad_3 = resultado_cantidad_3 if len(resultado_cantidad_3) > 0 else 0
+    
+    resultado_lista = [[resultado_cantidad_1,resultado_especie_1],
+                       [resultado_cantidad_2,resultado_especie_2],
+                       [resultado_cantidad_3,resultado_especie_3]]
+    
+    return resultado_lista
+
+def listJSON(resultado):
+    resultado_lista = getQuantity(resultado)
+    img_name = " ".join(resultado.split(" ")[0:2])
+    img_size = "".join(resultado.split(" ")[2])
+    strlista = ""
+    for resultado_lista, description in resultado_lista:
+        strlista += '{"quantity":"'+resultado_lista+'","description":"'+description+'"},'
+    strlista = strlista[:-1]
+    str_resultado_lista = '{"image":"'+str(img_name)+'","size":"'+str(img_size)+'","detail":['+strlista+']}'
+    json_string = json.loads(str_resultado_lista)
+    return json_string
+    
+def arrayLista(resultado):
+    resultado_lista = getQuantity(resultado)
+    df = pd.DataFrame(resultado_lista,columns=['Cantidad','Especie'])
+    return df
 
 def yolo(size, iou, conf, im):
     '''Wrapper fn for gradio'''
@@ -66,8 +99,12 @@ def yolo(size, iou, conf, im):
     results2 = model(im)  # inference
    
     results2.render()  # updates results.imgs with boxes and labels
-    return Image.fromarray(results2.ims[0])
-
+    
+    results_detail = str(results2)
+    lista = listJSON(results_detail)
+    lista2 = arrayLista(results_detail)
+    return Image.fromarray(results2.ims[0]), lista2, lista
+    
 #------------ Interface-------------
 
 
@@ -78,6 +115,8 @@ in3 = gr.inputs.Slider(minimum=0, maximum=1, step=0.05, default=0.50, label='Umb
 in4 = gr.inputs.Image(type='pil', label="Original Image")
 
 out2 = gr.outputs.Image(type="pil", label="YOLOv5")
+out3 = gr.outputs.Dataframe(label="Descripción", headers=['Cantidad','Especie'])  
+out4 = gr.outputs.JSON(label="JSON")
 #-------------- Text-----
 title = 'Trampas Barceló'
 description = """
@@ -94,15 +133,13 @@ examples = [['640',0.45, 0.75,'ejemplo1.jpg'], ['640',0.45, 0.75,'ejemplo2.jpg']
 
 iface = gr.Interface(yolo, 
                         inputs=[in1, in2, in3, in4], 
-                        outputs=out2, title=title, 
+                        outputs=[out2,out3,out4], title=title, 
                         description=description, 
                         article=article, 
                         examples=examples,
                         theme="huggingface", 
-                        analytics_enabled=False,
-                        allow_flagging="manual",
-                        flagging_callback=hf_writer
-                        ).launch(enable_queue=True,debug=True)
+                        analytics_enabled=False
+                        ).launch(debug=True)
 
 iface.launch()