Update get_dr_txt.py

get_dr_txt.py

@@ -17,11 +17,12 @@ class mAP_SSD(SSD):
 
         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)
-        
+        with torch.no_grad():
+            photo = Variable(torch.from_numpy(np.expand_dims(np.transpose(crop_img-MEANS,(2,0,1)),0)).type(torch.FloatTensor))
+            if self.cuda:
+                photo = photo.cuda()
+            preds = self.net(photo)
         top_conf = []
         top_label = []
         top_bboxes = []
@@ -77,4 +78,4 @@ for image_id in image_ids:
     print(image_id," done!")
     
 
-print("Conversion completed!")
+print("Conversion completed!")
\ No newline at end of file