提供了一个修改图片尺寸的imgs_resize.py函数，并修改了部分函数。

README.md

@@ -51,12 +51,13 @@ VOC拓展数据集的百度网盘如下：
 2、运行train.py进行训练，默认参数已经对应voc数据集所需要的参数了。  
 
 #### 二、训练自己的数据集
-1、本文使用VOC格式进行训练。  
-2、训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的SegmentationClass中。    
-3、训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。    
-4、在训练前利用voc_annotation.py文件生成对应的txt。    
-5、注意修改train.py的num_classes为分类个数+1。    
-6、运行train.py即可开始训练。  
+1、本文使用VOC格式进行训练。    
+2、如果需要修改图片尺寸请使用imgs_resize.py文件进行修改。    
+3、训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的SegmentationClass中。     
+4、训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。        
+5、在训练前利用voc_annotation.py文件生成对应的txt。        
+6、注意修改train.py的num_classes为分类个数+1。     
+7、运行train.py即可开始训练。  
 
 #### 三、训练医药数据集
 1、下载VGG的预训练权重到model_data下面。  

datasets/JPEGImages/README.MD

@@ -0,0 +1 @@
+

datasets/SegmentationClass/README.md

@@ -0,0 +1 @@
+#mask文件存放位置

datasets/before/README.md

@@ -0,0 +1,2 @@
+
+存放jpg和json

imgs_resize.py

@@ -0,0 +1,43 @@
+import os
+import argparse
+from tqdm import tqdm
+from PIL import Image
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--input_path', default='datasets/SegmentationClass',
+                        help='input path')
+    parser.add_argument('--output_path', default='datasets/ReSize_SegmentationClass',
+                        help='number of total epochs to run')
+    config = parser.parse_args()
+
+    return config
+
+"""程序功能是对图片进行裁剪，把图片多余的部分进行裁剪，留下含有数据的部分"""
+
+if __name__ == '__main__':
+
+    config = vars(parse_args())
+    imgs = os.listdir(config['input_path'])
+    if not os.path.exists(config['output_path']):
+        os.makedirs(config['output_path'])
+    for img in tqdm(imgs):
+        if img.endswith("jpg") or img.endswith("png"):
+            left = 645
+            top = 430
+            right = 645 + 1024
+            bottom = 430 + 1024
+            #left = math.floor(width/2) - 512
+            #right = math.floor(width/2) + 512
+            #top = math.floor(height/2) - 512
+            #bottom = math.floor(height/2) + 512
+            try:
+                im = Image.open(os.path.join(config['input_path'], img))
+                width, height = im.size
+                im1 = im.crop([left, top, right, bottom])
+                cut_name = os.path.join(config['output_path'], img)
+                im1.save(cut_name)
+            except RuntimeError as e:
+                print(e)
+    print("转换完成!!")

json_to_dataset.py

@@ -19,9 +19,9 @@
 if __name__ == '__main__':
     jpgs_path   = "datasets/JPEGImages"
     pngs_path   = "datasets/SegmentationClass"
