import gradio as gr
import cv2
from PIL import Image
from torchvision import transforms
import io
import rdkit
import numpy as np
import torch
import onnxruntime as ort
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]
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 image_to_numpy(image_path):

    image = Image.open(image_path)
    w, h = image.size
    
    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 img_normalized, w, h


def visualize_molecule(smiles):
    """
    使用RDKit将SMILES转换为分子结构图
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        img = Draw.MolToImage(mol)
        return img
    except:
        return None

def predict(input_image):
    """
    主要的推理函数
    """
    try:
        # 加载和初始化ONNX模型
        session = ort.InferenceSession("model.onnx")  # 替换为实际模型路径
        
        # 预处理图片
         # Example usage: #change thie image
        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(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(num_classes=23, use_focal_loss=True)
        result_ = postprocessor(preds, 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]
        }
        
        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 = 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)
        
        # 使用RDKit生成分子结构图
        mol_image = visualize_molecule(smiles)
        
        if mol_image is None:
            return "无效的SMILES字符串", None
        
        return smiles, mol_image
        
    except Exception as e:
        return f"发生错误: {str(e)}", None

# 创建Gradio界面
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Text(label="SMILES字符串"),
        gr.Image(label="分子结构图")
    ],
    title="化学结构OCR",
    description="上传一张包含化学结构的图片，获取对应的SMILES表示和分子结构图。",
    examples=[
        ["example.jpg"]
    ]
)

# 启动应用
if __name__ == "__main__":
    iface.launch()