yolov5保存中心点坐标和宽高到txt

注:在detect.py中下面的’–save-txt’参数设置为True,可以在检测时候,直接保存TXT,这是官方自带的,当然也可以自己diy这样保存。

打开detect.py文件
找预测图片的宽高
在这里插入图片描述
代码:

for path, im, im0s, vid_cap, s in dataset:
       print("test!!!",s)
       # 这里打印了s字符串,可以看出是很长的一串
       
       s_folder_name = s[-16:-13]
       print("s_folder_name:",s_folder_name)
       #裁剪出要下面保存的路径,是根据这个图片名字截取的后半段
       
       s_jpg_name = s[-12:-6]
       print("s_jpg_name:",s_jpg_name)
       #裁剪出要下面保存的txt名字,是根据这个图片名字截取的后半段
       #连接s_folder_name和s_jpg_name的话,也就是下面要保存的txt为./022/000022.txt
       #这样就拼接处了txt的路径和名字
       
       #拼接要保存结果的txt的路径和名字
       filename_txt = './result/' + s_folder_name + '/' + s_jpg_name + '.txt'
       #打开此txt
       file_txt = open(filename_txt, 'a')
       
       t1 = time_sync()
       im = torch.from_numpy(im).to(device)
       im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
       im /= 255  # 0 - 255 to 0.0 - 1.0
       if len(im.shape) == 3:
           im = im[None]  # expand for batch dim
       t2 = time_sync()
       dt[0] += t2 - t1

       # Inference
       visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
       pred = model(im, augment=augment, visualize=visualize)
       t3 = time_sync()
       dt[1] += t3 - t2

       # NMS
       pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
       dt[2] += time_sync() - t3

       # Second-stage classifier (optional)
       # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

       # Process predictions
       #f = open(txt_path + '.txt', 'a')
       #with open(txt_path + '.txt', 'a') as f:
       for i, det in enumerate(pred):  # per image
           seen += 1
           if webcam:  # batch_size >= 1
               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
           save_path = str(save_dir / p.name)  # im.jpg
           txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

		   # 这里其实在拼s字符串,终端打印的,和我们保存中心点坐标和宽高到txt操作没关系
           s += '%gx%g ' % im.shape[2:]  # print string 
           gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
           imc = im0.copy() if save_crop else im0  # for save_crop
           annotator = Annotator(im0, line_width=line_thickness, example=str(names))
           if len(det):
               # Rescale boxes from img_size to im0 size
               det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

               # 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,n box number
               # Write results
               for *xyxy, conf, cls in reversed(det):
               	   #这个if True地方,原来是if save_img,后来未来保存下面的xywh,我给强行改成trueif True:  # Write to file   
                   	   # xywh其实就是x_center,y_center,width,height,详细可以自己点进去xyxy2xywh这个函数看,文章后面我贴了           
                       xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                       line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                       #  这个file_txt.write里的东西是根据原来代码里改的,保存到上面拼接的txt里
                       file_txt.write(('%g ' * len(line)).rstrip() % line + '\n')

                   if save_img or save_crop or view_img:  # Add bbox to image
                       c = int(cls)  # integer class
                       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:
               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[i] != save_path:  # new video
                       vid_path[i] = save_path
                       if isinstance(vid_writer[i], cv2.VideoWriter):
                           vid_writer[i].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 = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
                       vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                   vid_writer[i].write(im0)
       #注意关闭这个文件
       file_txt.close()
       # Print time (inference-only)
       LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
        
        
    # Print results
    t = tuple(x / 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)  # update model (to fix SourceChangeWarning)

xyxy2xywh函数:
在这里插入图片描述

def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

猜你喜欢

转载自blog.csdn.net/Qingyou__/article/details/124108356
今日推荐