python 使用Yolact訓練自己的數據集

可能是由於yolact官方更新過其項目代碼,所以網上其他人的yolact訓練使用的config文件和我的稍微有區別。但總體還是差不多的。

1:提前準備好自己的數據集

使用labelme來制作分割數據集,但是得到的是一個個單獨的json文件。需要將其轉換成coco。
labelme2coco.py如下所示(代碼來源:github鏈接):

import os
import json
import numpy as np
import glob
import shutil
from sklearn.model_selection import train_test_split
np.random.seed(41)

#0為背景,此處根據你數據集的類別來修改key
classname_to_id = {"1": 1}

class Lableme2CoCo:

 def __init__(self):
  self.images = []
  self.annotations = []
  self.categories = []
  self.img_id = 0
  self.ann_id = 0

 def save_coco_json(self, instance, save_path):
  json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1) # indent=2 更加美觀顯示

 # 由json文件構建COCO
 def to_coco(self, json_path_list):
  self._init_categories()
  for json_path in json_path_list:
   obj = self.read_jsonfile(json_path)
   self.images.append(self._image(obj, json_path))
   shapes = obj['shapes']
   for shape in shapes:
    annotation = self._annotation(shape)
    self.annotations.append(annotation)
    self.ann_id += 1
   self.img_id += 1
  instance = {}
  instance['info'] = 'spytensor created'
  instance['license'] = ['license']
  instance['images'] = self.images
  instance['annotations'] = self.annotations
  instance['categories'] = self.categories
  return instance

 # 構建類別
 def _init_categories(self):
  for k, v in classname_to_id.items():
   category = {}
   category['id'] = v
   category['name'] = k
   self.categories.append(category)

 # 構建COCO的image字段
 def _image(self, obj, path):
  image = {}
  from labelme import utils
  img_x = utils.img_b64_to_arr(obj['imageData'])
  h, w = img_x.shape[:-1]
  image['height'] = h
  image['width'] = w
  image['id'] = self.img_id
  image['file_name'] = os.path.basename(path).replace(".json", ".jpg")
  return image

 # 構建COCO的annotation字段
 def _annotation(self, shape):
  label = shape['label']
  points = shape['points']
  annotation = {}
  annotation['id'] = self.ann_id
  annotation['image_id'] = self.img_id
  annotation['category_id'] = int(classname_to_id[label])
  annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
  annotation['bbox'] = self._get_box(points)
  annotation['iscrowd'] = 0
  annotation['area'] = 1.0
  return annotation

 # 讀取json文件,返回一個json對象
 def read_jsonfile(self, path):
  with open(path, "r", encoding='utf-8') as f:
   return json.load(f)

 # COCO的格式: [x1,y1,w,h] 對應COCO的bbox格式
 def _get_box(self, points):
  min_x = min_y = np.inf
  max_x = max_y = 0
  for x, y in points:
   min_x = min(min_x, x)
   min_y = min(min_y, y)
   max_x = max(max_x, x)
   max_y = max(max_y, y)
  return [min_x, min_y, max_x - min_x, max_y - min_y]


if __name__ == '__main__':
 labelme_path = "labelme/" # 此處根據你的數據集地址來修改
 saved_coco_path = "./"
 # 創建文件
 if not os.path.exists("%scoco/annotations/"%saved_coco_path):
  os.makedirs("%scoco/annotations/"%saved_coco_path)
 if not os.path.exists("%scoco/images/train2017/"%saved_coco_path):
  os.makedirs("%scoco/images/train2017"%saved_coco_path)
 if not os.path.exists("%scoco/images/val2017/"%saved_coco_path):
  os.makedirs("%scoco/images/val2017"%saved_coco_path)
 # 獲取images目錄下所有的joson文件列表
 json_list_path = glob.glob(labelme_path + "/*.json")
 # 數據劃分,這裡沒有區分val2017和tran2017目錄,所有圖片都放在images目錄下
 train_path, val_path = train_test_split(json_list_path, test_size=0.12)
 print("train_n:", len(train_path), 'val_n:', len(val_path))

 # 把訓練集轉化為COCO的json格式
 l2c_train = Lableme2CoCo()
 train_instance = l2c_train.to_coco(train_path)
 l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json'%saved_coco_path)
 for file in train_path:
  shutil.copy(file.replace("json","jpg"),"%scoco/images/train2017/"%saved_coco_path)
 for file in val_path:
  shutil.copy(file.replace("json","jpg"),"%scoco/images/val2017/"%saved_coco_path)

 # 把驗證集轉化為COCO的json格式
 l2c_val = Lableme2CoCo()
 val_instance = l2c_val.to_coco(val_path)
 l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json'%saved_coco_path)

隻需要修改兩個地方即可,然後放到data文件夾下。
最後,得到的coco格式的數據集如下所示:

至此,數據準備已經結束。

2:下載github存儲庫

網址:YOLACT

之後解壓,但是我解壓的時候不知道為啥沒有yolact.py這個文件。後來又建瞭一個py文件,復制瞭裡面的代碼。

下載權重文件,把權重文件放到yolact-master下的weights文件夾裡(沒有就新建):

3:修改config.py

文件所在位置:

修改類別,把原本的coco的類別全部註釋掉,修改成自己的(如紅色框),註意COCO_CLASSES裡有一個逗號。

修改數據集地址dataset_base

修改coco_base_config(下面第二個橫線max_iter並不是控制訓練輪數的,第二張圖中的max_iter才是)

4:訓練

cd到指定路徑下,執行下面命令即可

python train.py --config=yolact_base_config

剛開始:

因為我是租的雲服務器,在jupyter notebook裡訓練的。輸出的訓練信息比較亂。

訓練幾分鐘後:

主要看T後面的數字即可,好像他就是總的loss,如果它收斂瞭,按下Ctrl+C,即可中止訓練,保存模型權重。

第一個問題:

PytorchStreamReader failed reading zip archive: failed finding central directory

第二個問題:
(但是不知道為啥,我訓練時如果中斷,保存的模型不能用來測試,會爆出下面的錯誤)

RuntimeError: unexpected EOF, expected *** more bytes. The file might be corruptrd

沒辦法解決,所以隻能跑完,自動結束之後保存的模型拿來測試(自動保存的必中斷保存的要大十幾兆)

模型保存的格式:<config>_<epoch>_<iter>.pth。如果是中斷的:<config>_<epoch>_<iter>_interrupt.pth

5:測試

使用官網的測試命令即可

以上就是python 使用Yolact訓練自己的數據集的詳細內容,更多關於python 訓練數據集的資料請關註WalkonNet其它相關文章!