-    classes     = ["_background_","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
+    # classes     = ["_background_","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
     # classes     = ["_background_","cat","dog"]
-
+    classes     = ["_background_","quesun","youwu","huahen","maocao"] 
     count = os.listdir("./datasets/before/") 
     for i in range(0, len(count)):
         path = os.path.join("./datasets/before", count[i])

predict.py

@@ -31,8 +31,9 @@
     #   count、name_classes仅在mode='predict'时有效
     #-------------------------------------------------------------------------#
     count           = False
-    name_classes    = ["background","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
+    # name_classes    = ["background","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
     # name_classes    = ["background","cat","dog"]
+    name_classes     = ["_background_","quesun","youwu","huahen","maocao"]
     #----------------------------------------------------------------------------------------------------------#
     #   video_path          用于指定视频的路径，当video_path=0时表示检测摄像头
     #                       想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
@@ -90,14 +91,22 @@
         '''
         while True:
             img = input('Input image filename:')
+            img_name = img.split("/")[-1]
             try:
                 image = Image.open(img)
             except:
                 print('Open Error! Try again!')
                 continue
             else:
                 r_image = unet.detect_image(image, count=count, name_classes=name_classes)
-                r_image.show()
+                #r_image.show()
+                plt.figure(figsize=(24.48, 20.48))  # 设置窗口大小
+                plt.suptitle('predict result')  # 图片名称
+                plt.subplot(1, 2, 1), plt.title('Source: ' + img_name)
+                plt.imshow(image), plt.axis('off')
+                plt.subplot(1, 2, 2), plt.title('Result: ' + img_name)
+                plt.imshow(r_image), plt.axis('off')
+                plt.show()
 
     elif mode == "video":
         capture=cv2.VideoCapture(video_path)

requirements.txt

@@ -1,9 +1,10 @@
-scipy==1.2.1
-numpy==1.17.0
-matplotlib==3.1.2
-opencv_python==4.1.2.30
-torch==1.2.0
-torchvision==0.4.0
-tqdm==4.60.0
-Pillow==8.2.0
-h5py==2.10.0
+scipy
+numpy
+matplotlib
+opencv_python
+torch
+torchvision
+tqdm
+Pillow
+h5py
+

train.py

@@ -66,7 +66,7 @@
     #   num_classes     训练自己的数据集必须要修改的
     #                   自己需要的分类个数+1，如2+1
     #-----------------------------------------------------#
-    num_classes = 21
+    num_classes = 5
     #-----------------------------------------------------#
     #   主干网络选择
     #   vgg
@@ -102,7 +102,7 @@
     #-----------------------------------------------------#
     #   input_shape     输入图片的大小，32的倍数
     #-----------------------------------------------------#
-    input_shape = [512, 512]
+    input_shape = [1024, 1024]
 
     #----------------------------------------------------------------------------------------------------------------------------#
     #   训练分为两个阶段，分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
@@ -219,7 +219,7 @@
     #   种类多（十几类）时，如果batch_size比较大（10以上），那么设置为True
     #   种类多（十几类）时，如果batch_size比较小（10以下），那么设置为False
     #------------------------------------------------------------------#
-    dice_loss       = False
+    dice_loss       = True
     #------------------------------------------------------------------#
     #   是否使用focal loss来防止正负样本不平衡
     #------------------------------------------------------------------#
@@ -238,7 +238,7 @@
     #                   keras里开启多线程有些时候速度反而慢了许多
     #                   在IO为瓶颈的时候再开启多线程，即GPU运算速度远大于读取图片的速度。
     #------------------------------------------------------------------#
-    num_workers     = 4
+    num_workers     = 0
 
     #------------------------------------------------------#
     #   设置用到的显卡
@@ -471,9 +471,9 @@
                 if distributed:
                     batch_size = batch_size // ngpus_per_node
 
-                gen             = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
+                gen             = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=False,
                                             drop_last = True, collate_fn = unet_dataset_collate, sampler=train_sampler)
-                gen_val         = DataLoader(val_dataset  , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True, 
+                gen_val         = DataLoader(val_dataset  , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=False, 
                                             drop_last = True, collate_fn = unet_dataset_collate, sampler=val_sampler)
 
                 UnFreeze_flag = True

utils/callbacks.py

@@ -104,7 +104,7 @@ def __init__(self, net, input_shape, num_classes, image_ids, dataset_path, log_d
         self.eval_flag          = eval_flag
         self.period             = period
 
-        self.image_ids          = [image_id.split()[0] for image_id in image_ids]
+        self.image_ids          = [image_id.split('\n')[0] for image_id in image_ids]
         self.mious      = [0]
         self.epoches    = [0]
         if self.eval_flag:

utils/dataloader.py

@@ -24,7 +24,7 @@ def __len__(self):
 
     def __getitem__(self, index):
         annotation_line = self.annotation_lines[index]
-        name            = annotation_line.split()[0]
+        name            = annotation_line.split("\n")[0]
 
         #-------------------------------#
         #   从文件中读取图像

voc_annotation.py

@@ -64,7 +64,7 @@
 
     print("Check datasets format, this may take a while.")
     print("检查数据集格式是否符合要求，这可能需要一段时间。")
-    classes_nums        = np.zeros([256], np.int)
+    classes_nums        = np.zeros([256], np.int_)
     for i in tqdm(list):
         name            = total_seg[i]
         png_file_name   = os.path.join(segfilepath, name)
@@ -95,4 +95,4 @@
 
     print("JPEGImages中的图片应当为.jpg文件、SegmentationClass中的图片应当为.png文件。")
     print("如果格式有误，参考:")
-    print("https://github.com/bubbliiiing/segmentation-format-fix")
+    print("https://github.com/bubbliiiing/segmentation-format-fix")