Add files via upload

get_dr_txt.py

+from ssd import SSD
+from PIL import Image
+from utils.box_utils import letterbox_image,ssd_correct_boxes
+from torch.autograd import Variable
+import torch
+import numpy as np
+import os
+MEANS = (104, 117, 123)
+class mAP_SSD(SSD):
+    #---------------------------------------------------#
+    #   检测图片
+    #---------------------------------------------------#
+    def detect_image(self,image_id,image):
+        self.confidence = 0.05
+        f = open("./input/detection-results/"+image_id+".txt","w") 
+        image_shape = np.array(np.shape(image)[0:2])
+
+        crop_img = np.array(letterbox_image(image, (self.model_image_size[0],self.model_image_size[1])))
+        photo = np.array(crop_img,dtype = np.float64)
+
+        # 图片预处理，归一化
+        photo = Variable(torch.from_numpy(np.expand_dims(np.transpose(crop_img-MEANS,(2,0,1)),0)).cuda().type(torch.FloatTensor))
+        preds = self.net(photo)
+        
+        top_conf = []
+        top_label = []
+        top_bboxes = []
+        for i in range(preds.size(1)):
+            j = 0
+            while preds[0, i, j, 0] >= self.confidence:
+                score = preds[0, i, j, 0]
+                label_name = self.class_names[i-1]
+                pt = (preds[0, i, j, 1:]).detach().numpy()
+                coords = [pt[0], pt[1], pt[2], pt[3]]
+                top_conf.append(score)
+                top_label.append(label_name)
+                top_bboxes.append(coords)
+                j = j + 1
+        # 将预测结果进行解码
+        if len(top_conf)<=0:
+            return image
+        top_conf = np.array(top_conf)
+        top_label = np.array(top_label)
+        top_bboxes = np.array(top_bboxes)
+        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1)
+
+        # 去掉灰条
+        boxes = ssd_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape)
+
+
+        for i, c in enumerate(top_label):
+            predicted_class = c
+            score = str(float(top_conf[i]))
+
+            top, left, bottom, right = boxes[i]
+            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
+
+        f.close()
+        return 
+
+ssd = mAP_SSD()
+image_ids = open('VOCdevkit/VOC2007/ImageSets/Main/test.txt').read().strip().split()
+
+if not os.path.exists("./input"):
+    os.makedirs("./input")
+if not os.path.exists("./input/detection-results"):
+    os.makedirs("./input/detection-results")
+if not os.path.exists("./input/images-optional"):
+    os.makedirs("./input/images-optional")
+
+
+for image_id in image_ids:
+    image_path = "./VOCdevkit/VOC2007/JPEGImages/"+image_id+".jpg"
+    image = Image.open(image_path)
+    image.save("./input/images-optional/"+image_id+".jpg")
+    ssd.detect_image(image_id,image)
+    print(image_id," done!")
+    
+
+print("Conversion completed!")
\ No newline at end of file

get_gt_txt.py

+import sys
+import os
+import glob
+import xml.etree.ElementTree as ET
+
+image_ids = open('VOCdevkit/VOC2007/ImageSets/Main/test.txt').read().strip().split()
+
+if not os.path.exists("./input"):
+    os.makedirs("./input")
+if not os.path.exists("./input/ground-truth"):
+    os.makedirs("./input/ground-truth")
+
+for image_id in image_ids:
+    with open("./input/ground-truth/"+image_id+".txt", "w") as new_f:
+        root = ET.parse("VOCdevkit/VOC2007/Annotations/"+image_id+".xml").getroot()
+        for obj in root.findall('object'):
+            obj_name = obj.find('name').text
+            bndbox = obj.find('bndbox')
+            left = bndbox.find('xmin').text
+            top = bndbox.find('ymin').text
+            right = bndbox.find('xmax').text
+            bottom = bndbox.find('ymax').text
+            new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))
+print("Conversion completed!")
\ No newline at end of file

video.py

+#-------------------------------------#
+#       调用摄像头检测
+#-------------------------------------#
+from ssd import SSD
+from PIL import Image
+import numpy as np
+import cv2
+ssd = SSD()
+# 调用摄像头
+capture=cv2.VideoCapture(0) # capture=cv2.VideoCapture("1.mp4")
+
+while(True):
+    # 读取某一帧
+    ref,frame=capture.read()
+    # 格式转变，BGRtoRGB
+    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
+    # 转变成Image
+    frame = Image.fromarray(np.uint8(frame))
+
+    # 进行检测
+    frame = np.array(ssd.detect_image(frame))
+
+    # RGBtoBGR满足opencv显示格式
+    frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
+    cv2.imshow("video",frame)
+    c= cv2.waitKey(30) & 0xff 
+    if c==27:
+        capture.release()
+        break
+
+yolo.close_session()
+    
\ No newline at end of file