chuodinity/orisense

1import streamlit as st2import torch3import torch.nn as nn4from PIL import Image5from transformers import (6    AutoTokenizer, AutoConfig, AutoModel, PreTrainedModel, 7    pipeline, ViTImageProcessor, ViTForImageClassification8)9 10# --- DESKLIB TEXT DETECTOR ARCHITECTURE ---11class DesklibAIDetectionModel(PreTrainedModel):12    config_class = AutoConfig13    _tied_weights_keys = {}14 15    def __init__(self, config):16        super().__init__(config)17        self.model = AutoModel.from_config(config)18        self.classifier = nn.Linear(config.hidden_size, 1)19        20        # NEW: Always call post_init at the end of __init__21        self._tied_weights_keys = {}22        if not hasattr(self, "_keys_to_ignore_on_save"):23            self._keys_to_ignore_on_save = []24            25        self.post_init()26 27    def forward(self, input_ids, attention_mask=None):28        outputs = self.model(input_ids, attention_mask=attention_mask)29        last_hidden_state = outputs[0]30        input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()31        mean_pooled = torch.sum(last_hidden_state * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)32        return self.classifier(mean_pooled)33 34# --- LOAD SPECIALIZED MODELS ---35@st.cache_resource36def load_assets():37    device = "cuda" if torch.cuda.is_available() else "cpu"38    39    # Text Model (Desklib)40    text_model_id = "desklib/ai-text-detector-v1.01"41    t_tokenizer = AutoTokenizer.from_pretrained(text_model_id)42    t_model = DesklibAIDetectionModel.from_pretrained(text_model_id).to(device)43    44    # Image Model (Specialized ViT for AIGC)45    img_model_id = "capcheck/ai-image-detection"46    img_pipe = pipeline("image-classification", model=img_model_id, device=0 if device == "cuda" else -1)47    48    return t_tokenizer, t_model, img_pipe, device49 50tokenizer, text_model, img_pipeline, device = load_assets()51 52# --- UI INTERFACE ---53st.set_page_config(page_title="AIGC Late Fusion Detector", layout="wide")54st.title("OriSense")55 56col_in, col_out = st.columns([1, 1])57 58with col_in:59    st.subheader("Input Content")60    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])61    user_text = st.text_area("Input Text", placeholder="Paste article or caption...", height=200)62    63    if uploaded_file:64        st.image(Image.open(uploaded_file), caption="Uploaded Image", use_container_width=True)65 66# --- PROCESSING ---67if st.button("Run Multi-Modal Detection") and uploaded_file and user_text:68    with st.spinner("Analyzing artifacts in text and pixels..."):69        # 1. Text Score (Logit -> Sigmoid)70        t_inputs = tokenizer(user_text, return_tensors="pt", truncation=True, padding=True).to(device)71        with torch.no_grad():72            t_logit = text_model(t_inputs['input_ids'], t_inputs['attention_mask'])73            p_text = torch.sigmoid(t_logit).item()74 75        # 2. Image Score (AIGC ViT)76        img_results = img_pipeline(Image.open(uploaded_file))77        # Find the score for 'FAKE' (AI-generated), case-insensitive, with safe fallback78        p_image = next((item['score'] for item in img_results if item['label'].upper() == 'FAKE'), 0.0)79 80        # 3. Late Fusion (Weighted Average)81        # Using 0.5/0.5 for balanced multimodal detection82        fused_score = (0.5 * p_text) + (0.5 * p_image)83 84        with col_out:85            st.subheader("System Verdict")86            87            # Classification logic88            verdict = "AI-GENERATED" if fused_score > 0.5 else "HUMAN-ORIGIN"89            color = "red" if verdict == "AI-GENERATED" else "green"90            91            st.markdown(f"### Result: :{color}[{verdict}]")92            st.metric("Aggregate Confidence", f"{fused_score:.2%}")93            94            # Visual Breakdown95            st.write("**Modality Breakdown:**")96            st.progress(p_text, text=f"Text AI Probability: {p_text:.1%}")97            st.progress(p_image, text=f"Image AI Probability: {p_image:.1%}")98            99            # Brief Forensic Note100            if fused_score > 0.5:101                st.warning("Conclusion: High cross-modal artifact detection. The content shows patterns consistent with synthetic generation.")102            else:103                st.success("Conclusion: Low probability of AI generation. Features align with natural human patterns.")104