Update app.py

app.py

@@ -2,10 +2,13 @@ import gradio as gr
 import onnxruntime as ort
 import numpy as np
 from PIL import Image
+from torchvision import transforms
+import torchvision.transforms.v2 as T
 import io
 import rdkit
 from rdkit import Chem
 from rdkit.Chem import Draw
+from postprocessor import RTDETRPostProcessor
 from utils import bbox_to_graph_with_charge, mol_from_graph_with_chiral
 
 bond_labels = [13,14,15,16,17]
@@ -13,20 +16,26 @@ idx_to_labels = {0:'other',1:'C',2:'O',3:'N',4:'Cl',5:'Br',6:'S',7:'F',8:'B',
             9:'I',10:'P',11:'*',12:'Si',13:'NONE',14:'BEGINWEDGE',15:'BEGINDASH',
             16:'=',17:'#',18:'-4',19:'-2',20:'-1',21:'1',22:'+2',} #NONE is single ?
 
-def preprocess_image(image):
-    """
-    预处理输入图片
-    """
-    # 将图片调整为模型所需的输入尺寸
-    image = image.resize((640, 640))  # 根据实际模型需求调整尺寸
-    # 转换为numpy数组并归一化
-    img_array = np.array(image)
-    img_array = img_array.astype(np.float32) / 255.0
-    # 添加批次维度
-    img_array = np.expand_dims(img_array, axis=0)
-    # 根据模型训练时的预处理方式进行调整
-    img_array = img_array.transpose(0, 3, 1, 2)  # BHWC to BCHW
-    return img_array
+
+def image_to_tensor(image_path):
+    # Open the image using PIL
+    image = Image.open(image_path)
+    w, h = image.size
+    # print("width: {}, height: {}".format(w, h))
+    # Define a transform to convert the image to a tensor and normalize it
+    transform = transforms.Compose([
+        # transforms.Grayscale(num_output_channels=1),  # Convert to grayscale (1 channel)
+        T.Resize((640, 640)),  # Resize the image to 224x224
+        T.ToImageTensor(),  # Convert to Tensor (C x H x W)
+        T.ConvertDtype(dtype=torch.float32)
+        # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Optional normalization for pretrained models
+    ])
+    
+    # Apply the transform to the image
+    tensor = transform(image)
+    
+    return tensor,w,h
+
 
 def visualize_molecule(smiles):
     """
@@ -50,17 +59,41 @@ def predict(input_image):
         session = ort.InferenceSession("model.onnx")  # 替换为实际模型路径
         
         # 预处理图片
-        processed_image = preprocess_image(input_image)
+         # Example usage: #change thie image
+        tensor,w,h = image_to_tensor(input_image)
+        processed_image=tensor.unsqueeze(0)
         
         # 获取模型输入输出名称
         input_name = session.get_inputs()[0].name
         output_name = session.get_outputs()[0].name
         
         # 进行推理
-        predictions = session.run([output_name], {input_name: processed_image})
+        outputs = session.run([output_name], {input_name: processed_image})
+        ori_size=torch.Tensor([w,h]).long().unsqueeze(0)
+        postprocessor = RTDETRPostProcessor()
+        result_ = postprocessor(outputs, ori_size)
+        score_=result_[0]['scores']
+        boxe_=result_[0]['boxes']
+        label_=result_[0]['labels']
+        selected_indices =score_ > 0.5
+        output={
+            'labels': label_[selected_indices],
+            'boxes': boxe_[selected_indices],
+            'scores': score_[selected_indices]
+        }
         
-        # 假设模型输出是SMILES字符串
-        output = predictions  # 根据实际模型输出格式调整
+        filtered_output_dict={image_path: output
+        }
+        
+        
+        x_center = (output["boxes"][:, 0] + output["boxes"][:, 2]) / 2
+        y_center = (output["boxes"][:, 1] + output["boxes"][:, 3]) / 2
+        center_coords = torch.stack((x_center, y_center), dim=1)
+        output = {'bbox':         output["boxes"].to("cpu").numpy(),
+                    'bbox_centers': center_coords.to("cpu").numpy(),
+                    'scores':       output["scores"].to("cpu").numpy(),
+                    'pred_classes': output["labels"].to("cpu").numpy()}
+
         atoms_df, bonds_list,charge_list =bbox_to_graph_with_charge(output, idx_to_labels=idx_to_labels,
                                                     bond_labels=bond_labels,  result=[])
         smiles,mol_rebuit=mol_from_graph_with_chiral(atoms_df, bonds_list,charge_list )