AI Label

System_I.py

@@ -7,7 +7,7 @@ import json
 import random
 from configs.configs import model_head
 
-import os,sys
+import os,sys,shutil
 import torch
 from torch import nn
 from PIL import Image, ImageDraw, ImageFont
@@ -23,6 +23,7 @@ import numpy as np
 from fine.models.LoadModel import MainModel
 from fine.fineConfig import LoadConfig, load_data_transformers
 from torch.autograd import Variable
+from  creat_xml_fordetection import make_xml
 import threading
 from common import global_para as  gp
 import time,statistics
@@ -288,17 +289,14 @@ class Detection_S:
                 print("check final is ", (t5 - t4) * 1000)
         return flag,labels,boxes,scores
 class Detection_Y5:
-    def __init__(self,device,root):
+    def __init__(self,device,root,detect_thresh,DetWeight):
         self.device = device
         self.root = root
         self.img_size = 640
-        self.conf_thres = gp.detect_thresh
+        self.conf_thres = detect_thresh
         self.iou_thres = 0.4
         self.stride = 1
-        if gp.red_flag:
-            self.weights = '../data_MPKFNet/test/2_r.pt'
-        else:
-            self.weights = gp.DetWeight
+        self.weights = DetWeight
         # self.device = select_device(self.device)
         model = attempt_load(self.weights, map_location=self.device)
         model.to(self.device).eval()
@@ -320,6 +318,68 @@ class Detection_Y5:
             img = img.unsqueeze(0)
 
         return img0, img
+    def AIlabel(self,imagename):
+        with torch.no_grad():
+            image = cv2.imread(imagename)
+            image_or = image
+            img_detect = cv2.resize(image, (1024, 768))
+            t1 = time.time()
+            image_scale,image = self.preprocess(img_detect)
+            t2 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection data deal is ",(t2-t1)*1000)
+            pred = self.model(image, augment=False)[0]
+            pred = pred.float()
+            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
+            t3 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection inference is ",(t3-t2)*1000)
+            boxes = []
+            scores = []
+            labels = []
+            for det in pred:
+                if det is not None and len(det):
+                    det[:, :4] = scale_coords(
+                        image.shape[2:], det[:, :4], image_scale.shape).round()
+
+                    boxes.append(det[:, :4])
+                    scores.append(det[:, 4])
+                    labels.append((det[:, -1]))
+                    boxes = boxes[0]
+                    scores = scores[0]
+                    labels = labels[0]
+            t5 = time.time()
+            if gp.DebugPrint_flag:
+                print("Detection get result is ", (t5 - t3) * 1000)
+            if len(labels) > 0:
+                ratio_x = (image_or.shape[1] / 1024.0)
+                ratio_y = (image_or.shape[0] / 768.0)
+
+
+                for i in range(len(boxes)):
+                    boxes[i][0]=boxes[i][0] * ratio_x
+                    boxes[i][2]=boxes[i][2] * ratio_x
+                    boxes[i][1] = boxes[i][1] * ratio_y
+                    boxes[i][3] = boxes[i][3] * ratio_y
+                    if boxes[i][0]<0:
+                        boxes[i][0] = 0
+                    if boxes[i][1]<0:
+                        boxes[i][1] = 0
+                    if boxes[i][2]<0:
+                        boxes[i][2] = 0
+                    if boxes[i][3]<0:
+                        boxes[i][3] = 0
+                # for i in range(len(labels)):
+                #     labels[i] = "bird"
+                t4 = time.time()
+                if gp.DebugPrint_flag:
+                    print("Detection final is ", (t4 - t3) * 1000)
+                xml_outputpath = imagename.replace("imagesAILabel","LabelXML").replace(".jpg",".xml")
+                make_xml(imagename,image_or.shape[0],image_or.shape[1],boxes,labels,scores,xml_outputpath)
+                shutil.copy(imagename, imagename.replace("imagesAILabel", "YesResult"))
+            else:
+                shutil.copy(imagename,imagename.replace("imagesAILabel","NoResult"))
+
     def inference(self,image,imagetime,index,fine_grained):
         with torch.no_grad():
             rootpath = os.path.join(self.root, imagetime)
@@ -1578,7 +1638,7 @@ class Density_mul (threading.Thread):
         gp.destory_image = True
         if gp.DebugPrint_flag:
             print(self.name + self.imagetime," is over")
-class Detecntion_mul (threading.Thread):
+class Detection_mul (threading.Thread):
     def __init__(self, threadID, name,imagetime,DetectionClass,Fine_grainedClass):
         threading.Thread.__init__(self)
         self.threadID = threadID