Update detect.py

class  text 추출
text 추출한걸 음성으로 출력

Signed-off-by: Leedong414 <165615367+Leedong414@users.noreply.github.com>

detect.py

@@ -42,34 +42,27 @@ from gtts import gTTS
 def detect(opt):
     source, weights, conf, save_txt_path = opt.source, opt.weights, opt.conf, opt.save_txt_path
 
-    # 모델 로드
     model = torch.hub.load("ultralytics/yolov5", "custom", path=weights)
 
-    # 이미지 로드 및 추론 수행
     results = model(source)
 
-    # 결과를 pandas 데이터프레임으로 변환
     results_df = results.pandas().xyxy[0]
 
-    # 데이터프레임의 클래스 컬럼 추출
     classes = results_df["name"].tolist()
 
-    # 클래스 정보를 텍스트 파일로 저장
     with open(save_txt_path, "w") as f:
         for cls in classes:
             f.write(f"{cls}\n")
 
-    # 콘솔에 출력
     print("Detected Classes:")
     for cls in classes:
         print(cls)
 
-    # TTS를 사용하여 클래스 이름들을 음성으로 변환
     if classes:
         text_to_speak = "Detected classes are: " + ", ".join(classes)
         tts = gTTS(text=text_to_speak, lang="en")
         tts.save("detected_classes.mp3")
-        os.system("mpg321 detected_classes.mp3")  # mpg321 설치 필요 (리눅스의 경우)
+        os.system("mpg321 detected_classes.mp3") 
 
 
 if __name__ == "__main__":
