Update app.py

app.py

@@ -3,7 +3,6 @@ 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
@@ -17,24 +16,22 @@ idx_to_labels = {0:'other',1:'C',2:'O',3:'N',4:'Cl',5:'Br',6:'S',7:'F',8:'B',
             16:'=',17:'#',18:'-4',19:'-2',20:'-1',21:'1',22:'+2',} #NONE is single ?
 
 
-def image_to_tensor(image_path):
-    # Open the image using PIL
+def image_to_numpy(image_path):
+
     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)
+    img_array = np.array(image)
+    
+    img_resized = cv2.resize(img_array, (640, 640), interpolation=cv2.INTER_LINEAR)
+    
+    img_float = img_resized.astype(np.float32)
+    img_normalized = img_float / 255.0
+    
+    if len(img_normalized.shape) == 3:
+        img_normalized = img_normalized.transpose(2, 0, 1)
     
-    return tensor,w,h
+    return img_normalized, w, h
 
 
 def visualize_molecule(smiles):
@@ -60,18 +57,19 @@ def predict(input_image):
         
         # 预处理图片
          # Example usage: #change thie image
-        tensor,w,h = image_to_tensor(input_image)
-        processed_image=tensor.unsqueeze(0)
+        img_array,w,h = image_to_numpy(input_image)
+        processed_image=np.expand_dims(img_array, 0)
         
         # 获取模型输入输出名称
         input_name = session.get_inputs()[0].name
         output_name = session.get_outputs()[0].name
         
         # 进行推理
-        outputs = session.run([output_name], {input_name: processed_image})
+        outputs = session.run(None, {input_name: processed_image})
+        preds = {'pred_logits':torch.from_numpy(outputs[0]), 'pred_boxes':torch.from_numpy(outputs[1])}
         ori_size=torch.Tensor([w,h]).long().unsqueeze(0)
-        postprocessor = RTDETRPostProcessor()
-        result_ = postprocessor(outputs, ori_size)
+        postprocessor = RTDETRPostProcessor(num_classes=23, use_focal_loss=True)
+        result_ = postprocessor(preds, ori_size)
         score_=result_[0]['scores']
         boxe_=result_[0]['boxes']
         label_=result_[0]['labels']
@@ -94,9 +92,9 @@ def predict(input_image):
                     '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,
+        atoms_df, bonds_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 )
+        smiles, mol_rebuit = mol_from_graph_with_chiral(atoms_df, bonds_list)
         
         # 使用RDKit生成分子结构图
         mol_image = visualize_molecule(smiles)