added count feature

app.py

@@ -1,11 +1,9 @@
 import gradio as gr
-import os
 
 import argparse
 import time
 from pathlib import Path
 
-import cv2
 import torch
 import torch.backends.cudnn as cudnn
 from numpy import random
@@ -14,7 +12,6 @@ from models.experimental import attempt_load
 from utils.datasets import LoadStreams, LoadImages
 from utils.general import (
     check_img_size,
-    check_imshow,
     non_max_suppression,
     apply_classifier,
     scale_coords,
@@ -26,7 +23,6 @@ from utils.plots import plot_one_box
 from utils.torch_utils import (
     select_device,
     load_classifier,
-    time_synchronized,
     TracedModel,
 )
 from PIL import Image
@@ -50,9 +46,7 @@ def detect(img):
     parser.add_argument(
         "--weights", nargs="+", type=str, default=loaded_model, help="model.pt path(s)"
     )
-    parser.add_argument(
-        "--source", type=str, default="Inference/", help="source"
-    )  # file/folder, 0 for webcam
+    parser.add_argument("--source", type=str, default="Inference/", help="source")
     parser.add_argument(
         "--img-size", type=int, default=640, help="inference size (pixels)"
     )
@@ -105,11 +99,6 @@ def detect(img):
         opt.trace,
     )
     save_img = True  # save inference images
-    webcam = (
-        source.isnumeric()
-        or source.endswith(".txt")
-        or source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
-    )
 
     # Directories
     save_dir = Path(
@@ -144,13 +133,7 @@ def detect(img):
         ).to(device).eval()
 
     # Set Dataloader
-    vid_path, vid_writer = None, None
-    if webcam:
-        view_img = check_imshow()
-        cudnn.benchmark = True  # set True to speed up constant image size inference
-        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
-    else:
-        dataset = LoadImages(source, img_size=imgsz, stride=stride)
+    dataset = LoadImages(source, img_size=imgsz, stride=stride)
 
     # Get names and colors
     names = model.module.names if hasattr(model, "module") else model.names
@@ -170,7 +153,6 @@ def detect(img):
             img = img.unsqueeze(0)
 
         # Inference
-        t1 = time_synchronized()
         pred = model(img, augment=opt.augment)[0]
 
         # Apply NMS
@@ -181,7 +163,6 @@ def detect(img):
             classes=opt.classes,
             agnostic=opt.agnostic_nms,
         )
-        t2 = time_synchronized()
 
         # Apply Classifier
         if classify:
@@ -189,17 +170,12 @@ def detect(img):
 
         # Process detections
         for i, det in enumerate(pred):  # detections per image
-            if webcam:  # batch_size >= 1
-                p, s, im0, frame = path[i], "%g: " % i, im0s[i].copy(), dataset.count
-            else:
-                p, s, im0, frame = path, "", im0s, getattr(dataset, "frame", 0)
+            p, s, im0, frame = path, "", im0s, getattr(dataset, "frame", 0)
 
             p = Path(p)  # to Path
-            save_path = str(save_dir / p.name)  # img.jpg
             txt_path = str(save_dir / "labels" / p.stem) + (
                 "" if dataset.mode == "image" else f"_{frame}"
             )  # img.txt
-            s += "%gx%g " % img.shape[2:]  # print string
             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
             if len(det):
                 # Rescale boxes from img_size to im0 size
@@ -208,7 +184,7 @@ def detect(img):
                 # Print results
                 for c in det[:, -1].unique():
                     n = (det[:, -1] == c).sum()  # detections per class
-                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)} "  # add to string
 
                 # Write results
                 for *xyxy, conf, cls in reversed(det):
@@ -234,51 +210,15 @@ def detect(img):
                             line_thickness=3,
                         )
 
-            # Print time (inference + NMS)
-            # print(f'{s}Done. ({t2 - t1:.3f}s)')
-
-            # Stream results
-            if view_img:
-                cv2.imshow(str(p), im0)
-                cv2.waitKey(1)  # 1 millisecond
-
-            # Save results (image with detections)
-            if save_img:
-                if dataset.mode == "image":
-                    cv2.imwrite(save_path, im0)
-                else:  # 'video' or 'stream'
-                    if vid_path != save_path:  # new video
-                        vid_path = save_path
-                        if isinstance(vid_writer, cv2.VideoWriter):
-                            vid_writer.release()  # release previous video writer
-                        if vid_cap:  # video
-                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
-                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                        else:  # stream
-                            fps, w, h = 30, im0.shape[1], im0.shape[0]
-                            save_path += ".mp4"
-                        vid_writer = cv2.VideoWriter(
-                            save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
-                        )
-                    vid_writer.write(im0)
-
-    if save_txt or save_img:
-        s = (
-            f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}"
-            if save_txt
-            else ""
-        )
-        # print(f"Results saved to {save_dir}{s}")
-
     print(f"Done. ({time.time() - t0:.3f}s)")
 
-    return Image.fromarray(im0[:, :, ::-1])
+    return [Image.fromarray(im0[:, :, ::-1]), s]
 
 
 gr.Interface(
-    detect,
-    [gr.Image(type="pil")],
-    gr.Image(type="pil"),
+    fn=detect,
     title="Anything Counter",
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Image(label="detection", type="pil"), gr.Textbox(label="count")],
+    allow_flagging="manual",
 ).launch(debug=True)