pytorch 版fasterrcnn 跑通代碼增加保存xml的功能

1.下載（代碼和預訓練模型）

最近，領導說你找個pytorch版的fasterrcnn復現一下吧，以后都轉到pytorch上來吧，他終於肯放棄公司自己dragon框架了嗎。

問了同事說他有跑過一個代碼，下載jwyang/faster-rcnn.pytorch選擇pytorch1.0的分支，和res101的預訓練模型（又小又准）

2.環境，編譯

cuda10.0  python3.7   pytorch 1.0 

cd lib
python setup.py build develop

3.數據預處理和預訓練模型的放置

還是JPEGImages 文件夾放圖片， Annotations 放voc的標注文件，ImageSets里面划分訓練和測試集數據要放在data/VOCdevkit2007/VOC2007/路徑下，如下圖（代碼中也可以自己指定數據路徑，后文會寫），預訓練模型放在data/pretrained_model/路徑下

 

 

4.開始訓練

錯誤記錄

4.1改好參數之后run trainval_net.py開始報錯cannot import name '_mask' from 'lib.pycocotools

找了很多網上的解決方法

Just find out that you have to install the CoCO API.

cd data

git clone https://github.com/pdollar/coco.git cd coco/PythonAPI

我把整個coco文件夾復制覆蓋到原來的代碼中

刪除lib下的build文件夾重新編譯

python setup.py build develop

4.2繼續run的時候又報錯了UnboundLocalError: local variable 'pooled_feat' referenced before assignment


 

 

 

定位到faster_rcnn.py 82行，發現沒有什么問題，debug時候發現上面的if判斷語句都沒有進，然而cfg文件中POOLING_MODE 明明給的是align，看下面的信息發現cfg文件根本沒喲傳過來，所以兩種解決方法

1，cfg內容傳遞過來

2.如下圖，去掉判斷語句直接給pooled_feat賦值

 

 

至此終於可以正常run了

5.改進

5.1很多時候我不想在數據代碼混為一灘，數據讀取路徑可以直接賦值，在lib/datasets/pascal_voc.py中

把self._data_path修改為自己的數據路徑，下面就是自己的數據類別

 

 5.2demo.py中修改保存圖片的路徑，個人習慣吧，總是習慣把推斷過后的圖片另存一個文件夾

parse_args中新增一個output_images為存放新生成的預測后畫框的圖片

if vis and webcam_num == -1:
# cv2.imshow('test', im2show)
# cv2.waitKey(0)
result_path = os.path.join(args.output_images, imglist[num_images][:-4] + "_det.jpg") ###stt change output_images dir
cv2.imwrite(result_path, im2show)
 5.2demo.py中修改保存圖片的路徑，預測后的圖片xml文件保存出來（減少后期標注量）
新增較多

 

pred_det = []
for j in xrange(1, len(pascal_classes)):
 inds = torch.nonzero(scores[:,j]>thresh).view(-1)
 # if there is det
 if inds.numel() > 0: ### stt 判斷預測出來的類別中是否有框
 cls_scores = scores[:,j][inds] ### stt預測出來的所有得分
 _, order = torch.sort(cls_scores, 0, True)
 if args.class_agnostic:
 cls_boxes = pred_boxes[inds, :]
 else:
 cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4] ###stt 預測出來的所有框
 
 cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1) ### stt框和得分降成1維
 # cls_dets = torch.cat((cls_boxes, cls_scores), 1)
 cls_dets = cls_dets[order]
 # keep = nms(cls_dets, cfg.TEST.NMS, force_cpu=not cfg.USE_GPU_NMS)
 keep = nms(cls_boxes[order, :], cls_scores[order], cfg.TEST.NMS) #### stt預測出來的框nms
 cls_dets = cls_dets[keep.view(-1).long()]
 ### stt add
 arr_2 =[]
 np.set_printoptions(suppress=True) ###不以科學計數法保存
 arr_cls_det = cls_dets.cpu().numpy() ###張量轉數組
 arr_cls_det.astype(float) ###列表數據轉為float格式
 for i in arr_cls_det:
 arr_2.append(pascal_classes[j]) ###類別加入新的數組中
 arr_2 += map(int, i.tolist()[:4]) ###預測出來的框（四個坐標點）加入數組中
 pred_det.append(arr_2) ###類比，框保存為一個列表，每張圖片可能有多個預測框
 #print(arr_cls_det)
 #print((np.round(cls_dets)))
 if vis:
 im2show = vis_detections(im2show, pascal_classes[j], cls_dets.cpu().numpy(), 0.5)
#print(pred_det)
#######stt add save_txt and xml
save_txt = True
save_xml = True
### Write results:
for i in pred_det:
 if save_txt: # Write to file
 img_name = im_file.split('/')[-1].split('.')[0]
 txt_path = os.path.join(args.output_images, 'txt_path')
 if not os.path.exists(txt_path):
 os.makedirs(txt_path)
 with open(txt_path + '/' + img_name + '.txt', 'a') as file:
 file.write(('%s, %g , %g, %g, %g ' + '\n') % (i[0], i[1], i[2], i[3], i[4]))
 txtfile = os.path.join(txt_path, img_name + '.txt')
 ###stt _add start write_xml
 if save_txt and save_xml:
 Anno_path = os.path.join(args.output_images, 'Annotations')
 if not os.path.exists(Anno_path):
 os.makedirs(Anno_path)
 outdir = os.path.join(Anno_path,
 img_name + '.xml')
 txt2xml(pascal_classes, txtfile, im_file, outdir)
####stt add end
misc_toc = time.time()
nms_time = misc_toc - misc_tic

 

還有一個save_xml的文件夾放在下面鏈接中（自己修改過后代碼）