app.py

import gradio as gr
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import cv2
import numpy as np
import pandas as pd
import ast
from collections import Counter
from io import BytesIO
from io import StringIO


# # Ensure GPU usage if available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize the model and tokenizer
model = AutoModelForCausalLM.from_pretrained("ManishThota/CustomModel",
                                             torch_dtype=torch.float16, 
                                             device_map="auto",
                                             trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("ManishThota/CustomModel", trust_remote_code=True)




def video_to_frames(video, fps=1):
    """Converts a video file into frames and stores them as PNG images in a list."""
    frames_png = []
    cap = cv2.VideoCapture(video)
    
    if not cap.isOpened():
        print("Error opening video file")
        return frames_png
    
    frame_count = 0
    frame_interval = int(cap.get(cv2.CAP_PROP_FPS)) // fps  # Calculate frame interval
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print("Can't receive frame (stream end?). Exiting ...")
            break
        
        if frame_count % frame_interval == 0:
            is_success, buffer = cv2.imencode(".png", frame)
            if is_success:
                frames_png.append(np.array(buffer).tobytes())
        
        frame_count += 1
    
    cap.release()
    return frames_png

def extract_frames(frame):

    # Convert binary data to a numpy array
    frame_np = np.frombuffer(frame, dtype=np.uint8)

    # Decode the PNG image
    image_rgb = cv2.imdecode(frame_np, flags=cv2.IMREAD_COLOR)  # Assuming it's in RGB format

    # Convert RGB to BGR
    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    return image_bgr

def predict_answer(video, image, question):

    text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{question}? ASSISTANT:"
    input_ids = tokenizer(text, return_tensors='pt').input_ids.to(device)

    
    if image:
        # Process as an image
        image = image.convert("RGB")
        image_tensor = model.image_preprocess(image)
        
        #Generate the answer
        output_ids = model.generate(
            input_ids,
            max_new_tokens=25,
            images=image_tensor,
            use_cache=True)[0]
        
        return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
        
    elif video:
        frames = video_to_frames(video)
        image = extract_frames(frames[2])
        image_tensor = model.image_preprocess([image])
        # Generate the answer
        output_ids = model.generate(
                input_ids,
                max_new_tokens=25,
                images=image_tensor,
                use_cache=True)[0]
        answer = tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()

        return answer
        # # Process as a video
        # frames = video_to_frames(video)
        # answers = []
        # for frame in frames:
        #     image = extract_frames(frame)
        #     image_tensor = model.image_preprocess([image])
            
        #     # Generate the answer
        #     output_ids = model.generate(
        #         input_ids,
        #         max_new_tokens=25,
        #         images=image_tensor,
        #         use_cache=True)[0]
            
        #     answer = tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
        #     answers.append(answer)

        # # Modify this logic based on your specific needs
        # most_common_answer = Counter(answers).most_common(1)[0][0]
        
        # # Safely evaluate the most common answer assuming it's a string representation of a Python literal
        # try:
        #     evaluated_answer = ast.literal_eval(most_common_answer)
        # except (ValueError, SyntaxError):
        #     # Handle malformed answer string
        #     evaluated_answer = f"Error evaluating answer: {most_common_answer}"
        
        # return evaluated_answer
    
    #     return ast.literal_eval(answers[0])
        
    # else:
    #     return "Unsupported file type. Please upload an image or video."
        
promt_cat_dog = """
    Annotate this image with this schema:
    {
        “description”: “Is there a cat in the image?”,
        “value”: “Cat”
    },
    {
        “description”: “Is there a dog in the image?”,
        “value”: “Dog”,
    },
    {
        “description”: “Is there a horse in the image?”,
        “value”: “Horse”,
    },
provide me the answers as a dictionary with key as the string value of the variable value on top and its value should be boolean value        
"""
promt_bus_people = """
    Annotate this image with this schema:
    {
        “description”: “Is there a bus in the image?”,
        “value”: “Bus”,
    },
    {
        “description”: “Is there a bike in the image?”,
        “value”: “Bike”,
    },
provide me the answers as a dictionary with key as the string value of the variable value on top and its value should be boolean value        
"""


promt_video = """
Annotate this image with this schema:
    {
        “description”: “Is there a person standing in the image?”,
        “value”: “standing”,
    },
    {
        “description”: “Is the person's hands free in the image?”,
        “value”: “hands-free”,
    },
    {
         “description”: “Is it indoors?”,
          “value”: “Indoors”
    },
provide me the answers as a dictionary with key as the string value of the variable value on top and its value should be boolean value.

"""


test_examples = [[None, "Images/cat_dog.jpeg", promt_cat_dog], 
            [None,"Images/bus_people.jpeg", promt_bus_people], 
            ["videos/v1_new.mp4",None,promt_video],
            ["videos/v3.mp4",None,promt_video]]


def gradio_predict(video,image, question):
    answer = predict_answer(video,image, question)

    return answer

def export_csv(d):
    print(d)
    d_dict = ast.literal_eval(d)
    
    print(type(d_dict))  # This will show <class 'dict'>
    print(d_dict)
    # df = pd.DataFrame(columns=list(d_dict.keys()))
    # df.loc[0] = list(d.values())
    df = pd.DataFrame([d_dict])
    df.to_csv("output.csv", sep=',')
    print(df)
    return gr.File(value="output.csv", visible=True)
    
css = """
#container{
    display: block;
    margin-left: auto;
    margin-right: auto;
    width: 60%;
}
#intro{
    max-width: 100%;
    margin: 0 auto;
    text-align: center;
}

"""
with gr.Blocks(css = css) as app:
    with gr.Row(elem_id="container"):
        gr.Image("gsoc_redhen.png",min_width=60, label="GSOC 2024")
    gr.Markdown("""  
    ## This Gradio app serves as four folds:
    ### 1. My ability and experience to design a customizable Gradio application with Interface/Blocks structure.
    ### 2. One of my Multimodel Vision-Language model's capabilities with the LLaVA framework.
    ### 3. Demo for annotating random images and 4 second videos provided at Notion (https://shorturl.at/givyC) 
    ### 4. Ability to integrate a Large Language Model and Vision Encoder & Downloading annotation CSV    
    ### 5. Click on the below examples to import the cached annotations for images and videos 
                 """)
    with gr.Row():
        video = gr.Video(label="Video")
        image = gr.Image(type="pil", label="Image")
    with gr.Row():
        with gr.Column():
            question = gr.Textbox(label="Annotate", placeholder="Annotate prompt", lines=4.3)
            btn = gr.Button("Annotate")
        with gr.Column():
            answer = gr.TextArea(label="Answer")
            save_btn = gr.Button("Save as CSV")
            csv = gr.File(interactive=False, visible=False)

    # Make sure the inputs and outputs match in your click function
    
    
    btn.click(gradio_predict, inputs=[video,image, question], outputs=answer)

    # Button to save the answer as CSV
    save_btn.click(export_csv, answer, csv)

    gr.Examples(
        examples=test_examples,
        inputs=[video,image, question],
        outputs= answer,
        fn=gradio_predict,
        cache_examples=True,
    )

app.launch(debug=True)