Update app.py

app.py

@@ -1,22 +1,18 @@
-# app.py (Enhanced for Medical Models & Image Analysis in PDFs)
+# app.py (Enhanced for Medical Models, Image Analysis, & OSINT)
 
 import os
 import re
 import json
 import torch
 import fitz  # PyMuPDF
-from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, DataCollatorForLanguageModeling, CLIPProcessor, CLIPModel, AutoModelForSequenceClassification, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, CLIPProcessor, CLIPModel, DataCollatorForLanguageModeling
 from sklearn.model_selection import train_test_split
 from io import BytesIO
 from PIL import Image
 import gradio as gr
 from torch.utils.data import Dataset, DataLoader
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained('microsoft/BioGPT-Large')
-model = AutoModelForCausalLM.from_pretrained('microsoft/BioGPT-Large')
-
+import pytesseract
+from PIL.ExifTags import TAGS
 
 
 # Custom Dataset Class for Better Handling
@@ -35,6 +31,22 @@ class ConversationDataset(Dataset):
         }
 
 
+def extract_metadata(image):
+    """
+    Extract metadata from an image file.
+    """
+    metadata = {}
+    try:
+        exif_data = image._getexif()
+        if exif_data:
+            for tag, value in exif_data.items():
+                tag_name = TAGS.get(tag, tag)
+                metadata[tag_name] = value
+    except Exception as e:
+        metadata['error'] = str(e)
+    return metadata
+
+
 def extract_text_and_images_from_pdf(pdf_stream):
     doc = fitz.open(stream=pdf_stream, filetype='pdf')
     text = ''
@@ -44,7 +56,6 @@ def extract_text_and_images_from_pdf(pdf_stream):
         page = doc.load_page(page_index)
         text += page.get_text() + '\n'
 
-        # Extracting Images
         for img_index, img in enumerate(page.get_images(full=True)):
             xref = img[0]
             base_image = doc.extract_image(xref)
@@ -56,72 +67,17 @@ def extract_text_and_images_from_pdf(pdf_stream):
     return text, images
 
 
-def load_and_preprocess_data(file):
-    content = file.read()
-    if file.name.endswith('.pdf'):
-        pdf_stream = BytesIO(content)
-        text, images = extract_text_and_images_from_pdf(pdf_stream)
-    else:
-        text = content.decode('utf-8')
-        images = []
-
-    lines = text.split('\n')
-    conversations = []
-    conversation = {'prompt': '', 'response': ''}
-    user_turn = True
-
-    for line in lines:
-        line = line.strip()
-        if line == '':
-            continue
-        if user_turn:
-            conversation['prompt'] += line + ' '
-        else:
-            conversation['response'] += line + ' '
-
-        user_turn = not user_turn
-
-        if not user_turn:
-            conversations.append(conversation)
-            conversation = {'prompt': '', 'response': ''}
-
-    return conversations, images
-
-
-def analyze_images_in_pdf(images):
-    descriptions = []
-
-    for image in images:
-        inputs = clip_processor(images=image, return_tensors="pt")
-        outputs = clip_model.get_image_features(**inputs)
-        descriptions.append(f"Image processed successfully. Feature vector shape: {outputs.shape}")
-
-    return descriptions
-
+def analyze_images(images):
+    results = {}
 
-def train_model(inputs, labels, model, tokenizer):
-    dataset = ConversationDataset(inputs, labels)
-    training_args = TrainingArguments(
-        output_dir='./results',
-        overwrite_output_dir=True,
-        num_train_epochs=1,
-        per_device_train_batch_size=2,
-        save_steps=500,
-        save_total_limit=2,
-        logging_dir='./logs',
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=dataset,
-        data_collator=data_collator,
-    )
-
-    trainer.train()
-
-    model.save_pretrained('./trained_model')
-    tokenizer.save_pretrained('./trained_model')
+    for index, image in enumerate(images):
+        metadata = extract_metadata(image)
+        ocr_text = pytesseract.image_to_string(image)
+        results[f"Image_{index+1}"] = {
+            'metadata': metadata,
+            'ocr_text': ocr_text
+        }
+    return results
 
 
 def generate_response(prompt, model, tokenizer, max_length=150):
@@ -133,47 +89,29 @@ def generate_response(prompt, model, tokenizer, max_length=150):
         return response
 
 
-def train_interface(file):
-    try:
-        conversations, images = load_and_preprocess_data(file)
-        inputs, labels = tokenize_data(conversations, tokenizer)
-        train_model(inputs, labels, model, tokenizer)
-
-        if images:
-            image_analysis = analyze_images_in_pdf(images)
-            return f"✅ Model training completed successfully. {len(images)} images processed."
-        else:
-            return "✅ Model training completed successfully. No images found."
-    except Exception as e:
-        return f"❌ Error during training: {str(e)}"
-
-
-def generate_interface(prompt):
-    try:
-        return generate_response(prompt, model, tokenizer)
-    except Exception as e:
-        return f"❌ Error during generation: {str(e)}"
-
-
-tokenizer = AutoTokenizer.from_pretrained('microsoft/BioGPT')
-model = AutoModelForSequenceClassification.from_pretrained('microsoft/BioGPT')
-data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+tokenizer = AutoTokenizer.from_pretrained('microsoft/BioGPT-Large')
+model = AutoModelForCausalLM.from_pretrained('microsoft/BioGPT-Large')
 
-clip_model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32')
-clip_processor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Medical LLM Model Training & PDF Image Analysis")
-
-    with gr.Tab("Train LLM"):
-        file_input = gr.File(label="Upload PDF or Text File")
-        train_output = gr.Textbox(label="Training Status")
-        file_input.change(train_interface, inputs=file_input, outputs=train_output)
+    gr.Markdown("# Medical LLM Model Training, PDF Image Analysis & OSINT")
 
     with gr.Tab("Generate Text"):
         prompt_input = gr.Textbox(label="Enter Medical Prompt")
         generate_output = gr.Textbox(label="Generated Medical Text")
-        prompt_input.change(generate_interface, inputs=prompt_input, outputs=generate_output)
+        prompt_input.change(generate_response, inputs=[prompt_input, model, tokenizer], outputs=generate_output)
+
+    with gr.Tab("PDF Analysis"):
+        pdf_file = gr.File(label="Upload PDF")
+        pdf_output = gr.Textbox(label="Extracted Text")
+        image_analysis_output = gr.Textbox(label="Image Metadata & OCR")
+
+        def process_pdf(file):
+            pdf_content, images = extract_text_and_images_from_pdf(BytesIO(file.read()))
+            image_analysis_results = analyze_images(images)
+            return pdf_content, json.dumps(image_analysis_results, indent=4)
+
+        pdf_file.upload(process_pdf, inputs=pdf_file, outputs=[pdf_output, image_analysis_output])
 
 
 demo.launch()