add output txt

app.py

@@ -18,13 +18,13 @@ def download_model(url, local_path):
         response = requests.get(url, stream=True)
 
         # Get the total file size in bytes
-        total_size = int(response.headers.get('content-length', 0))
+        total_size = int(response.headers.get("content-length", 0))
 
         # Initialize the tqdm progress bar
-        progress_bar = tqdm(total=total_size, unit='B', unit_scale=True)
+        progress_bar = tqdm(total=total_size, unit="B", unit_scale=True)
 
         # Open a local file with write-binary mode
-        with open(local_path, 'wb') as file:
+        with open(local_path, "wb") as file:
             for data in response.iter_content(chunk_size=1024):
                 # Update the progress bar
                 progress_bar.update(len(data))
@@ -40,7 +40,7 @@ def download_model(url, local_path):
 
 def inference(filename):
     if not filename:
-        filename = './examples/butterfly.jpg'
+        filename = "./examples/butterfly.jpg"
 
     detector = InsectDetector()
     identifier = InsectIdentifier()
@@ -55,6 +55,7 @@ def inference(filename):
     image_for_draw = image.copy()
     image_height, image_width = image.shape[:2]
     boxes, confs, classes = detector.detect(image)
+    text = "Unknown"
 
     for box, _, _ in zip(boxes, confs, classes):
         box = box.astype(np.int32)
@@ -67,36 +68,26 @@ def inference(filename):
         cropped = khandy.crop_or_pad(image, box[0], box[1], box[2], box[3])
         results = identifier.identify(cropped)
         print(results[0])
-        prob = results[0]['probability']
-
-        if prob < 0.10:
-            text = 'Unknown'
-        else:
-            text = '{}: {:.2f}%'.format(
-                results[0]['latin_name'],
-                100.0 * results[0]['probability']
+        prob = results[0]["probability"]
+
+        if prob >= 0.10:
+            text = "{}: {:.2f}%".format(
+                results[0]["latin_name"], 100.0 * results[0]["probability"]
             )
 
         position = [box[0] + 2, box[1] - 20]
         position[0] = min(max(position[0], 0), image_width)
         position[1] = min(max(position[1], 0), image_height)
         cv2.rectangle(
-            image_for_draw,
-            (box[0], box[1]),
-            (box[2], box[3]),
-            (0, 255, 0),
-            2
+            image_for_draw, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2
         )
 
         image_for_draw = khandy.draw_text(
-            image_for_draw,
-            text,
-            position,
-            font='simsun.ttc',
-            font_size=15
+            image_for_draw, text, position, font="simsun.ttc", font_size=15
         )
 
-    return Image.fromarray(image_for_draw[:, :, ::-1], mode='RGB')
+    outxt = text.split(":")[0] if ":" in text else text
+    return Image.fromarray(image_for_draw[:, :, ::-1], mode="RGB"), outxt
 
 
 domain = "https://huggingface.co/MuGeminorum/insecta/resolve/main/"
@@ -104,27 +95,30 @@ domain_zh = "https://www.modelscope.cn/api/v1/models/MuGeminorum/insecta/repo?Re
 try:
     download_model(
         f"{domain}quarrying_insect_detector.onnx",
-        "./insectid/models/quarrying_insect_detector.onnx"
+        "./insectid/models/quarrying_insect_detector.onnx",
     )
     download_model(
         f"{domain}quarrying_insect_identifier.onnx",
-        "./insectid/models/quarrying_insect_identifier.onnx"
+        "./insectid/models/quarrying_insect_identifier.onnx",
     )
 except Exception:
     download_model(
         f"{domain_zh}quarrying_insect_detector.onnx",
-        "./insectid/models/quarrying_insect_detector.onnx"
+        "./insectid/models/quarrying_insect_detector.onnx",
     )
     download_model(
         f"{domain_zh}quarrying_insect_identifier.onnx",
-        "./insectid/models/quarrying_insect_identifier.onnx"
+        "./insectid/models/quarrying_insect_identifier.onnx",
     )
 
 iface = gr.Interface(
     fn=inference,
-    inputs=gr.Image(label='Upload insect photo', type='filepath'),
-    outputs=gr.Image(label='Detection result'),
-    examples=['./examples/butterfly.jpg', './examples/beetle.jpg']
+    inputs=gr.Image(label="Upload insect photo", type="filepath"),
+    outputs=[
+        gr.Image(label="Detection result"),
+        gr.Textbox(label="Most probable species", show_copy_button=True),
+    ],
+    examples=["./examples/butterfly.jpg", "./examples/beetle.jpg"],
 )
 
 iface.launch()