Update app.py

app.py

@@ -1,7 +1,523 @@
 import gradio as gr
+import torch
+import boto3
+import requests
+import re
+from bs4 import BeautifulSoup
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+import os
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import hashlib
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+#First output the thinking process in <think> </think> tags and then output the final answer in <answer> </answer> tags.
+#Answer the following question based on the information above and the given image, and provide citations for your response.
+
+
+# Cache for storing recognition results
+recognition_cache = {}
+
+# Function to calculate image hash for caching
+def calculate_image_hash(image):
+    """Calculate a hash for the image to use as a cache key"""
+    if image is None:
+        return None
+    img_array = np.array(image)
+    # Use a simple hash of the image data
+    return hashlib.md5(img_array.tobytes()).hexdigest()
+
+# Helper function to remove numbers from brackets
+def remove_bracketed_nums(input_string):
+    """
+    Removes all occurrences of [num] from the input string.
+    """
+    return re.sub(r'\[\d+\]', '', input_string)
+
+# Function to get Wikipedia links
+def fetch_wikipedia_links(entity_id):
+    url = f"https://www.wikidata.org/wiki/Special:EntityData/{entity_id}.json"
+    links = {}
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        data = response.json()
+        entity_data = data.get("entities", {}).get(entity_id, {})
+        if not entity_data:
+            print("No data found for the given entity ID.")
+            return links
+
+        sitelinks = entity_data.get("sitelinks", {})
+        if not sitelinks:
+            print("No sitelinks found for the given entity ID.")
+            return links
+
+        wikipedia_links = {}
+        for site, site_info in sitelinks.items():
+            if site.endswith("wiki"):
+                language = site.replace("wiki", "")
+                url = site_info.get("url")
+                wikipedia_links[language] = url
+
+        for lang, link in wikipedia_links.items():
+            links[lang] = link
+
+    except requests.RequestException as e:
+        print(f"Error fetching data: {e}")
+    return links
+
+# Function to get Wikipedia information
+def fetch_wikipedia_info(url):
+    try:
+        # Send HTTP request
+        response = requests.get(url)
+        response.raise_for_status()
+
+        # Parse HTML content using BeautifulSoup
+        soup = BeautifulSoup(response.content, 'html.parser')
+        
+        # Try multiple methods to find the title
+        title = None
+        
+        # Method 1: Find title in the <title> tag of the page
+        if not title:
+            title_tag = soup.find('title')
+            if title_tag:
+                title = title_tag.text.replace(' - Wikipedia', '').strip()
+        
+        # Method 2: Find title in h1 tag
+        if not title:
+            title_tag = soup.find('h1', {'id': 'firstHeading'})
+            if title_tag:
+                title = title_tag.text.strip()
+        
+        # Method 3: If previous methods fail, use a more flexible method
+        if not title:
+            title_tag = soup.find('h1')
+            if title_tag:
+                title = title_tag.text.strip()
+        
+        # Extract sections and their content
+        sections = {}
+        current_section = None
+        paragraphs = soup.find_all('p')
+        
+        # Traverse all headings and paragraphs, mapping sections to content
+        for element in soup.find_all(['h2', 'p']):
+            if element.name == 'h2':
+                # New section starts, get section title
+                section_title = element.text.strip()
+                current_section = section_title
+                sections[current_section] = []
+            elif element.name == 'p' and current_section:
+                # Add paragraph to current section
+                paragraph_text = element.text.strip()
+                if paragraph_text:
+                    sections[current_section].append(paragraph_text)
+        
+        # Convert section content to format where each section is connected into a single string
+        sections = {section: '\n'.join(content) for section, content in sections.items()}
+        
+        return {
+            'title': title or 'No title found',
+            'sections': sections
+        }
+    
+    except requests.RequestException as e:
+        return {'error': f'Request failed: {e}'}
+    except Exception as e:
+        return {'error': f'An error occurred: {e}'}
+
+# Use AWS Rekognition to recognize celebrities in the image
+def recognize_celebrities(image_path, confidence_threshold=90):
+    client = boto3.client(
+        "rekognition",
+        aws_access_key_id=,
+        aws_secret_access_key=,
+        region_name='us-east-1'
+    )
+    
+    with open(image_path, 'rb') as image:
+        response = client.recognize_celebrities(Image={'Bytes': image.read()})
+    
+    # Filter out celebrities with None or low confidence
+    filtered_celebs = [cele for cele in response.get('CelebrityFaces', []) 
+                      if cele.get('MatchConfidence') is not None and cele.get('MatchConfidence') > confidence_threshold]
+    
+    names = [cele['Name'] for cele in filtered_celebs]
+    bounding_boxes = [cele['Face']['BoundingBox'] for cele in filtered_celebs]
+    wikidatas = [cele['Urls'][0] for cele in filtered_celebs]
+    
+    return names, bounding_boxes, wikidatas
+
+# Draw bounding boxes on the image
+def draw_bounding_boxes(image_path, bounding_boxes, names):
+    img = Image.open(image_path)
+    draw = ImageDraw.Draw(img)
+    
+    width, height = img.size
+    
+    for i, bbox in enumerate(bounding_boxes):
+        left = int(bbox['Left'] * width)
+        top = int(bbox['Top'] * height)
+        right = int((bbox['Left'] + bbox['Width']) * width)
+        bottom = int((bbox['Top'] + bbox['Height']) * height)
+        
+        # Draw rectangle
+        draw.rectangle([(left, top), (right, bottom)], outline="red", width=3)
+        
+        # Add name label
+        text = f"[{i+1}]: {names[i]}"
+        font = ImageFont.truetype("arial.ttf", 20)  # Adjust font and size as needed
+        text_bbox = draw.textbbox((left, top - 20), text, font=font)
+        draw.rectangle(text_bbox, fill="white")
+        draw.text((left, top - 20), text, fill="red", font=font)
+    return np.array(img)
+
+# Resize the image to a smaller size for display
+def resize_image_for_display(img_array, max_width=500):
+    img = Image.fromarray(img_array)
+    width, height = img.size
+    if width > max_width:
+        ratio = max_width / width
+        new_height = int(height * ratio)
+        img = img.resize((max_width, new_height), Image.LANCZOS)
+    return np.array(img)
+
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    "weihongliang/RC-Qwen2VL-2b", 
+    torch_dtype=torch.float16, 
+    device_map="auto"
+)
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+
+# Use Qwen model for Q&A
+def qwen_qa(image_path, question, names, bounding_boxes, en_wiki_pedia_links, en_intros):
+    # Prepare prompt text
+    prompt = '\n'.join([
+        f"[{i+1}]: The information of the person located at <|box_start|>({int(bbox['Left']*1000)},{int(bbox['Top']*1000)}),"
+        f"({int((bbox['Left']+bbox['Width'])*1000)},{int((bbox['Top']+bbox['Height'])*1000)})<|box_end|> in the image: {intro}"
+        for i, (name, intro, bbox) in enumerate(zip(
+            names, 
+            [remove_bracketed_nums(intro).replace('\n','') for intro in en_intros], 
+            bounding_boxes
+        ))
+    ])
+    print(prompt)
+    # Prepare messages
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "image": image_path,
+                },
+                {"type": "text", "text": f"{prompt}\nAnswer the following question based on the information above and the given image, and provide citations for your response.\n{question}"},
+            ],
+        }
+    ]
+    
+    # Prepare inference
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=None,
+        padding=True,
+        return_tensors="pt",
+    )
+    
+    inputs = inputs.to("cuda")
+    
+    # Inference: generate output
+    generated_ids = model.generate(**inputs, max_new_tokens=4096)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=False, clean_up_tokenization_spaces=False
+    )[0]
+    print(output_text)
+    return output_text, en_wiki_pedia_links, prompt
+
+# Compare image with examples by filename or content when available
+def is_example_image(image, examples):
+    if isinstance(image, str) and os.path.exists(image):
+        # If we have a path, try to match by filename
+        image_name = os.path.basename(image)
+        for example in examples:
+            example_name = os.path.basename(example[0])
+            if image_name == example_name:
+                print(f"Matched example by filename: {image_name}")
+                return True, example
+    
+    # If we couldn't match by filename or if image is an object, return False
+    return False, None
+
+# Main processing function
+def process_image(image, question, confidence_threshold, examples=None):
+    if image is None:
+        return None, "Please upload an image", []
+    
+    # Calculate image hash for caching
+    image_hash = calculate_image_hash(image)
+    
+    # Check if we already have cached recognition results for this image
+    cache_key = f"{image_hash}_{confidence_threshold}"
+    if cache_key in recognition_cache:
+        print("Using cached recognition results")
+        names, bounding_boxes, en_wiki_pedia_links, en_intros, result_image = recognition_cache[cache_key]
+    else:
+        print("Performing new recognition")
+        # Save uploaded image
+        temp_image_path = "temp_image.jpg"
+        image.save(temp_image_path)
+        
+        # Check if this image matches one of our example images
+        # First, get original filename from gradio image object if available
+        original_filename = None
+        if hasattr(image, 'name'):
+            original_filename = image.name
+        
+        # For each example, check if the file exists and compare
+        example_match = None
+        if examples:
+            for example in examples:
+                example_path = example[0]
+                if os.path.exists(example_path):
+                    example_img = Image.open(example_path)
+                    # Compare images by size first (quick check)
+                    if example_img.size == image.size:
+                        # More thorough comparison - hash the pixels
+                        example_hash = calculate_image_hash(example_img)
+                        if example_hash == image_hash:
+                            example_match = example
+                            break
+                    example_img.close()
+        
+        if example_match:
+            print("Using example data instead of recognition")
+            names = example_match[2]
+            bounding_boxes = example_match[3]
+            wikidatas = example_match[4]
+            print(f"Example data: {names}, {bounding_boxes}, {wikidatas}")
+        else:
+            # Recognize celebrities with adjusted confidence threshold
+            names, bounding_boxes, wikidatas = recognize_celebrities(temp_image_path, confidence_threshold)
+        
+        print(f'识别_{names}_{bounding_boxes}_{wikidatas}')
+        en_wiki_pedia_links = []
+        en_intros = []
+        
+        for wikidata_url in wikidatas:
+            entity_id = wikidata_url.split('/')[-1]
+            links = fetch_wikipedia_links(entity_id)
+            
+            if 'en' in links:
+                en_link = links['en']
+                en_wiki_pedia_links.append(en_link)
+                
+                wiki_info = fetch_wikipedia_info(en_link)
+                # Try to get 'Contents' section, if it doesn't exist, try to get the first section
+                intro = wiki_info['sections'].get('Contents', '')
+                if not intro and wiki_info['sections']:
+                    intro = next(iter(wiki_info['sections'].values()), '')
+                
+                en_intros.append(intro)
+            else:
+                en_wiki_pedia_links.append("No English Wikipedia link found")
+                en_intros.append("No information available")
+        
+        if not names:
+            if os.path.exists(temp_image_path):
+                os.remove(temp_image_path)
+            return None, f"No celebrities could be recognized with confidence threshold {confidence_threshold}%", []
+        
+        indices = list(range(len(en_intros)))
+        
+        # Sort indices by the length of en_intros elements
+        indices.sort(key=lambda i: len(en_intros[i]))
+        
+        # Use sorted indices to rearrange the three lists
+        names = [names[i] for i in indices]
+        bounding_boxes = [bounding_boxes[i] for i in indices]
+        en_intros = [en_intros[i] for i in indices]
+        en_wiki_pedia_links = [en_wiki_pedia_links[i] for i in indices]
+        
+        # Draw bounding boxes on the image
+        result_image = draw_bounding_boxes(temp_image_path, bounding_boxes, names)
+        
+        # Cache the recognition results
+        recognition_cache[cache_key] = (names, bounding_boxes, en_wiki_pedia_links, en_intros, result_image)
+        
+        # Clean up temporary files
+        if os.path.exists(temp_image_path):
+            os.remove(temp_image_path)
+    
+    # Save image for qwen_qa (it requires a file path)
+    temp_image_path = "temp_image_for_qa.jpg"
+    Image.fromarray(result_image if isinstance(image, np.ndarray) else np.array(image)).save(temp_image_path)
+    
+    # Resize the result image for display
+    resized_image = resize_image_for_display(result_image, max_width=500)
+    
+    # Use Qwen model for Q&A
+    answer, wiki_links, prompt = qwen_qa(temp_image_path, question, names, bounding_boxes, en_wiki_pedia_links, en_intros)
+    
+    # Prepare recognized person information for frontend
+    people_info = [[prompt]]
+    [people_info.append([f"{name}: {link}"]) for name, link in zip(names, en_wiki_pedia_links)]
+   
+    # Clean up temporary files
+    if os.path.exists(temp_image_path):
+        os.remove(temp_image_path)
+    
+    return resized_image, answer, people_info
+
+# Create Gradio interface
+def create_interface():
+    # Define examples - with proper wikidata URLs
+    examples = [
+        ["./musk_and_huang.jpg", "What is Jensen doing?", 
+         ['Jen-Hsun Huang', 'Elon Musk'], 
+         [{'Width': 0.14037726819515228, 'Height': 0.20855841040611267, 'Left': 0.20549052953720093, 'Top': 0.05287889018654823}, 
+          {'Width': 0.07806053012609482, 'Height': 0.12407417595386505, 'Left': 0.7048892974853516, 'Top': 0.12597641348838806}], 
+         ['https://www.wikidata.org/wiki/Q305177', 'https://www.wikidata.org/wiki/Q317521']],
+        ["./dtms.jpg", "Describe the image.", 
+         ['Claude Shannon', 'John McCarthy', 'Marvin Minsky'], 
+         [{'Width': 0.06745696812868118, 'Height': 0.1393456608057022, 'Left': 0.8298415541648865, 'Top': 0.2935926616191864}, 
+          {'Width': 0.058158036321401596, 'Height': 0.10835719853639603, 'Left': 0.7029165029525757, 'Top': 0.16815270483493805}, 
+          {'Width': 0.05310939624905586, 'Height': 0.1121278703212738, 'Left': 0.4884793162345886, 'Top': 0.19025743007659912}], 
+         ['https://www.wikidata.org/wiki/Q92760', 'https://www.wikidata.org/wiki/Q92739', 'https://www.wikidata.org/wiki/Q204815']],
+        ["./hinton.jpeg", "Describe the image in detail.", 
+         ['Yoshua Bengio', 'Geoffrey Hinton', 'Yann LeCun'], 
+         [{'Width': 0.077728271484375, 'Height': 0.16116079688072205, 'Left': 0.8799420595169067, 'Top': 0.4856656789779663}, 
+          {'Width': 0.07422236353158951, 'Height': 0.15943190455436707, 'Left': 0.4633428454399109, 'Top': 0.07901764661073685}, 
+          {'Width': 0.07562466710805893, 'Height': 0.13936467468738556, 'Left': 0.025178398936986923, 'Top': 0.4953641891479492}], 
+         ['https://www.wikidata.org/wiki/Q3572699', 'https://www.wikidata.org/wiki/Q92894', 'https://www.wikidata.org/wiki/Q3571662']],
+        ["./clinton.jpg", "Who are the people in the picture, and what is the relationship between them?", 
+         ['Bill Clinton', 'Monica Lewinsky'], 
+         [{'Width': 0.07620880007743835, 'Height': 0.16198107600212097, 'Left': 0.5074607729911804, 'Top': 0.14220821857452393}, 
+          {'Width': 0.0722670778632164, 'Height': 0.1512720286846161, 'Left': 0.3914872407913208, 'Top': 0.24376636743545532}], 
+         ['https://www.wikidata.org/wiki/Q1124', 'https://www.wikidata.org/wiki/Q212659']],
+        ["./epst.jpeg", "Provide image description.", 
+         ['Lisa Randall', 'Kip S. Thorne', 'David Gross', 'Stephen Hawking', 'Brenda Chapman'], 
+         [{'Width': 0.09916354715824127, 'Height': 0.166521355509758, 'Left': 0.7962431311607361, 'Top': 0.4121580123901367}, 
+          {'Width': 0.07940348237752914, 'Height': 0.1626593917608261, 'Left': 0.6891748905181885, 'Top': 0.33117005228996277}, 
+          {'Width': 0.06350294500589371, 'Height': 0.12757645547389984, 'Left': 0.544218122959137, 'Top': 0.1575603038072586}, 
+          {'Width': 0.06830617040395737, 'Height': 0.1128319799900055, 'Left': 0.2937725782394409, 'Top': 0.30404558777809143}, 
+          {'Width': 0.03966952860355377, 'Height': 0.11658532917499542, 'Left': 0.18093101680278778, 'Top': 0.31299835443496704}], 
+         ['https://www.wikidata.org/wiki/Q450404', 'https://www.wikidata.org/wiki/Q323320', 'https://www.wikidata.org/wiki/Q40262', 'https://www.wikidata.org/wiki/Q17714', 'https://www.wikidata.org/wiki/Q429715']]
+    ]
+    
+    # Filter examples to only include files that exist
+    existing_examples = []
+    for example in examples:
+        if os.path.exists(example[0]):
+            existing_examples.append(example)
+    
+    with gr.Blocks(title="Celebrity Recognition and Q&A System") as app:
+        gr.Markdown("<div style='text-align: center;'><h1 style=' font-size: 28px; '>Celebrity Recognition and Q&A System</h1></div>")
+        gr.Markdown("**RC-MLLM** model is developed based on the Qwen2-VL model through a novel method called **RCVIT (Region-level Context-aware Visual Instruction Tuning)**, using the specially constructed **RCMU dataset** for training. Its core feature is the capability for **Region-level Context-aware Multimodal Understanding (RCMU)**. This means it can simultaneously understand both the visual content of specific regions/objects within an image and their associated textual information (utilizing bounding boxes coordinates), allowing it to respond to user instructions in a more context-aware manner. Simply put, RC-MLLM not only understands images but can also integrate the textual information linked to specific objects within the image for understanding. It achieves outstanding performance on RCMU tasks and is suitable for applications like personalized conversation.")
+
+        markdown_content = """
+            📑 [Arxiv](https://arxiv.org/abs/your-paper-id) | 
+            🤗 [Checkpoint]() | 
+            📁 [Dataset](https://huggingface.co/your-model-name) |
+            [Github](https://github.com/your-username/your-repo) | 
+            🚀 [Personalized Conversation](https://your-project-url.com)
+            """
+        gr.Markdown(markdown_content)
+        gr.Markdown("📌 Upload an image containing celebrities, the system will recognize them and provide Wikipedia-based Q&A using the RC-Qwen2-VL model.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(type="pil", label="Upload Image")
+                question_input = gr.Textbox(label="Question", placeholder="Enter your question...")
+                confidence_slider = gr.Slider(
+                    minimum=50, 
+                    maximum=100, 
+                    value=90, 
+                    step=1, 
+                    label="Confidence Threshold (%)", 
+                    info="Adjust the minimum confidence level for celebrity recognition"
+                )
+                submit_button = gr.Button("Ask RC-Qwen2-VL")
+                
+                # Add examples section
+                if existing_examples:
+                    gr.Examples(
+                        examples=[[example[0], example[1]] for example in existing_examples],
+                        inputs=[image_input, question_input],
+                        label="Example Images with Questions",
+                        fn=None
+                    )
+            with gr.Column(scale=1):
+                image_output = gr.Image(label="Recognition Result")
+                answer_output = gr.Textbox(label="RC-Qwen2-VL Answer")
+
+                people_info = gr.Dataframe(
+                    headers=["Recognized Person Information and Wikipedia Links"],
+                    datatype=["str"],
+                    label="Person Information",
+                    wrap=True,  # Enable text wrapping
+                )
+        
+        # Function to clear the cache
+        def clear_cache():
+            global recognition_cache
+            old_count = len(recognition_cache)
+            recognition_cache = {}
+            return f"Cache cleared. {old_count} entries were removed."
+        
+        # Handle image upload - clear cache when new image is uploaded
+        def on_image_upload(image):
+            if image is not None:
+                clear_cache()
+                return f"New image detected. Cache has been cleared."
+            return f"No image uploaded. Cache status unchanged."
+        
+        # Modified submit function to handle examples without the current_example parameter
+        def submit_fn(image, question, confidence_threshold):
+            # Process image using the recognition system
+            return process_image(image, question, confidence_threshold, examples=examples)
+        
+        # Connect buttons to functions
+        submit_button.click(
+            fn=submit_fn,
+            inputs=[image_input, question_input, confidence_slider],
+            outputs=[image_output, answer_output, people_info]
+        )
+        
+        # Clear cache automatically when new image is uploaded
+        image_input.change(
+            fn=on_image_upload,
+            inputs=[image_input],
+        )
+        
+        gr.Markdown("## Instructions")
+        gr.Markdown("""
+        1. Upload an image containing celebrities
+        2. Enter your question, for example:
+           - "Who are the people in the image?"
+           - "What achievements does the person on the left have?"
+           - "What is the relationship between these people?"
+        3. Adjust the confidence threshold slider if needed (lower values will recognize more faces but might be less accurate)
+        4. Click the submit button to get the answer
+        5. Or try one of the examples below
+        6. The system caches recognition results for each image and confidence threshold combination
+        7. Cache is automatically cleared when you upload a new image
+        """)
+    
+    return app
+
+# Launch the application
+if __name__ == "__main__":
+    app = create_interface()
+    app.launch(share=True)