@@ -79,7 +72,7 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--source", type=str, default="/HongRyeon_test01.jpg", help="source"
-    )  # file/folder, 0 for webcam
+    )  
     parser.add_argument("--conf", type=float, default=0.5, help="object confidence threshold")
     parser.add_argument(
         "--save_txt_path", type=str, default="detected_classes.txt", help="path to save detected classes txt file"
@@ -91,10 +84,10 @@ if __name__ == "__main__":
 
 
 FILE = Path(__file__).resolve()
-ROOT = FILE.parents[0]  # YOLOv5 root directory
+ROOT = FILE.parents[0]  
 if str(ROOT) not in sys.path:
-    sys.path.append(str(ROOT))  # add ROOT to PATH
+    sys.path.append(str(ROOT))  
-ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  
 
 from ultralytics.utils.plotting import Annotator, colors, save_one_box
 
@@ -121,56 +114,56 @@ from utils.torch_utils import select_device, smart_inference_mode
 
 @smart_inference_mode()
 def run(
-    weights=ROOT / "yolov5s.pt",  # model path or triton URL
+    weights=ROOT / "yolov5s.pt", 
-    source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
+    source=ROOT / "data/images",  
-    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
+    data=ROOT / "data/coco128.yaml", 
-    imgsz=(640, 640),  # inference size (height, width)
+    imgsz=(640, 640), 
-    conf_thres=0.25,  # confidence threshold
+    conf_thres=0.25,  
-    iou_thres=0.45,  # NMS IOU threshold
+    iou_thres=0.45, 
-    max_det=1000,  # maximum detections per image
+    max_det=1000, 
-    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    device="",  
-    view_img=False,  # show results
+    view_img=False, 
-    save_txt=False,  # save results to *.txt
+    save_txt=False, 
-    save_csv=False,  # save results in CSV format
+    save_csv=False,  
-    save_conf=False,  # save confidences in --save-txt labels
+    save_conf=False, 
-    save_crop=False,  # save cropped prediction boxes
+    save_crop=False,  
-    nosave=False,  # do not save images/videos
+    nosave=False,  
-    classes=None,  # filter by class: --class 0, or --class 0 2 3
+    classes=None,  
-    agnostic_nms=False,  # class-agnostic NMS
+    agnostic_nms=False, 
-    augment=False,  # augmented inference
+    augment=False, 
-    visualize=False,  # visualize features
+    visualize=False,  
-    update=False,  # update all models
+    update=False, 
-    project=ROOT / "runs/detect",  # save results to project/name
+    project=ROOT / "runs/detect",  
-    name="exp",  # save results to project/name
+    name="exp",  
-    exist_ok=False,  # existing project/name ok, do not increment
+    exist_ok=False, 
-    line_thickness=3,  # bounding box thickness (pixels)
+    line_thickness=3,  
-    hide_labels=False,  # hide labels
+    hide_labels=False,  
-    hide_conf=False,  # hide confidences
+    hide_conf=False,  
-    half=False,  # use FP16 half-precision inference
+    half=False, 
-    dnn=False,  # use OpenCV DNN for ONNX inference
+    dnn=False,  
-    vid_stride=1,  # video frame-rate stride
+    vid_stride=1,  
 ):
     source = str(source)
-    save_img = not nosave and not source.endswith(".txt")  # save inference images
+    save_img = not nosave and not source.endswith(".txt")  
     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
     is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
     webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
     screenshot = source.lower().startswith("screen")
     if is_url and is_file:
-        source = check_file(source)  # download
+        source = check_file(source) 
+
+  
+    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  
+    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  
 
-    # Directories
-    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
-    # Load model
     device = select_device(device)
     model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
     stride, names, pt = model.stride, model.names, model.pt
-    imgsz = check_img_size(imgsz, s=stride)  # check image size
+    imgsz = check_img_size(imgsz, s=stride)  
 
-    # Dataloader
+  
-    bs = 1  # batch_size
+    bs = 1  
     if webcam:
         view_img = check_imshow(warn=True)
         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
@@ -181,20 +174,20 @@ def run(
         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
     vid_path, vid_writer = [None] * bs, [None] * bs
 
-    # Run inference
+
-    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  
     seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
     for path, im, im0s, vid_cap, s in dataset:
         with dt[0]:
             im = torch.from_numpy(im).to(model.device)
-            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+            im = im.half() if model.fp16 else im.float()  
-            im /= 255  # 0 - 255 to 0.0 - 1.0
+            im /= 255  
             if len(im.shape) == 3:
-                im = im[None]  # expand for batch dim
+                im = im[None]  
             if model.xml and im.shape[0] > 1:
                 ims = torch.chunk(im, im.shape[0], 0)
 
-        # Inference
+
         with dt[1]:
             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
             if model.xml and im.shape[0] > 1:
@@ -207,17 +200,13 @@ def run(
                 pred = [pred, None]
             else:
                 pred = model(im, augment=augment, visualize=visualize)
-        # NMS
+ 
         with dt[2]:
             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
 
-        # Second-stage classifier (optional)
-        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
 
-        # Define the path for the CSV file
         csv_path = save_dir / "predictions.csv"
 
-        # Create or append to the CSV file
         def write_to_csv(image_name, prediction, confidence):
             """Writes prediction data for an image to a CSV file, appending if the file exists."""
             data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
@@ -227,34 +216,33 @@ def run(
                     writer.writeheader()
                 writer.writerow(data)
 
-        # Process predictions
+        for i, det in enumerate(pred):  
-        for i, det in enumerate(pred):  # per image
             seen += 1
-            if webcam:  # batch_size >= 1
+            if webcam:  
                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
                 s += f"{i}: "
             else:
                 p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
 
-            p = Path(p)  # to Path
+            p = Path(p)  
-            save_path = str(save_dir / p.name)  # im.jpg
+            save_path = str(save_dir / p.name) 
-            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
+            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  
-            s += "%gx%g " % im.shape[2:]  # print string
+            s += "%gx%g " % im.shape[2:] 
-            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  
-            imc = im0.copy() if save_crop else im0  # for save_crop
+            imc = im0.copy() if save_crop else im0 
             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
             if len(det):
-                # Rescale boxes from img_size to im0 size
+
                 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
 
-                # Print results
+         
                 for c in det[:, 5].unique():
-                    n = (det[:, 5] == c).sum()  # detections per class
+                    n = (det[:, 5] == c).sum()  
-                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " 
 
-                # Write results
+            
                 for *xyxy, conf, cls in reversed(det):
-                    c = int(cls)  # integer class
+                    c = int(cls)  
                     label = names[c] if hide_conf else f"{names[c]}"
                     confidence = float(conf)
                     confidence_str = f"{confidence:.2f}"
@@ -262,59 +250,56 @@ def run(
                     if save_csv:
                         write_to_csv(p.name, label, confidence_str)
 
-                    if save_txt:  # Write to file
+                    if save_txt:  
-                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  
-                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
+                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  
                         with open(f"{txt_path}.txt", "a") as f:
                             f.write(("%g " * len(line)).rstrip() % line + "\n")
 
-                    if save_img or save_crop or view_img:  # Add bbox to image
+                    if save_img or save_crop or view_img: 
-                        c = int(cls)  # integer class
+                        c = int(cls)  
                         label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
                         annotator.box_label(xyxy, label, color=colors(c, True))
                     if save_crop:
                         save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
 
-            # Stream results
+
             im0 = annotator.result()
             if view_img:
                 if platform.system() == "Linux" and p not in windows:
                     windows.append(p)
-                    cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
+                    cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) 
                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
                 cv2.imshow(str(p), im0)
-                cv2.waitKey(1)  # 1 millisecond
+                cv2.waitKey(1)  
 
-            # Save results (image with detections)
             if save_img:
                 if dataset.mode == "image":
                     cv2.imwrite(save_path, im0)
-                else:  # 'video' or 'stream'
+                else:  
-                    if vid_path[i] != save_path:  # new video
+                    if vid_path[i] != save_path:  
                         vid_path[i] = save_path
                         if isinstance(vid_writer[i], cv2.VideoWriter):
-                            vid_writer[i].release()  # release previous video writer
+                            vid_writer[i].release()  
-                        if vid_cap:  # video
+                        if vid_cap: 
                             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
+                        else: 
                             fps, w, h = 30, im0.shape[1], im0.shape[0]
-                        save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
+                        save_path = str(Path(save_path).with_suffix(".mp4"))  
                         vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
                     vid_writer[i].write(im0)
 
-        # Print time (inference-only)
         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
 
-    # Print results
+    t = tuple(x.t / seen * 1e3 for x in dt)
-    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
     LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
     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 ""
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
-        strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
+        strip_optimizer(weights[0])  
 
 
 def parse_opt():
@@ -349,7 +334,7 @@ def parse_opt():
     parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
     parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
     opt = parser.parse_args()
-    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  
     print_args(vars(opt))
     return opt